evobench_tools/util/
tree.rs

1//! A simple, heap allocation based hence slow, string key based tree
2//! representation.
3
4use std::collections::{BTreeMap, btree_map::Entry};
5
6#[derive(Debug)]
7pub struct Tree<'t, T> {
8    pub value: Option<T>,
9    pub children: BTreeMap<&'t str, Tree<'t, T>>,
10}
11
12impl<'key, T> Tree<'key, T> {
13    pub fn new() -> Self {
14        Self {
15            value: None,
16            children: Default::default(),
17        }
18    }
19
20    pub fn get(&self, mut path: impl Iterator<Item = &'key str>) -> Option<&T> {
21        if let Some(key) = path.next() {
22            self.children.get(key)?.get(path)
23        } else {
24            self.value.as_ref()
25        }
26    }
27
28    pub fn insert(&mut self, mut path: impl Iterator<Item = &'key str>, value: T) {
29        if let Some(key) = path.next() {
30            match self.children.entry(key) {
31                Entry::Vacant(vacant_entry) => {
32                    let mut node = Self {
33                        value: None,
34                        children: Default::default(),
35                    };
36                    node.insert(path, value);
37                    vacant_entry.insert(node);
38                }
39                Entry::Occupied(mut occupied_entry) => {
40                    occupied_entry.get_mut().insert(path, value);
41                }
42            }
43        } else {
44            self.value = Some(value);
45        }
46    }
47
48    pub fn from_key_val(
49        data: impl IntoIterator<Item = (impl IntoIterator<Item = &'key str>, T)>,
50    ) -> Self {
51        let mut top = Self::new();
52        for (key, val) in data {
53            top.insert(key.into_iter(), val);
54        }
55        top
56    }
57
58    pub fn into_map_values<U>(self, f: &impl Fn(T) -> U) -> Tree<'key, U> {
59        let Self { value, children } = self;
60        Tree {
61            value: value.map(f),
62            children: children
63                .into_iter()
64                .map(|(key, child)| (key, child.into_map_values(f)))
65                .collect(),
66        }
67    }
68
69    pub fn map_values<U>(&self, f: &impl Fn(&T) -> U) -> Tree<'key, U> {
70        let Self { value, children } = self;
71        Tree {
72            value: value.as_ref().map(f),
73            children: children
74                .iter()
75                .map(|(key, child)| (*key, child.map_values(f)))
76                .collect(),
77        }
78    }
79
80    /// Return a sequence of path, value pairs, sorted by path, the
81    /// path built by joining the keys with `separator`.
82    pub fn into_joined_key_val(self, separator: &str) -> Vec<(String, T)> {
83        let mut path: Vec<&str> = Vec::new();
84        let mut out = Vec::new();
85        into_joined_key_val(self, separator, &mut path, &mut out);
86        out
87    }
88}
89
90fn into_joined_key_val<'key, T>(
91    tree: Tree<'key, T>,
92    separator: &'key str,
93    path: &mut Vec<&'key str>,
94    out: &mut Vec<(String, T)>,
95) {
96    let Tree { value, children } = tree;
97    if let Some(value) = value {
98        out.push((path.join(separator), value));
99    }
100    for (key, child) in children.into_iter() {
101        path.push(key);
102        into_joined_key_val(child, separator, path, out);
103        path.pop();
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn t_forth_and_back() {
113        let vals = &[("a", 2), ("a:b", 1), ("c:d", 3), ("d:e:f", 4), ("d", 5)];
114        let mut tree = Tree::new();
115        for (key, val) in vals {
116            tree.insert(key.split(':'), *val);
117        }
118        let vals2 = tree.into_joined_key_val("--");
119
120        let s = |s: &str| s.to_string();
121        assert_eq!(
122            vals2,
123            [
124                (s("a"), 2),
125                (s("a--b"), 1),
126                (s("c--d"), 3),
127                (s("d"), 5),
128                (s("d--e--f"), 4),
129            ]
130        );
131    }
132}