mlx_rs/
nested.rs

1//! Implements a nested hashmap
2
3use std::{collections::HashMap, fmt::Display, rc::Rc};
4
5const DELIMITER: char = '.';
6
7/// A nested value that can be either a value or a map of nested values
8#[derive(Debug, Clone)]
9pub enum NestedValue<K, T> {
10    /// A value
11    Value(T),
12
13    /// A map of nested values
14    Map(HashMap<K, NestedValue<K, T>>),
15}
16
17impl<K, V> NestedValue<K, V> {
18    /// Flattens the nested value into a hashmap
19    pub fn flatten(self, prefix: &str) -> HashMap<Rc<str>, V>
20    where
21        K: Display,
22    {
23        match self {
24            NestedValue::Value(array) => {
25                let mut map = HashMap::new();
26                map.insert(prefix.into(), array);
27                map
28            }
29            NestedValue::Map(entries) => entries
30                .into_iter()
31                .flat_map(|(key, value)| value.flatten(&format!("{prefix}{DELIMITER}{key}")))
32                .collect(),
33        }
34    }
35}
36
37/// A nested hashmap
38#[derive(Debug, Clone)]
39pub struct NestedHashMap<K, V> {
40    /// The internal hashmap
41    pub entries: HashMap<K, NestedValue<K, V>>,
42}
43
44impl<K, V> From<NestedHashMap<K, V>> for NestedValue<K, V> {
45    fn from(map: NestedHashMap<K, V>) -> Self {
46        NestedValue::Map(map.entries)
47    }
48}
49
50impl<K, V> Default for NestedHashMap<K, V> {
51    fn default() -> Self {
52        Self::new()
53    }
54}
55
56impl<K, V> NestedHashMap<K, V> {
57    /// Creates a new nested hashmap
58    pub fn new() -> Self {
59        Self {
60            entries: HashMap::new(),
61        }
62    }
63
64    /// Inserts a new entry into the nested hashmap
65    pub fn insert(&mut self, key: K, value: NestedValue<K, V>)
66    where
67        K: Eq + std::hash::Hash,
68    {
69        self.entries.insert(key, value);
70    }
71
72    /// Flattens the nested hashmap into a hashmap
73    pub fn flatten(self) -> HashMap<Rc<str>, V>
74    where
75        K: AsRef<str> + Display,
76    {
77        self.entries
78            .into_iter()
79            .flat_map(|(key, value)| value.flatten(key.as_ref()))
80            .collect()
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use crate::{
87        array,
88        test_utils::{assert_array_eq, tolerances},
89    };
90
91    use super::*;
92
93    #[test]
94    fn test_flatten_nested_hash_map_of_owned_arrays() {
95        let first_entry = NestedValue::Value(array!([1, 2, 3]));
96        let second_entry = NestedValue::Map({
97            let mut map = HashMap::new();
98            map.insert("a", NestedValue::Value(array!([4, 5, 6])));
99            map.insert("b", NestedValue::Value(array!([7, 8, 9])));
100            map
101        });
102
103        let map = NestedHashMap {
104            entries: {
105                let mut map = HashMap::new();
106                map.insert("first", first_entry);
107                map.insert("second", second_entry);
108                map
109            },
110        };
111
112        let flattened = map.flatten();
113
114        assert_eq!(flattened.len(), 3);
115        assert_array_eq(
116            &flattened["first"],
117            array!([1, 2, 3]),
118            tolerances::EXACT.rtol,
119            tolerances::EXACT.atol,
120        );
121        assert_array_eq(
122            &flattened["second.a"],
123            array!([4, 5, 6]),
124            tolerances::EXACT.rtol,
125            tolerances::EXACT.atol,
126        );
127        assert_array_eq(
128            &flattened["second.b"],
129            array!([7, 8, 9]),
130            tolerances::EXACT.rtol,
131            tolerances::EXACT.atol,
132        );
133    }
134
135    #[test]
136    fn test_flatten_nested_hash_map_of_borrowed_arrays() {
137        let first_entry_content = array!([1, 2, 3]);
138        let first_entry = NestedValue::Value(&first_entry_content);
139
140        let second_entry_content_a = array!([4, 5, 6]);
141        let second_entry_content_b = array!([7, 8, 9]);
142        let second_entry = NestedValue::Map({
143            let mut map = HashMap::new();
144            map.insert("a", NestedValue::Value(&second_entry_content_a));
145            map.insert("b", NestedValue::Value(&second_entry_content_b));
146            map
147        });
148
149        let map = NestedHashMap {
150            entries: {
151                let mut map = HashMap::new();
152                map.insert("first", first_entry);
153                map.insert("second", second_entry);
154                map
155            },
156        };
157
158        let flattened = map.flatten();
159
160        assert_eq!(flattened.len(), 3);
161        assert_array_eq(
162            &**flattened.get("first").unwrap(),
163            &first_entry_content,
164            tolerances::EXACT.rtol,
165            tolerances::EXACT.atol,
166        );
167        assert_array_eq(
168            &**flattened.get("second.a").unwrap(),
169            &second_entry_content_a,
170            tolerances::EXACT.rtol,
171            tolerances::EXACT.atol,
172        );
173        assert_array_eq(
174            &**flattened.get("second.b").unwrap(),
175            &second_entry_content_b,
176            tolerances::EXACT.rtol,
177            tolerances::EXACT.atol,
178        );
179    }
180
181    #[test]
182    fn test_flatten_nested_hash_map_of_mut_borrowed_arrays() {
183        let mut first_entry_content = array!([1, 2, 3]);
184        let first_entry = NestedValue::Value(&mut first_entry_content);
185
186        let mut second_entry_content_a = array!([4, 5, 6]);
187        let mut second_entry_content_b = array!([7, 8, 9]);
188        let second_entry = NestedValue::Map({
189            let mut map = HashMap::new();
190            map.insert("a", NestedValue::Value(&mut second_entry_content_a));
191            map.insert("b", NestedValue::Value(&mut second_entry_content_b));
192            map
193        });
194
195        let map = NestedHashMap {
196            entries: {
197                let mut map = HashMap::new();
198                map.insert("first", first_entry);
199                map.insert("second", second_entry);
200                map
201            },
202        };
203
204        let flattened = map.flatten();
205
206        assert_eq!(flattened.len(), 3);
207        assert_array_eq(
208            &**flattened.get("first").unwrap(),
209            &mut array!([1, 2, 3]),
210            tolerances::EXACT.rtol,
211            tolerances::EXACT.atol,
212        );
213        assert_array_eq(
214            &**flattened.get("second.a").unwrap(),
215            &mut array!([4, 5, 6]),
216            tolerances::EXACT.rtol,
217            tolerances::EXACT.atol,
218        );
219        assert_array_eq(
220            &**flattened.get("second.b").unwrap(),
221            &mut array!([7, 8, 9]),
222            tolerances::EXACT.rtol,
223            tolerances::EXACT.atol,
224        );
225    }
226
227    #[test]
228    fn test_flatten_empty_nested_hash_map() {
229        let map = NestedHashMap::<&str, i32>::new();
230        let flattened = map.flatten();
231
232        assert!(flattened.is_empty());
233
234        // Insert another empty map
235        let mut map = NestedHashMap::<&str, i32>::new();
236        let empty_map = NestedValue::Map(HashMap::new());
237        map.insert("empty", empty_map);
238
239        let flattened = map.flatten();
240        assert!(flattened.is_empty());
241    }
242}
mlx_rs/nested.rs

mlx_rs/
nested.rs
