bitwarden_collections/
tree.rs

1use std::{collections::HashMap, fmt::Debug};
2
3use uuid::Uuid;
4
5#[allow(missing_docs)]
6pub trait TreeItem: Clone + Debug {
7    fn id(&self) -> Uuid;
8    /*
9    This is the name that will be output when getting the tree nodes
10     */
11    fn short_name(&self) -> &str;
12    /*
13    This is the path that the item is stored into a tree
14     */
15    fn path(&self) -> Vec<&str>;
16    const DELIMITER: char;
17}
18
19#[allow(missing_docs)]
20#[derive(Clone, Debug)]
21pub struct TreeIndex<T: TreeItem> {
22    pub id: usize, // location in the tree
23    pub data: T,   // this will be the raw value
24    pub path: Vec<String>,
25}
26
27impl<T: TreeItem> TreeIndex<T> {
28    #[allow(missing_docs)]
29    pub fn new(id: usize, data: &T) -> Self {
30        TreeIndex {
31            id,
32            data: data.clone(),
33            path: data.path().iter().map(|s| s.to_string()).collect(),
34        }
35    }
36}
37
38#[allow(missing_docs)]
39pub struct NodeItem<T: TreeItem> {
40    pub item: T,
41    pub parent: Option<T>,
42    pub children: Vec<T>,
43    pub ancestors: HashMap<Uuid, String>,
44}
45
46#[allow(missing_docs)]
47pub struct TreeNode {
48    pub id: usize,
49    pub item_id: Uuid,
50    pub parent_idx: Option<usize>,
51    pub children_idx: Vec<usize>,
52    pub path: Vec<String>,
53}
54
55impl TreeNode {
56    #[allow(missing_docs)]
57    pub fn new<T: TreeItem>(
58        id: usize,
59        parent_idx: Option<usize>,
60        children_idx: Vec<usize>,
61        index: TreeIndex<T>,
62    ) -> Self {
63        TreeNode {
64            id,
65            item_id: index.data.id(),
66            parent_idx,
67            children_idx,
68            path: index.path,
69        }
70    }
71}
72
73#[allow(missing_docs)]
74pub struct Tree<T: TreeItem> {
75    pub nodes: Vec<TreeNode>,
76    pub items: HashMap<Uuid, TreeIndex<T>>,
77    path_to_node: HashMap<Vec<String>, usize>,
78}
79
80impl<T: TreeItem> Tree<T> {
81    /// Takes vector of TreeItem and stores them into a tree structure.
82    pub fn from_items(items: Vec<T>) -> Self {
83        let mut tree = Tree {
84            nodes: Vec::new(),
85            items: HashMap::new(),
86            path_to_node: HashMap::new(),
87        };
88
89        // sort items
90        let sorted_items = {
91            let mut i = items.clone();
92            i.sort_by(|a, b| a.path().cmp(&b.path()));
93            i
94        };
95
96        // add items
97        for (index, item) in sorted_items.iter().enumerate() {
98            let tree_index = TreeIndex::new(index, item);
99            tree.items.insert(item.id(), tree_index.clone());
100            tree.add_item(tree_index);
101        }
102
103        tree
104    }
105
106    /// This inserts an item into the tree and sets any look-up information that may be needed.
107    fn add_item(&mut self, index: TreeIndex<T>) {
108        let parent_path = index.path[0..index.path.len() - 1].to_vec();
109
110        let parent_id = self.path_to_node.get(&parent_path).map(|&id| {
111            let parent = &mut self.nodes[id];
112            parent.children_idx.push(index.id);
113            parent.id
114        });
115
116        // add new node
117        let node = TreeNode::new(index.id, parent_id, vec![], index);
118        self.path_to_node.insert(node.path.clone(), node.id);
119        self.nodes.push(node);
120    }
121
122    /// Returns an optional node item for a given tree item id.
123    ///
124    /// This contains the item, its children (or an empty vector), and its parent (if it has one)
125    pub fn get_item_by_id(&self, tree_item_id: Uuid) -> Option<NodeItem<T>> {
126        let item = self.items.get(&tree_item_id)?;
127
128        self.get_relatives(item)
129    }
130
131    fn get_relatives(&self, item: &TreeIndex<T>) -> Option<NodeItem<T>> {
132        let node = self.nodes.get(item.id)?;
133
134        // Get the parent if it exists
135        let parent = node
136            .parent_idx
137            .and_then(|pid| self.nodes.get(pid))
138            .and_then(|p| self.items.get(&p.item_id))
139            .map(|p| p.data.clone());
140
141        // Get any children
142        let children: Vec<T> = node
143            .children_idx
144            .iter()
145            .filter_map(|&child_id| self.nodes.get(child_id))
146            .filter_map(|child| self.items.get(&child.item_id))
147            .map(|i| i.data.clone())
148            .collect();
149
150        // Get names and ids of ancestors
151        let ancestors = std::iter::successors(node.parent_idx, |&parent_id| {
152            self.nodes.get(parent_id).and_then(|node| node.parent_idx)
153        })
154        .filter_map(|parent_id| {
155            self.nodes
156                .get(parent_id)
157                .and_then(|parent_node| self.items.get(&parent_node.item_id))
158                .map(|parent_item| {
159                    (
160                        parent_item.data.id(),
161                        parent_item.data.short_name().to_string(),
162                    )
163                })
164        })
165        .collect();
166
167        Some(NodeItem {
168            item: item.data.clone(),
169            parent,
170            children,
171            ancestors,
172        })
173    }
174
175    /// Returns the list of root nodes with their children
176    pub fn get_root_items(&self) -> Vec<NodeItem<T>> {
177        self.nodes
178            .iter()
179            .filter(|n| n.parent_idx.is_none())
180            .filter_map(|n| self.get_item_by_id(n.item_id))
181            .collect()
182    }
183
184    /// Returns a flat list of all items in the tree with relationships.
185    pub fn get_flat_items(&self) -> Vec<NodeItem<T>> {
186        self.items
187            .values()
188            .filter_map(|i| self.get_relatives(i))
189            .collect()
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use uuid::Uuid;
196
197    use super::*;
198
199    #[derive(Clone, Debug)]
200    pub struct TestItem {
201        pub id: Uuid,
202        pub name: String,
203    }
204
205    impl TreeItem for TestItem {
206        fn id(&self) -> Uuid {
207            self.id
208        }
209
210        fn short_name(&self) -> &str {
211            self.path().last().unwrap()
212        }
213
214        fn path(&self) -> Vec<&str> {
215            self.name
216                .split(Self::DELIMITER)
217                .filter(|s| !s.is_empty())
218                .collect::<Vec<&str>>()
219        }
220
221        const DELIMITER: char = '/';
222    }
223
224    #[test]
225    fn given_collection_with_one_parent_and_two_children_when_getting_parent_then_parent_is_returned_with_children_and_no_parent(
226    ) {
227        let parent_id = Uuid::new_v4();
228        let items = vec![
229            TestItem {
230                id: Uuid::new_v4(),
231                name: "parent/child1".to_string(),
232            },
233            TestItem {
234                id: parent_id,
235                name: "parent".to_string(),
236            },
237            TestItem {
238                id: Uuid::new_v4(),
239                name: "parent/child2".to_string(),
240            },
241        ];
242
243        let node = Tree::from_items(items)
244            .get_item_by_id(parent_id)
245            .expect("Node not found");
246
247        let item = node.item;
248        let parent = node.parent;
249        let children = node.children;
250        let ancestors = node.ancestors;
251
252        assert_eq!(children.len(), 2);
253        assert_eq!(item.id(), parent_id);
254        assert_eq!(item.short_name(), "parent");
255        assert_eq!(item.path(), ["parent"]);
256        assert!(parent.is_none());
257        assert!(ancestors.is_empty());
258    }
259
260    #[test]
261    fn given_collection_with_one_parent_and_two_children_when_getting_child1_then_child1_is_returned_with_no_children_and_a_parent(
262    ) {
263        let child_1_id = Uuid::new_v4();
264        let parent_id = Uuid::new_v4();
265        let items = vec![
266            TestItem {
267                id: child_1_id,
268                name: "parent/child1".to_string(),
269            },
270            TestItem {
271                id: parent_id,
272                name: "parent".to_string(),
273            },
274            TestItem {
275                id: Uuid::new_v4(),
276                name: "parent/child2".to_string(),
277            },
278        ];
279
280        let node = Tree::from_items(items)
281            .get_item_by_id(child_1_id)
282            .expect("Node not found");
283
284        let item = node.item;
285        let parent = node.parent;
286        let children = node.children;
287        let ancestors = node.ancestors;
288
289        assert_eq!(children.len(), 0);
290        assert_eq!(item.id(), child_1_id);
291        assert_eq!(item.short_name(), "child1");
292        assert_eq!(item.path(), ["parent", "child1"]);
293        assert_eq!(parent.unwrap().id, parent_id);
294        assert_eq!(ancestors.len(), 1);
295        assert_eq!(ancestors.get(&parent_id).unwrap(), "parent");
296    }
297
298    #[test]
299    fn given_collection_with_child_who_has_parent_and_grandparent_returns_correct_ancestors() {
300        let child_1_id = Uuid::new_v4();
301        let parent_id = Uuid::new_v4();
302        let grandparent_id = Uuid::new_v4();
303        let items = vec![
304            TestItem {
305                id: child_1_id,
306                name: "grandparent/parent/child".to_string(),
307            },
308            TestItem {
309                id: parent_id,
310                name: "grandparent/parent".to_string(),
311            },
312            TestItem {
313                id: grandparent_id,
314                name: "grandparent".to_string(),
315            },
316        ];
317
318        let node = Tree::from_items(items)
319            .get_item_by_id(child_1_id)
320            .expect("Node not found");
321
322        let item = node.item;
323        let parent = node.parent;
324        let children = node.children;
325        let ancestors = node.ancestors;
326
327        assert_eq!(children.len(), 0);
328        assert_eq!(item.id(), child_1_id);
329        assert_eq!(item.short_name(), "child");
330        assert_eq!(item.path(), ["grandparent", "parent", "child"]);
331        assert_eq!(parent.unwrap().id, parent_id);
332        assert_eq!(ancestors.len(), 2);
333        assert_eq!(ancestors.get(&parent_id).unwrap(), "parent");
334        assert_eq!(ancestors.get(&grandparent_id).unwrap(), "grandparent");
335    }
336}
bitwarden_collections/tree.rs

bitwarden_collections/
tree.rs
