1use std::collections::BTreeMap;
2use std::fmt::{Debug, Display, Error, Formatter};
3use std::hash::Hash;
4use std::ops::Deref;
5use std::sync::{Arc, RwLock, Weak};
6use crate::mixed::{DceErr, DceResult};
7
8pub enum TreeTraverBreak {
9 Stop,
10 Break,
11 Skip,
12 Continue,
13}
14
15pub trait KeyFactory<K> {
16 fn key(&self) -> K;
17 fn child_of(&self, parent: &Self) -> bool;
18}
19
20#[derive(Debug)]
21pub struct ATree<E, K> {
22 element: RwLock<E>,
23 children: RwLock<BTreeMap<K, Arc<ATree<E, K>>>>,
24 parent: Weak<ATree<E, K>>,
25 own: RwLock<Weak<ATree<E, K>>>,
26}
27
28impl<E, K> ATree<E, K>
29where E: PartialEq + Debug,
30 K: Hash + Clone + Ord + Display + Debug
31{
32 pub fn set(&self, key: K, element: E) -> DceResult<Arc<ATree<E, K>>> {
33 let child = Self::new_with_parent(element, Arc::downgrade(&self.own.read()
34 .map_err(DceErr::closed0)?.upgrade().ok_or_else(|| DceErr::closed0("Failed to update Weak to Arc"))?))?;
35 self.children.write().map_err(DceErr::closed0)?.insert(key, child.clone());
36 Ok(child)
37 }
38
39 pub fn set_if_absent(&self, key: K, element: E) -> DceResult<Arc<ATree<E, K>>> {
40 if let Some(exists) = self.get(&key) {
41 return Ok(exists);
42 }
43 self.set(key, element)
44 }
45
46 pub fn set_by_path(&self, path: Vec<K>, element: E) -> DceResult<Arc<ATree<E, K>>> {
47 self.actual_set_by_path(path, element, true)
48 }
49
50 pub fn set_by_path_if_absent(&self, path: Vec<K>, element: E) -> DceResult<Arc<ATree<E, K>>> {
51 self.actual_set_by_path(path, element, false)
52 }
53
54 fn actual_set_by_path(&self, mut path: Vec<K>, element: E, force: bool) -> DceResult<Arc<ATree<E, K>>> {
55 let key = path.pop().ok_or_else(|| DceErr::closed0("Cannot get by an empty path"))?;
56 let parent = self.get_by_path(&path).ok_or_else(|| DceErr::closed0("Parent not found"))?;
57 Ok(if force { parent.set(key, element)? } else { parent.set_if_absent(key, element)? })
58 }
59
60 pub fn get(&self, key: &K) -> Option<Arc<ATree<E, K>>> {
61 return self.children.read().ok()?.get(key).map(Clone::clone);
62 }
63
64 pub fn get_by_path(&self, path: &[K]) -> Option<Arc<ATree<E, K>>> {
65 let mut child = self.own.read().ok()?.upgrade()?;
66 for key in path { child = child.get(key)?; }
67 Some(child)
68 }
69
70 pub fn parent(&self) -> Option<Arc<ATree<E, K>>> {
71 self.parent.upgrade()
72 }
73
74 pub fn parents(&self) -> DceResult<Vec<Arc<ATree<E, K>>>> {
75 self.parents_until(None, true)
76 }
77
78 pub fn parents_until(&self, until: Option<Arc<ATree<E, K>>>, elder_first: bool) -> DceResult<Vec<Arc<ATree<E, K>>>> {
79 let parent = self.own.read().map_err(DceErr::closed0)?.upgrade();
80 if parent.is_none() {
81 return Ok(vec![]);
82 }
83 let mut parent = parent.ok_or_else(|| DceErr::closed0("Failed to update Weak to Arc"))?;
84 let mut parents = vec![parent.clone()];
85 while Some(parent.clone()) != until && match parent.parent() {
86 Some(p) => {
87 parents.push(p.clone());
88 parent = p;
89 true
90 },
91 _ => false
92 } {}
93 if elder_first { parents.reverse(); }
94 Ok(parents)
95 }
96
97 pub fn children(&self) -> &RwLock<BTreeMap<K, Arc<ATree<E, K>>>> {
98 &self.children
99 }
100
101 pub fn contains_key(&self, key: K) -> DceResult<bool> {
102 self.children.read().map_err(DceErr::closed0).map(|r| r.contains_key(&key))
103 }
104
105 pub fn is_empty(&self) -> DceResult<bool> {
106 self.children.read().map_err(DceErr::closed0).map(|r| r.is_empty())
107 }
108
109 pub fn remove(&mut self, key: &K) -> Option<Arc<ATree<E, K>>> {
110 self.children.write().ok()?.remove(key)
111 }
112
113 pub fn traversal(
114 &self,
115 callback: fn(Arc<ATree<E, K>>) -> DceResult<TreeTraverBreak>,
116 ) -> DceResult<()> {
117 let mut nodes = vec![self.own.read().map_err(DceErr::closed0)?.upgrade().ok_or_else(|| DceErr::closed0("Failed to update Weak to Arc"))?];
118 'outer: while let Some(parent) = nodes.pop() {
119 let nodes_len = nodes.len();
120 for child in parent.children.read().map_err(DceErr::closed0)?.values() {
121 match callback(child.clone())? {
122 TreeTraverBreak::Stop => break 'outer,
123 TreeTraverBreak::Break => break,
124 TreeTraverBreak::Skip => continue,
125 _ => nodes.insert(nodes_len.clone(), child.clone()),
126 }
127 }
128 }
129 Ok(())
130 }
131
132 pub fn new(element: E) -> DceResult<Arc<ATree<E, K>>> {
133 Self::new_with_parent(element, Weak::new())
134 }
135
136 fn new_with_parent(element: E, parent: Weak<ATree<E, K>>) -> DceResult<Arc<ATree<E, K>>> {
137 let rc = Arc::new(ATree {
138 element: RwLock::new(element),
139 children: RwLock::new(BTreeMap::new()),
140 parent,
141 own: RwLock::new(Weak::new()),
142 });
143 *rc.own.write().map_err(DceErr::closed0)? = Arc::downgrade(&rc);
144 Ok(rc)
145 }
146
147 pub fn build(
151 &self,
152 mut elements: Vec<E>,
153 remains_handler: Option<fn(&ATree<E, K>, Vec<E>)>,
154 ) -> DceResult<()>
155 where E: KeyFactory<K>,
156 {
157 let mut parents = vec![self.own.write().map_err(DceErr::closed0)?.upgrade().ok_or_else(|| DceErr::closed0("Failed to update Weak to Arc"))?];
158 while let Some(pa) = parents.pop() {
159 for i in (0 .. elements.len()).filter(|i| pa.element.read().map_or(false, |e| elements[*i].child_of(&e))).rev().collect::<Vec<_>>() {
160 let elem = elements.remove(i);
161 parents.push(pa.set(elem.key(), elem)?);
162 }
163 }
164 if let Some(remains_handler) = remains_handler {
165 remains_handler(self, elements);
166 }
167 Ok(())
168 }
169}
170
171impl<E, K> Deref for ATree<E, K> {
172 type Target = RwLock<E>;
173
174 fn deref(&self) -> &Self::Target {
175 &self.element
176 }
177}
178
179impl<E: PartialEq, K> PartialEq for ATree<E, K> {
180 fn eq(&self, other: &Self) -> bool {
181 self.element.read().map_or(false, |s| other.element.read().map_or(false, |o| s.eq(o.deref())))
182 }
183}
184
185impl<E: Display, K> Display for ATree<E, K> {
186 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
187 write!(f, "{}", self.element.read().map_err(|_| Error)?)
188 }
189}
190
191
192#[cfg(test)]
193mod tests {
194 use super::*;
195
196 impl KeyFactory<u8> for (u8, u8, String) {
197 fn key(&self) -> u8 {
198 self.0
199 }
200
201 fn child_of(&self, parent: &Self) -> bool {
202 self.1 == parent.0
203 }
204 }
205
206 #[test]
207 fn build() {
208 let root = ATree::new((0, 0, "x".to_string())).unwrap();
209 root.set(1, (1, 0, "a".to_string())).unwrap();
210 root.set(2, (2, 0, "b".to_string())).unwrap();
211 root.set_by_path(vec![8], (8, 0, "h".to_string())).unwrap();
212 root.set_by_path(vec![1, 3], (3, 1, "c".to_string())).unwrap();
213 root.set_by_path(vec![1, 4], (4, 1, "d".to_string())).unwrap();
214 root.set_by_path(vec![8, 5], (5, 8, "e".to_string())).unwrap();
215 root.set_by_path(vec![8, 5, 6], (6, 5, "f".to_string())).unwrap();
216 root.set_by_path(vec![8, 5, 6], (7, 5, "g".to_string())).unwrap();
217 root.set_by_path(vec![1, 3, 9], (9, 3, "i".to_string())).unwrap();
218 root.set_by_path(vec![1, 3, 10], (10, 3, "j".to_string())).unwrap();
219 println!("{:?}", root);
220 root.traversal(|t| {
221 eprintln!("t = {:?}", **t);
222 Ok(TreeTraverBreak::Continue)
223 }).unwrap();
224 }
225
226
227 impl KeyFactory<&'static str> for &'static str {
228 fn key(&self) -> &'static str {
229 if let Some(index) = self.rfind('/') {
230 return &self[index+1 ..]
231 }
232 self
233 }
234
235 fn child_of(&self, parent: &Self) -> bool {
236 if let Some(index) = self.rfind('/') {
237 &&self[..index] == parent
238 } else {
239 parent.is_empty()
240 }
241 }
242 }
243
244 #[test]
245 fn get_child() {
246 let tree = ATree::new("").unwrap();
247 tree.build(vec![
248 "hello",
249 "hello/world",
250 "hello/world/!",
251 "hello/rust!",
252 "hello/dce/for/rust!",
253 ], Some(|tree: &ATree<&'static str, &'static str>, mut remains: Vec<&'static str>| {
254 let mut fills: BTreeMap<Vec<&'static str>, &'static str> = BTreeMap::new();
255 while let Some(element) = remains.pop() {
256 let paths: Vec<_> = element.split("/").collect();
257 let last_index = paths.len();
258 for i in 0..last_index {
259 let path = paths[..=i].to_vec();
260 if matches!(tree.get_by_path(&path), None) && ! fills.contains_key(&path) {
261 let element = Box::leak(path.clone().join("/").into_boxed_str());
262 fills.insert(path, element);
263 }
264 }
265 }
266 while let Some((paths, nb)) = fills.pop_first() {
267 tree.set_by_path(paths, nb).unwrap();
268 }
269 }),).unwrap();
270 let t = tree.get_by_path(&["hello", "world"]).unwrap();
271 let t2 = t.get(&"!").unwrap();
272 let parents = t2.parents_until(tree.get(&"hello"), true);
273 println!("{:#?}", t);
274 println!("{:#?}", tree);
275 println!("{:#?}", t2);
276 println!("{:#?}", parents);
277 tree.traversal(|t| {
278 eprintln!("t = {:?}", t);
279 Ok(TreeTraverBreak::Continue)
280 }).unwrap();
281 }
282
283}