lrcat/
keywordtree.rs

1/*
2 This Source Code Form is subject to the terms of the Mozilla Public
3 License, v. 2.0. If a copy of the MPL was not distributed with this
4 file, You can obtain one at http://mozilla.org/MPL/2.0/.
5*/
6
7use std::collections::{BTreeMap, HashMap};
8
9use super::keywords::Keyword;
10use super::lrobject::LrObject;
11
12/// Keyword tree
13/// Operate as a hash multimap of parent -> `Vec<child>`
14#[derive(Default)]
15pub struct KeywordTree {
16    // HashMap. Key is the parent id. Values: the children ids.
17    map: HashMap<i64, Vec<i64>>,
18}
19
20impl KeywordTree {
21    pub fn new() -> KeywordTree {
22        KeywordTree::default()
23    }
24
25    /// Get children for keyword with `id`
26    pub fn children_for(&self, id: i64) -> Vec<i64> {
27        if let Some(children) = self.map.get(&id) {
28            return children.clone();
29        }
30        vec![]
31    }
32
33    fn add_child(&mut self, keyword: &Keyword) {
34        self.map.entry(keyword.parent).or_default();
35        self.map
36            .get_mut(&keyword.parent)
37            .unwrap()
38            .push(keyword.id());
39    }
40
41    /// Add children to the tree node.
42    pub fn add_children(&mut self, children: &BTreeMap<i64, Keyword>) {
43        for child in children.values() {
44            self.add_child(child);
45        }
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use std::collections::BTreeMap;
52
53    use super::{Keyword, KeywordTree};
54
55    #[test]
56    fn keyword_tree_test() {
57        let mut keywords: BTreeMap<i64, Keyword> = BTreeMap::new();
58        keywords.insert(1, Keyword::new(1, 0, "", ""));
59        keywords.insert(2, Keyword::new(2, 1, "", ""));
60        keywords.insert(3, Keyword::new(3, 2, "", ""));
61        keywords.insert(4, Keyword::new(4, 0, "", ""));
62        keywords.insert(5, Keyword::new(5, 2, "", ""));
63
64        let mut tree = KeywordTree::new();
65        tree.add_children(&keywords);
66
67        assert_eq!(tree.map.len(), 3);
68        assert_eq!(tree.map[&0].len(), 2);
69        assert_eq!(tree.map[&1].len(), 1);
70        assert_eq!(tree.map[&2].len(), 2);
71
72        let children = tree.children_for(0);
73        assert_eq!(children, vec![1, 4]);
74    }
75}
lrcat/keywordtree.rs

lrcat/
keywordtree.rs
