muse2/

id.rs

//! Code for handling IDs
use anyhow::Result;
use indexmap::{IndexMap, IndexSet};
use std::borrow::Borrow;
use std::collections::HashSet;
use std::error::Error;
use std::fmt::{Debug, Display};
use std::hash::Hash;
use std::marker::PhantomData;

/// A trait for ID types
pub trait ID: Eq + Hash + Borrow<str> + Clone + Display + Debug + From<String> {
    /// Get the name of this type of ID (e.g. "commodity ID", "region ID" etc.)
    fn get_type_name() -> &'static str;
}

macro_rules! define_id_type {
    ($name:ident, $type_name:expr) => {
        #[derive(
            Clone,
            derive_more::Display,
            std::hash::Hash,
            PartialOrd,
            Ord,
            PartialEq,
            Eq,
            Debug,
            serde::Serialize,
        )]
        /// An ID type (e.g. `AgentID`, `CommodityID`, etc.)
        pub struct $name(pub std::rc::Rc<str>);

        impl std::borrow::Borrow<str> for $name {
            fn borrow(&self) -> &str {
                &self.0
            }
        }

        impl From<&str> for $name {
            fn from(s: &str) -> Self {
                $name(std::rc::Rc::from(s))
            }
        }

        impl From<String> for $name {
            fn from(s: String) -> Self {
                $name(std::rc::Rc::from(s))
            }
        }

        impl<'de> serde::Deserialize<'de> for $name {
            fn deserialize<D>(deserialiser: D) -> std::result::Result<Self, D::Error>
            where
                D: serde::Deserializer<'de>,
            {
                use serde::de::Error;
                const FORBIDDEN_IDS: [&str; 2] = ["all", "annual"];

                let id: String = serde::Deserialize::deserialize(deserialiser)?;
                let id = id.trim();
                if id.is_empty() {
                    return Err(D::Error::custom("IDs cannot be empty"));
                }

                for forbidden in FORBIDDEN_IDS.iter() {
                    if id.eq_ignore_ascii_case(forbidden) {
                        return Err(D::Error::custom(format!(
                            "'{id}' is an invalid value for an ID"
                        )));
                    }
                }

                Ok(id.into())
            }
        }

        impl crate::id::ID for $name {
            fn get_type_name() -> &'static str {
                $type_name
            }
        }

        impl $name {
            /// Create a new ID from a string slice
            pub fn new(id: &str) -> Self {
                $name(std::rc::Rc::from(id))
            }
        }
    };
}
pub(crate) use define_id_type;

#[cfg(test)]
define_id_type!(GenericID, "generic ID");

/// Indicates that the struct has an ID field
pub trait HasID<T: ID> {
    /// Get the struct's ID
    fn get_id(&self) -> &T;
}

/// Implement the `HasID` trait for the given type, assuming it has a field called `id`
macro_rules! define_id_getter {
    ($t:ty, $id_ty:ty) => {
        impl crate::id::HasID<$id_ty> for $t {
            fn get_id(&self) -> &$id_ty {
                &self.id
            }
        }
    };
}
pub(crate) use define_id_getter;

/// Indicates that the specified ID was not found in a given collection
#[derive(Debug)]
pub struct MissingIDError<T: ID> {
    missing_id: String,
    _phantom: PhantomData<fn() -> T>,
}

impl<T: ID> MissingIDError<T> {
    /// Create a new `MissingIDError`
    pub fn new(missing_id: &str) -> MissingIDError<T> {
        MissingIDError::<T> {
            missing_id: missing_id.to_string(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: ID> Display for MissingIDError<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Unknown {} '{}'", T::get_type_name(), self.missing_id)
    }
}

impl<T: ID> Error for MissingIDError<T> {}

/// A data structure containing a set of IDs
pub trait IDCollection<T: ID> {
    /// Check if the ID is in the collection, returning a reference to it if found.
    ///
    /// # Arguments
    ///
    /// * `id` - The ID to check (can be string or ID type)
    ///
    /// # Returns
    ///
    /// A reference to the ID in `self`, or an error if not found.
    fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>>;
}

macro_rules! define_id_methods {
    () => {
        fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>> {
            self.get(id.borrow())
                .ok_or_else(|| MissingIDError::new(id.borrow()))
        }
    };
}

impl<T: ID> IDCollection<T> for HashSet<T> {
    define_id_methods!();
}

impl<T: ID> IDCollection<T> for IndexSet<T> {
    define_id_methods!();
}

impl<T: ID, V> IDCollection<T> for IndexMap<T, V> {
    fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>> {
        let (found, _) = self
            .get_key_value(id.borrow())
            .ok_or_else(|| MissingIDError::new(id.borrow()))?;
        Ok(found)
    }
}

/// A trait for getting an ID and a value from a map
pub trait GetIDValue<K: ID, V> {
    /// Get the ID and value, if any, for the given collection
    fn get_id_value<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<(&K, &V), MissingIDError<K>>;
}

impl<K: ID + Borrow<str>, V> GetIDValue<K, V> for IndexMap<K, V> {
    fn get_id_value<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<(&K, &V), MissingIDError<K>> {
        self.get_key_value(id.borrow())
            .ok_or_else(|| MissingIDError::new(id.borrow()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;

    use serde::Deserialize;

    #[derive(Debug, Deserialize)]
    struct Record {
        id: GenericID,
    }

    fn deserialise_id(id: &str) -> Result<Record> {
        Ok(toml::from_str(&format!("id = \"{id}\""))?)
    }

    #[rstest]
    #[case("commodity1")]
    #[case("some commodity")]
    #[case("PROCESS")]
    #[case("café")] // unicode supported
    fn deserialise_id_valid(#[case] id: &str) {
        assert_eq!(deserialise_id(id).unwrap().id.to_string(), id);
    }

    #[rstest]
    #[case("")]
    #[case("all")]
    #[case("annual")]
    #[case("ALL")]
    #[case(" ALL ")]
    fn deserialise_id_invalid(#[case] id: &str) {
        deserialise_id(id).unwrap_err();
    }
}