evobench_tools/util/
tree.rs1use 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 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}