1use std::collections::HashMap;
16use std::error::Error;
17use std::fmt;
18use std::ops::Deref;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct Node<T> {
29 id: u64,
30 depth: usize,
31 value: T,
32 children: Vec<Node<T>>,
33}
34
35impl<T> Node<T> {
36 fn new(id: u64, depth: usize, value: T) -> Self {
37 Self {
38 id,
39 depth,
40 value,
41 children: Vec::new(),
42 }
43 }
44
45 pub fn id(&self) -> u64 {
47 self.id
48 }
49
50 pub fn depth(&self) -> usize {
52 self.depth
53 }
54
55 pub fn value(&self) -> &T {
57 &self.value
58 }
59
60 pub fn value_mut(&mut self) -> &mut T {
62 &mut self.value
63 }
64
65 pub fn children(&self) -> &[Node<T>] {
67 &self.children
68 }
69
70 pub fn children_mut(&mut self) -> &mut Vec<Node<T>> {
72 &mut self.children
73 }
74
75 pub fn child_count(&self) -> usize {
77 self.children.len()
78 }
79
80 pub fn is_leaf(&self) -> bool {
82 self.children.is_empty()
83 }
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct ParentNotFound(u64);
94
95impl ParentNotFound {
96 pub fn parent_id(&self) -> u64 {
98 self.0
99 }
100}
101
102impl fmt::Display for ParentNotFound {
103 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104 write!(f, "parent node {} not found in tree", self.0)
105 }
106}
107
108impl Error for ParentNotFound {}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub enum PrunePolicy {
117 KeepLastN(usize),
119 KeepDepthUnder(usize),
121}
122
123#[derive(Debug, Clone)]
134pub struct Tree<T> {
135 root: Node<T>,
136 next_id: u64,
137 index: HashMap<u64, Vec<usize>>,
138}
139
140impl<T> Tree<T> {
141 pub fn new(root_value: T) -> Self {
143 let root = Node::new(0, 0, root_value);
144 let mut index = HashMap::new();
145 index.insert(0, Vec::new());
146 Self {
147 root,
148 next_id: 1,
149 index,
150 }
151 }
152
153 pub fn root(&self) -> &Node<T> {
155 &self.root
156 }
157
158 pub fn root_mut(&mut self) -> &mut Node<T> {
160 &mut self.root
161 }
162
163 pub fn next_id(&self) -> u64 {
165 self.next_id
166 }
167
168 pub fn push_child(&mut self, parent_id: u64, value: T) -> Result<u64, ParentNotFound> {
173 let parent_path = self
174 .index
175 .get(&parent_id)
176 .cloned()
177 .ok_or(ParentNotFound(parent_id))?;
178
179 let depth = self
180 .node_at_path(&parent_path)
181 .expect("index points to existing node")
182 .depth()
183 + 1;
184 let child_index = self
185 .node_at_path(&parent_path)
186 .expect("index points to existing node")
187 .child_count();
188
189 let child_id = self.next_id;
190 let parent = self
191 .node_mut_at_path(&parent_path)
192 .expect("index points to existing node");
193 parent.children.push(Node::new(child_id, depth, value));
194 self.next_id += 1;
195
196 let mut child_path = parent_path;
197 child_path.push(child_index);
198 self.index.insert(child_id, child_path);
199
200 Ok(child_id)
201 }
202
203 pub fn find(&self, id: u64) -> Option<&Node<T>> {
205 let path = self.index.get(&id)?;
206 self.node_at_path(path)
207 }
208
209 pub fn find_mut(&mut self, id: u64) -> Option<&mut Node<T>> {
211 let path = self.index.get(&id)?.clone();
212 self.node_mut_at_path(&path)
213 }
214
215 pub fn find_path_to(&self, id: u64) -> Option<Vec<usize>> {
217 self.index.get(&id).cloned()
218 }
219
220 pub fn dfs(&self) -> DfsIter<'_, T> {
222 DfsIter {
223 stack: vec![(0, &self.root)],
224 }
225 }
226
227 pub fn flatten(&self) -> Vec<(usize, &Node<T>)> {
229 self.dfs().collect()
230 }
231
232 pub fn remove_subtree(&mut self, id: u64) -> Option<Node<T>> {
238 let path = self.find_path_to(id)?;
239 if path.is_empty() {
240 return None; }
242 let mut parent_path = path;
243 let child_index = parent_path.pop()?;
244 let parent = self.node_mut_at_path(&parent_path)?;
245 let removed = parent.children.remove(child_index);
246 self.rebuild_index();
247 Some(removed)
248 }
249
250 pub fn prune(&mut self, policy: PrunePolicy) {
256 match policy {
257 PrunePolicy::KeepLastN(n) => Self::prune_keep_last_n(&mut self.root, n),
258 PrunePolicy::KeepDepthUnder(d) => Self::prune_keep_depth_under(&mut self.root, d),
259 }
260 self.rebuild_index();
261 }
262
263 fn node_at_path(&self, path: &[usize]) -> Option<&Node<T>> {
266 let mut node = &self.root;
267 for &i in path {
268 node = node.children.get(i)?;
269 }
270 Some(node)
271 }
272
273 fn node_mut_at_path(&mut self, path: &[usize]) -> Option<&mut Node<T>> {
274 let mut node = &mut self.root;
275 for &i in path {
276 node = node.children.get_mut(i)?;
277 }
278 Some(node)
279 }
280
281 fn rebuild_index(&mut self) {
282 let mut index = HashMap::new();
283 let mut path = Vec::new();
284 Self::rebuild_from_node(&mut self.root, &mut index, &mut path, 0);
285 self.index = index;
286 }
287
288 fn rebuild_from_node(
289 node: &mut Node<T>,
290 index: &mut HashMap<u64, Vec<usize>>,
291 path: &mut Vec<usize>,
292 depth: usize,
293 ) {
294 node.depth = depth;
295 index.insert(node.id, path.clone());
296 for (i, child) in node.children.iter_mut().enumerate() {
297 path.push(i);
298 Self::rebuild_from_node(child, index, path, depth + 1);
299 path.pop();
300 }
301 }
302
303 fn prune_keep_last_n(node: &mut Node<T>, limit: usize) {
304 if node.children.len() > limit {
305 let drop = node.children.len() - limit;
306 node.children.drain(0..drop);
307 }
308 for child in &mut node.children {
309 Self::prune_keep_last_n(child, limit);
310 }
311 }
312
313 fn prune_keep_depth_under(node: &mut Node<T>, limit: usize) {
314 if limit == 0 || node.depth + 1 >= limit {
315 node.children.clear();
316 return;
317 }
318 for child in &mut node.children {
319 Self::prune_keep_depth_under(child, limit);
320 }
321 }
322}
323
324pub struct DfsIter<'a, T> {
332 stack: Vec<(usize, &'a Node<T>)>,
333}
334
335impl<'a, T> Iterator for DfsIter<'a, T> {
336 type Item = (usize, &'a Node<T>);
337
338 fn next(&mut self) -> Option<Self::Item> {
339 let (depth, node) = self.stack.pop()?;
340 for child in node.children.iter().rev() {
341 self.stack.push((depth + 1, child));
342 }
343 Some((depth, node))
344 }
345}
346
347#[derive(Debug, Clone)]
361pub struct CursoredTree<T> {
362 tree: Tree<T>,
363 cursor: Vec<usize>,
364}
365
366impl<T> CursoredTree<T> {
367 pub fn new(root_value: T) -> Self {
369 Self {
370 tree: Tree::new(root_value),
371 cursor: Vec::new(),
372 }
373 }
374
375 pub fn cursor_path(&self) -> &[usize] {
379 &self.cursor
380 }
381
382 pub fn cursor(&self) -> &[usize] {
384 &self.cursor
385 }
386
387 pub fn current(&self) -> &Node<T> {
389 self.tree
390 .node_at_path(&self.cursor)
391 .expect("cursor points to existing node")
392 }
393
394 pub fn current_mut(&mut self) -> &mut Node<T> {
396 let path = self.cursor.clone();
397 self.tree
398 .node_mut_at_path(&path)
399 .expect("cursor points to existing node")
400 }
401
402 pub fn current_id(&self) -> u64 {
404 self.current().id()
405 }
406
407 pub fn has_parent(&self) -> bool {
409 !self.cursor.is_empty()
410 }
411
412 pub fn has_children(&self) -> bool {
414 !self.current().is_leaf()
415 }
416
417 pub fn push(&mut self, value: T) -> u64 {
423 let parent_id = self.current().id();
424 let child_id = self
425 .tree
426 .push_child(parent_id, value)
427 .expect("current node exists");
428 let child_index = self.current().child_count() - 1;
429 self.cursor.push(child_index);
430 child_id
431 }
432
433 pub fn go_parent(&mut self) -> bool {
435 self.cursor.pop().is_some()
436 }
437
438 pub fn go_child(&mut self, index: usize) -> bool {
441 if self.current().children().get(index).is_none() {
442 return false;
443 }
444 self.cursor.push(index);
445 true
446 }
447
448 pub fn go_child_last(&mut self) -> bool {
451 let count = self.current().child_count();
452 if count == 0 {
453 return false;
454 }
455 self.cursor.push(count - 1);
456 true
457 }
458
459 pub fn go_sibling_next(&mut self) -> bool {
462 let Some(current_index) = self.cursor.last().copied() else {
463 return false;
464 };
465 let parent_child_count = {
466 let parent_path = &self.cursor[..self.cursor.len() - 1];
467 match self.tree.node_at_path(parent_path) {
468 Some(p) => p.children().len(),
469 None => return false,
470 }
471 };
472 if current_index + 1 >= parent_child_count {
473 return false;
474 }
475 *self.cursor.last_mut().expect("checked above") += 1;
476 true
477 }
478
479 pub fn go_sibling_prev(&mut self) -> bool {
482 let Some(last) = self.cursor.last_mut() else {
483 return false;
484 };
485 if *last == 0 {
486 return false;
487 }
488 *last -= 1;
489 true
490 }
491
492 pub fn go_to(&mut self, id: u64) -> bool {
495 match self.tree.find_path_to(id) {
496 Some(path) => {
497 self.cursor = path;
498 true
499 }
500 None => false,
501 }
502 }
503
504 pub fn go_root(&mut self) -> bool {
506 if self.cursor.is_empty() {
507 return false;
508 }
509 self.cursor.clear();
510 true
511 }
512
513 pub fn push_child(&mut self, parent_id: u64, value: T) -> Result<u64, ParentNotFound> {
520 self.tree.push_child(parent_id, value)
521 }
522
523 pub fn remove_subtree(&mut self, id: u64) -> Option<Node<T>> {
526 let removed = self.tree.remove_subtree(id);
527 self.repair_cursor();
528 removed
529 }
530
531 pub fn prune(&mut self, policy: PrunePolicy) {
534 let current_id = self.current().id();
535 self.tree.prune(policy);
536 if let Some(path) = self.tree.find_path_to(current_id) {
537 self.cursor = path;
538 } else {
539 self.repair_cursor();
540 }
541 }
542
543 pub fn tree(&self) -> &Tree<T> {
547 &self.tree
548 }
549
550 pub fn tree_mut(&mut self) -> &mut Tree<T> {
555 &mut self.tree
556 }
557
558 fn repair_cursor(&mut self) {
561 while self.tree.node_at_path(&self.cursor).is_none() {
562 if self.cursor.pop().is_none() {
563 break;
564 }
565 }
566 }
567}
568
569impl<T> Deref for CursoredTree<T> {
570 type Target = Tree<T>;
571
572 fn deref(&self) -> &Self::Target {
573 &self.tree
574 }
575}
576
577#[cfg(test)]
582mod tests {
583 use super::*;
584
585 #[test]
588 fn tree_new_creates_root_with_id_zero() {
589 let tree = Tree::new("root");
590 assert_eq!(tree.root().id(), 0);
591 assert_eq!(tree.root().depth(), 0);
592 assert_eq!(tree.root().value(), &"root");
593 assert!(tree.root().is_leaf());
594 assert_eq!(tree.root().child_count(), 0);
595 assert_eq!(tree.next_id(), 1);
596 assert_eq!(tree.find_path_to(0), Some(vec![]));
597 }
598
599 #[test]
600 fn push_child_assigns_incrementing_ids_and_updates_index() {
601 let mut tree = Tree::new("root");
602 let a = tree.push_child(0, "a").unwrap();
603 let b = tree.push_child(0, "b").unwrap();
604 let a1 = tree.push_child(a, "a1").unwrap();
605
606 assert_eq!(a, 1);
607 assert_eq!(b, 2);
608 assert_eq!(a1, 3);
609
610 assert_eq!(tree.find(a).unwrap().depth(), 1);
611 assert_eq!(tree.find(a1).unwrap().depth(), 2);
612 assert_eq!(tree.find_path_to(b), Some(vec![1]));
613 assert_eq!(tree.find_path_to(a1), Some(vec![0, 0]));
614 }
615
616 #[test]
617 fn push_child_returns_error_for_missing_parent() {
618 let mut tree = Tree::new("root");
619 let err = tree.push_child(99, "x").unwrap_err();
620 assert_eq!(err.parent_id(), 99);
621 assert_eq!(err.to_string(), "parent node 99 not found in tree");
622 }
623
624 #[test]
625 fn find_mut_allows_value_mutation() {
626 let mut tree = Tree::new(String::from("root"));
627 let a = tree.push_child(0, String::from("a")).unwrap();
628 tree.find_mut(a).unwrap().value_mut().push_str("_updated");
629 assert_eq!(tree.find(a).unwrap().value(), "a_updated");
630 }
631
632 #[test]
633 fn find_returns_none_for_nonexistent_id() {
634 let tree = Tree::new("root");
635 assert!(tree.find(42).is_none());
636 assert!(tree.find_path_to(42).is_none());
637 }
638
639 #[test]
640 fn dfs_visits_nodes_in_preorder() {
641 let mut tree = Tree::new("r");
642 let a = tree.push_child(0, "a").unwrap();
643 let b = tree.push_child(0, "b").unwrap();
644 tree.push_child(a, "a1").unwrap();
645 tree.push_child(b, "b1").unwrap();
646
647 let order: Vec<(&str, usize)> = tree.dfs().map(|(d, n)| (*n.value(), d)).collect();
648 assert_eq!(
649 order,
650 vec![("r", 0), ("a", 1), ("a1", 2), ("b", 1), ("b1", 2)]
651 );
652 }
653
654 #[test]
655 fn flatten_matches_dfs_output() {
656 let mut tree = Tree::new("r");
657 let a = tree.push_child(0, "a").unwrap();
658 tree.push_child(a, "a1").unwrap();
659
660 let dfs: Vec<_> = tree.dfs().map(|(d, n)| (d, n.id())).collect();
661 let flat: Vec<_> = tree.flatten().iter().map(|(d, n)| (*d, n.id())).collect();
662 assert_eq!(dfs, flat);
663 }
664
665 #[test]
666 fn remove_subtree_compacts_siblings() {
667 let mut tree = Tree::new("root");
668 let a = tree.push_child(0, "a").unwrap();
669 let b = tree.push_child(0, "b").unwrap();
670 let c = tree.push_child(0, "c").unwrap();
671
672 let removed = tree.remove_subtree(b).unwrap();
673 assert_eq!(removed.id(), b);
674 assert!(tree.find(b).is_none());
675 assert_eq!(tree.find_path_to(c), Some(vec![1]));
677 assert_eq!(tree.root().children().len(), 2);
678 assert_eq!(tree.root().children()[0].id(), a);
679 assert_eq!(tree.root().children()[1].id(), c);
680 }
681
682 #[test]
683 fn remove_subtree_returns_none_for_root() {
684 let mut tree = Tree::new("root");
685 assert!(tree.remove_subtree(0).is_none());
686 }
687
688 #[test]
689 fn prune_keep_last_n_drops_oldest_children() {
690 let mut tree = Tree::new("root");
691 tree.push_child(0, "a").unwrap();
692 let b = tree.push_child(0, "b").unwrap();
693 let c = tree.push_child(0, "c").unwrap();
694 tree.push_child(c, "c1").unwrap();
695 let c2 = tree.push_child(c, "c2").unwrap();
696 let c3 = tree.push_child(c, "c3").unwrap();
697
698 tree.prune(PrunePolicy::KeepLastN(2));
699
700 let root_ids: Vec<u64> = tree.root().children().iter().map(Node::id).collect();
701 assert_eq!(root_ids, vec![b, c]);
702 let c_ids: Vec<u64> = tree
703 .find(c)
704 .unwrap()
705 .children()
706 .iter()
707 .map(Node::id)
708 .collect();
709 assert_eq!(c_ids, vec![c2, c3]);
710 }
711
712 #[test]
713 fn prune_keep_depth_under_truncates_deep_branches() {
714 let mut tree = Tree::new("root");
715 let a = tree.push_child(0, "a").unwrap();
716 let a1 = tree.push_child(a, "a1").unwrap();
717 tree.push_child(a1, "a2").unwrap();
718 let b = tree.push_child(0, "b").unwrap();
719
720 tree.prune(PrunePolicy::KeepDepthUnder(2));
721
722 assert!(tree.find(a).is_some());
723 assert!(tree.find(b).is_some());
724 assert!(tree.find(a1).is_none());
725 }
726
727 #[test]
730 fn cursored_push_adds_child_and_moves_cursor() {
731 let mut ct = CursoredTree::new("root");
732 let a = ct.push("a");
733 assert_eq!(ct.current_id(), a);
734 assert_eq!(ct.cursor_path(), &[0]);
735 assert!(ct.has_parent());
736
737 let a1 = ct.push("a1");
738 assert_eq!(ct.current_id(), a1);
739 assert_eq!(ct.cursor_path(), &[0, 0]);
740 }
741
742 #[test]
743 fn cursored_go_parent_stops_at_root() {
744 let mut ct = CursoredTree::new("root");
745 assert!(!ct.go_parent());
746
747 ct.push("child");
748 assert!(ct.go_parent());
749 assert_eq!(ct.current_id(), 0);
750 }
751
752 #[test]
753 fn cursored_go_child_and_go_child_last() {
754 let mut ct = CursoredTree::new("root");
755 ct.push_child(0, "a").unwrap();
756 let b = ct.push_child(0, "b").unwrap();
757 let b1 = ct.push_child(b, "b1").unwrap();
758
759 assert!(ct.go_child_last());
760 assert_eq!(ct.current_id(), b);
761 assert!(ct.go_child_last());
762 assert_eq!(ct.current_id(), b1);
763 assert!(!ct.go_child_last()); }
765
766 #[test]
767 fn cursored_go_to_jumps_to_any_node() {
768 let mut ct = CursoredTree::new("root");
769 let a = ct.push_child(0, "a").unwrap();
770 let b = ct.push_child(0, "b").unwrap();
771 ct.push_child(b, "b1").unwrap();
772
773 assert!(ct.go_to(a));
774 assert_eq!(ct.current_id(), a);
775 assert!(ct.go_to(0));
776 assert_eq!(ct.current_id(), 0);
777 assert!(!ct.go_to(999));
778 }
779
780 #[test]
781 fn cursored_sibling_navigation() {
782 let mut ct = CursoredTree::new("root");
783 ct.push_child(0, "a").unwrap();
784 let b = ct.push_child(0, "b").unwrap();
785 ct.push_child(0, "c").unwrap();
786
787 assert!(ct.go_child(1)); assert_eq!(ct.current_id(), b);
789 assert!(ct.go_sibling_next()); assert!(!ct.go_sibling_next()); assert!(ct.go_sibling_prev()); assert_eq!(ct.current_id(), b);
793 assert!(ct.go_sibling_prev()); assert!(!ct.go_sibling_prev()); }
796
797 #[test]
798 fn cursored_go_root() {
799 let mut ct = CursoredTree::new("root");
800 assert!(!ct.go_root());
801 ct.push("a");
802 ct.push("a1");
803 assert!(ct.go_root());
804 assert_eq!(ct.current_id(), 0);
805 }
806
807 #[test]
808 fn cursored_has_parent_and_has_children() {
809 let mut ct = CursoredTree::new("root");
810 assert!(!ct.has_parent());
811 assert!(!ct.has_children());
812
813 ct.push("child");
814 assert!(ct.has_parent());
815 assert!(!ct.has_children());
816
817 ct.go_parent();
818 assert!(!ct.has_parent());
819 assert!(ct.has_children());
820 }
821
822 #[test]
823 fn cursored_prune_repairs_cursor() {
824 let mut ct = CursoredTree::new("root");
825 let a = ct.push_child(0, "a").unwrap();
826 ct.push_child(a, "a1").unwrap();
827
828 assert!(ct.go_child(0));
829 assert!(ct.go_child(0)); ct.prune(PrunePolicy::KeepDepthUnder(2));
831 assert_eq!(ct.current_id(), a);
832 assert_eq!(ct.cursor(), &[0]);
833 }
834
835 #[test]
836 fn cursored_remove_subtree_repairs_cursor() {
837 let mut ct = CursoredTree::new("root");
838 let a = ct.push_child(0, "a").unwrap();
839 let b = ct.push_child(0, "b").unwrap();
840
841 assert!(ct.go_child(1)); let removed = ct.remove_subtree(b).unwrap();
843 assert_eq!(removed.id(), b);
844 assert_eq!(ct.current_id(), 0);
845 assert!(ct.find(a).is_some());
846 assert!(ct.find(b).is_none());
847 }
848
849 #[test]
850 fn cursored_tree_mut_gives_inner_access() {
851 let mut ct = CursoredTree::new("root");
852 ct.push("a");
853 ct.tree_mut().root_mut().value_mut();
854 assert_eq!(ct.root().child_count(), 1);
855 }
856
857 #[test]
858 fn deref_exposes_tree_methods() {
859 let ct = CursoredTree::new("root");
860 assert_eq!(ct.root().id(), 0);
862 assert!(ct.find(0).is_some());
863 }
864}