nix_query_tree_viewer/
tree.rs

1use std::collections::{HashMap, VecDeque};
2use std::hash::Hash;
3
4#[derive(Clone, Debug, Eq, PartialEq)]
5pub struct Tree<T> {
6    pub item: T,
7    pub children: Vec<Tree<T>>,
8}
9
10impl<T> Tree<T> {
11    pub fn new(item: T, children: Vec<Tree<T>>) -> Tree<T> {
12        Tree { item, children }
13    }
14
15    pub fn singleton(item: T) -> Tree<T> {
16        Tree::new(item, vec![])
17    }
18
19    /// Lookup the item in the `Tree` that corresponds to the given `Path`.
20    pub fn lookup(&self, path: Path) -> Option<&T> {
21        match path.split_front() {
22            None => Some(&self.item),
23            Some((index, child_path)) => match self.children.get(index) {
24                None => None,
25                Some(child_tree) => child_tree.lookup(child_path),
26            },
27        }
28    }
29
30    /// Similar to `path_map`, but take a function for mapping an item in the tree to an
31    /// alternative type to use to construct the `TreePathMap`.
32    ///
33    /// This is useful when there is some information in the item for the `Tree` that
34    /// can be thrown away when constructing the `TreePathMap`.
35    pub fn path_map_map<U>(&self, f: &dyn Fn(&T) -> U) -> TreePathMap<U>
36    where
37        U: Eq + Hash,
38    {
39        let mut map = TreePathMap::new();
40        let root_path = Path::new();
41        map.insert(f(&self.item), root_path.clone());
42        map.insert_children_map(&self.children, &root_path, f);
43        map
44    }
45}
46
47impl<T> Tree<T>
48where
49    T: Clone + Eq + Hash,
50{
51    /// Create a `TreePathMap` for the elements in the `Tree`.
52    pub fn path_map(&self) -> TreePathMap<T> {
53        self.path_map_map(&|i| i.clone())
54    }
55}
56
57/// This represents the path through a `Tree<T>` to a given node.
58#[derive(Clone, Debug, Default, Eq, PartialEq)]
59pub struct Path(pub VecDeque<usize>);
60
61impl Path {
62    pub fn split_front(mut self) -> Option<(usize, Path)> {
63        let option_front_elem: Option<usize> = self.0.pop_front();
64        match option_front_elem {
65            None => None,
66            Some(i) => Some((i, self)),
67        }
68    }
69
70    pub fn push_back(&mut self, value: usize) {
71        self.0.push_back(value);
72    }
73
74    pub fn new() -> Self {
75        Path(VecDeque::new())
76    }
77}
78
79impl<T> From<T> for Path
80where
81    T: Into<VecDeque<usize>>,
82{
83    fn from(other: T) -> Path {
84        Path(other.into())
85    }
86}
87
88/// This is a mapping of items in `Tree` to their `Path`s.  A single item in the `Tree` can have
89/// multiple `Path`s to it if it is in the `Tree` multiple times.
90#[derive(Clone, Debug, Default, Eq, PartialEq)]
91pub struct TreePathMap<U>(HashMap<U, Vec<Path>>)
92where
93    U: Eq + Hash;
94
95impl<U> TreePathMap<U>
96where
97    U: Eq + Hash,
98{
99    pub fn new() -> TreePathMap<U> {
100        TreePathMap(HashMap::new())
101    }
102
103    /// Insert a mapping from `U` to `Path`.
104    pub fn insert(&mut self, k: U, path: Path) {
105        self.0
106            .entry(k)
107            .and_modify(|paths| paths.push(path.clone()))
108            .or_insert_with(|| vec![path]);
109    }
110
111    /// Lookup the first `Path` for a given item.
112    pub fn lookup_first(&self, k: &U) -> Option<&Path> {
113        let option_paths: Option<&Vec<Path>> = self.0.get(k);
114        option_paths.and_then(|vec: &Vec<Path>| vec.first())
115    }
116}
117
118impl<U> TreePathMap<U>
119where
120    U: Eq + Hash,
121{
122    /// Insert child `Tree`s starting at `Path`.
123    ///
124    /// The function `f` will map the items in the `T` to an alternative type.
125    fn insert_children_map<T>(
126        &mut self,
127        children: &[Tree<T>],
128        path: &Path,
129        f: &dyn Fn(&T) -> U,
130    ) {
131        for (i, child) in children.iter().enumerate() {
132            let mut child_path = path.clone();
133            child_path.push_back(i);
134            self.insert(f(&child.item), child_path.clone());
135            self.insert_children_map(&child.children, &child_path, f);
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    use std::ops::Deref;
145
146    #[test]
147    fn test_path_split_front_empty() {
148        let res = Path::new().split_front();
149        assert_eq!(res, None);
150    }
151
152    #[test]
153    fn test_path_push_back() {
154        let mut path = Path::new();
155        path.push_back(1);
156        path.push_back(2);
157        path.push_back(3);
158
159        let mut actual_vec = VecDeque::new();
160        actual_vec.push_back(1);
161        actual_vec.push_back(2);
162        actual_vec.push_back(3);
163        let actual_path = Path(actual_vec);
164
165        assert_eq!(path, actual_path);
166    }
167
168    #[test]
169    fn test_path_split_front_nonempty() {
170        let mut path = Path::new();
171        path.push_back(1);
172        path.push_back(2);
173        path.push_back(3);
174        let res = path.split_front();
175
176        let mut actual_path = Path::new();
177        actual_path.push_back(2);
178        actual_path.push_back(3);
179
180        let actual = Some((1, actual_path));
181
182        assert_eq!(res, actual);
183    }
184
185    #[test]
186    fn test_lookup_no_item() {
187        let tree = Tree::new(
188            "root",
189            vec![
190                Tree::singleton("0"),
191                Tree::singleton("1"),
192                Tree::new(
193                    "2",
194                    vec![Tree::singleton("2-0"), Tree::singleton("2-1")],
195                ),
196            ],
197        );
198
199        let path1 = vec![3].into();
200        let path2 = vec![0, 1].into();
201        let path3 = vec![1, 0].into();
202        let path4 = vec![2, 2].into();
203        let path5 = vec![2, 0, 3].into();
204        let path6 = vec![0, 1, 2, 3, 4].into();
205
206        assert_eq!(tree.lookup(path1), None);
207        assert_eq!(tree.lookup(path2), None);
208        assert_eq!(tree.lookup(path3), None);
209        assert_eq!(tree.lookup(path4), None);
210        assert_eq!(tree.lookup(path5), None);
211        assert_eq!(tree.lookup(path6), None);
212    }
213
214    #[test]
215    fn test_lookup_find_item() {
216        let tree: Tree<String> = Tree::new(
217            "root".into(),
218            vec![
219                Tree::singleton("0".into()),
220                Tree::singleton("1".into()),
221                Tree::new(
222                    "2".into(),
223                    vec![
224                        Tree::singleton("2-0".into()),
225                        Tree::new(
226                            "2-1".into(),
227                            vec![
228                                Tree::singleton("2-1-0".into()),
229                                Tree::singleton("2-1-1".into()),
230                            ],
231                        ),
232                    ],
233                ),
234            ],
235        );
236
237        let path_root = vec![].into();
238        let path0 = vec![0].into();
239        let path1 = vec![1].into();
240        let path2 = vec![2].into();
241        let path2_0 = vec![2, 0].into();
242        let path2_1 = vec![2, 1].into();
243        let path2_1_0 = vec![2, 1, 0].into();
244        let path2_1_1 = vec![2, 1, 1].into();
245
246        assert_eq!(tree.lookup(path_root).map(String::deref), Some("root"));
247        assert_eq!(tree.lookup(path0).map(String::deref), Some("0"));
248        assert_eq!(tree.lookup(path1).map(String::deref), Some("1"));
249        assert_eq!(tree.lookup(path2).map(String::deref), Some("2"));
250        assert_eq!(tree.lookup(path2_0).map(String::deref), Some("2-0"));
251        assert_eq!(tree.lookup(path2_1).map(String::deref), Some("2-1"));
252        assert_eq!(tree.lookup(path2_1_0).map(String::deref), Some("2-1-0"));
253        assert_eq!(tree.lookup(path2_1_1).map(String::deref), Some("2-1-1"));
254    }
255
256    #[test]
257    fn test_tree_path_map_from_tree_all_unique() {
258        let tree: Tree<String> = Tree::new(
259            "root".into(),
260            vec![
261                Tree::singleton("0".into()),
262                Tree::singleton("1".into()),
263                Tree::new(
264                    "2".into(),
265                    vec![
266                        Tree::singleton("2-0".into()),
267                        Tree::new(
268                            "2-1".into(),
269                            vec![
270                                Tree::singleton("2-1-0".into()),
271                                Tree::singleton("2-1-1".into()),
272                            ],
273                        ),
274                    ],
275                ),
276            ],
277        );
278
279        let res_tree_path_map: TreePathMap<String> = tree.path_map();
280
281        let mut actual_tree_path_map: HashMap<String, Vec<Path>> =
282            HashMap::new();
283
284        actual_tree_path_map.insert("root".into(), vec![Path::new()]);
285        actual_tree_path_map.insert("0".into(), vec![vec![0].into()]);
286        actual_tree_path_map.insert("1".into(), vec![vec![1].into()]);
287        actual_tree_path_map.insert("2".into(), vec![vec![2].into()]);
288        actual_tree_path_map.insert("2-0".into(), vec![vec![2, 0].into()]);
289        actual_tree_path_map.insert("2-1".into(), vec![vec![2, 1].into()]);
290        actual_tree_path_map.insert("2-1-0".into(), vec![vec![2, 1, 0].into()]);
291        actual_tree_path_map.insert("2-1-1".into(), vec![vec![2, 1, 1].into()]);
292
293        assert_eq!(res_tree_path_map, TreePathMap(actual_tree_path_map));
294    }
295
296    #[test]
297    fn test_tree_path_map_from_tree() {
298        let tree: Tree<String> = Tree::new(
299            "cat".into(), // root
300            vec![
301                Tree::singleton("dog".into()), // 0
302                Tree::singleton("cat".into()), // 1
303                Tree::new(
304                    "mouse".into(), // 2
305                    vec![
306                        Tree::singleton("fish".into()), // 2-0
307                        Tree::new(
308                            "fish".into(), // 2-1
309                            vec![
310                                Tree::singleton("dog".into()), // 2-1-0
311                                Tree::singleton("cat".into()), // 2-1-1
312                            ],
313                        ),
314                    ],
315                ),
316            ],
317        );
318
319        let res_tree_path_map: TreePathMap<String> = tree.path_map();
320
321        let mut actual_tree_path_map: HashMap<String, Vec<Path>> =
322            HashMap::new();
323
324        actual_tree_path_map.insert(
325            "cat".into(),
326            vec![Path::new(), vec![1].into(), vec![2, 1, 1].into()],
327        );
328        actual_tree_path_map
329            .insert("dog".into(), vec![vec![0].into(), vec![2, 1, 0].into()]);
330        actual_tree_path_map.insert("mouse".into(), vec![vec![2].into()]);
331        actual_tree_path_map
332            .insert("fish".into(), vec![vec![2, 0].into(), vec![2, 1].into()]);
333
334        assert_eq!(res_tree_path_map, TreePathMap(actual_tree_path_map));
335    }
336}
nix_query_tree_viewer/tree.rs

nix_query_tree_viewer/
tree.rs
