Skip to main content

muse2/
id.rs

1//! Code for handling IDs
2use anyhow::Result;
3use indexmap::{IndexMap, IndexSet};
4use std::borrow::Borrow;
5use std::collections::HashSet;
6use std::error::Error;
7use std::fmt::{Debug, Display};
8use std::hash::Hash;
9use std::marker::PhantomData;
10
11/// A trait for ID types
12pub trait ID: Eq + Hash + Borrow<str> + Clone + Display + Debug + From<String> {
13    /// The name of this type of ID (e.g. "commodity ID", "region ID" etc.)
14    const TYPE_NAME: &'static str;
15}
16
17macro_rules! define_id_type {
18    ($name:ident, $type_name:expr) => {
19        #[derive(
20            Clone,
21            derive_more::Display,
22            derive_more::From,
23            std::hash::Hash,
24            PartialOrd,
25            Ord,
26            PartialEq,
27            Eq,
28            Debug,
29            serde::Serialize,
30        )]
31        /// An ID type (e.g. `AgentID`, `CommodityID`, etc.)
32        #[from(forward)]
33        pub struct $name(pub std::rc::Rc<str>);
34
35        impl std::borrow::Borrow<str> for $name {
36            fn borrow(&self) -> &str {
37                &self.0
38            }
39        }
40
41        impl<'de> serde::Deserialize<'de> for $name {
42            fn deserialize<D>(deserialiser: D) -> std::result::Result<Self, D::Error>
43            where
44                D: serde::Deserializer<'de>,
45            {
46                use serde::de::Error;
47                const FORBIDDEN_IDS: [&str; 2] = ["all", "annual"];
48
49                let id: String = serde::Deserialize::deserialize(deserialiser)?;
50                let id = id.trim();
51                if id.is_empty() {
52                    return Err(D::Error::custom("IDs cannot be empty"));
53                }
54
55                for forbidden in FORBIDDEN_IDS.iter() {
56                    if id.eq_ignore_ascii_case(forbidden) {
57                        return Err(D::Error::custom(format!(
58                            "'{id}' is an invalid value for an ID"
59                        )));
60                    }
61                }
62
63                Ok(id.into())
64            }
65        }
66
67        impl crate::id::ID for $name {
68            const TYPE_NAME: &'static str = $type_name;
69        }
70
71        impl $name {
72            /// Create a new ID from a string slice
73            pub fn new(id: &str) -> Self {
74                $name(std::rc::Rc::from(id))
75            }
76        }
77    };
78}
79pub(crate) use define_id_type;
80
81#[cfg(test)]
82define_id_type!(GenericID, "generic ID");
83
84/// Indicates that the struct has an ID field
85pub trait HasID<T: ID> {
86    /// Get the struct's ID
87    fn get_id(&self) -> &T;
88}
89
90/// Implement the `HasID` trait for the given type, assuming it has a field called `id`
91macro_rules! define_id_getter {
92    ($t:ty, $id_ty:ty) => {
93        impl crate::id::HasID<$id_ty> for $t {
94            fn get_id(&self) -> &$id_ty {
95                &self.id
96            }
97        }
98    };
99}
100pub(crate) use define_id_getter;
101
102/// Indicates that the specified ID was not found in a given collection
103#[derive(Debug, derive_more::Display)]
104#[display("Unknown {} '{missing_id}'", T::TYPE_NAME)]
105pub struct MissingIDError<T: ID> {
106    missing_id: String,
107    _phantom: PhantomData<fn() -> T>,
108}
109
110impl<T: ID> MissingIDError<T> {
111    /// Create a new `MissingIDError`
112    pub fn new(missing_id: &str) -> MissingIDError<T> {
113        MissingIDError::<T> {
114            missing_id: missing_id.to_string(),
115            _phantom: std::marker::PhantomData,
116        }
117    }
118}
119
120impl<T: ID> Error for MissingIDError<T> {}
121
122/// A data structure containing a set of IDs
123pub trait IDCollection<T: ID> {
124    /// Check if the ID is in the collection, returning a reference to it if found.
125    ///
126    /// # Arguments
127    ///
128    /// * `id` - The ID to check (can be string or ID type)
129    ///
130    /// # Returns
131    ///
132    /// A reference to the ID in `self`, or an error if not found.
133    fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>>;
134}
135
136macro_rules! define_id_methods {
137    () => {
138        fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>> {
139            self.get(id.borrow())
140                .ok_or_else(|| MissingIDError::new(id.borrow()))
141        }
142    };
143}
144
145impl<T: ID> IDCollection<T> for HashSet<T> {
146    define_id_methods!();
147}
148
149impl<T: ID> IDCollection<T> for IndexSet<T> {
150    define_id_methods!();
151}
152
153impl<T: ID, V> IDCollection<T> for IndexMap<T, V> {
154    fn get_id<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<&T, MissingIDError<T>> {
155        let (found, _) = self
156            .get_key_value(id.borrow())
157            .ok_or_else(|| MissingIDError::new(id.borrow()))?;
158        Ok(found)
159    }
160}
161
162/// A trait for getting an ID and a value from a map
163pub trait GetIDValue<K: ID, V> {
164    /// Get the ID and value, if any, for the given collection
165    fn get_id_value<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<(&K, &V), MissingIDError<K>>;
166}
167
168impl<K: ID + Borrow<str>, V> GetIDValue<K, V> for IndexMap<K, V> {
169    fn get_id_value<S: Borrow<str> + ?Sized>(&self, id: &S) -> Result<(&K, &V), MissingIDError<K>> {
170        self.get_key_value(id.borrow())
171            .ok_or_else(|| MissingIDError::new(id.borrow()))
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use rstest::rstest;
179
180    use serde::Deserialize;
181
182    #[derive(Debug, Deserialize)]
183    struct Record {
184        id: GenericID,
185    }
186
187    fn deserialise_id(id: &str) -> Result<Record> {
188        Ok(toml::from_str(&format!("id = \"{id}\""))?)
189    }
190
191    #[rstest]
192    #[case("commodity1")]
193    #[case("some commodity")]
194    #[case("PROCESS")]
195    #[case("café")] // unicode supported
196    fn deserialise_id_valid(#[case] id: &str) {
197        assert_eq!(deserialise_id(id).unwrap().id.to_string(), id);
198    }
199
200    #[rstest]
201    #[case("")]
202    #[case("all")]
203    #[case("annual")]
204    #[case("ALL")]
205    #[case(" ALL ")]
206    fn deserialise_id_invalid(#[case] id: &str) {
207        deserialise_id(id).unwrap_err();
208    }
209}