From 4a8aabb3ce4c125f5da70294bfd611391a3f6fcb Mon Sep 17 00:00:00 2001 From: remimimimimi Date: Thu, 17 Apr 2025 23:51:51 +0300 Subject: [PATCH] Add traverse functions for NodeMut --- src/lib.rs | 115 ++++++++++++++++++++++++++++++++++++++++++++++ tests/node_mut.rs | 71 ++++++++++++++++++++++++++++ 2 files changed, 186 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 13ed4b3..3b42f0f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -369,6 +369,11 @@ impl<'a, T: 'a> NodeMut<'a, T> { &mut self.node().value } + /// Downcast `NodeMut` to `NodeRef`. + pub fn as_ref(&mut self) -> NodeRef<'_, T> { + unsafe { self.tree.get_unchecked(self.id) } + } + fn axis(&mut self, f: F) -> Option> where F: FnOnce(&mut Node) -> Option, @@ -463,6 +468,116 @@ impl<'a, T: 'a> NodeMut<'a, T> { unsafe { self.tree.get_unchecked(self.id).has_children() } } + /// Apply function for each ancestor mutable node reference. + pub fn for_each_ancestor<'b, F>(&'b mut self, mut f: F) + where + F: FnMut(&mut NodeMut<'b, T>), + { + let mut current = self.parent(); + while let Some(mut node) = current { + f(&mut node); + current = node.into_parent().ok(); + } + } + + /// Apply function for each next sibling mutable node reference. + pub fn for_each_next_sibling<'b, F>(&'b mut self, mut f: F) + where + F: FnMut(&mut NodeMut<'b, T>), + { + let mut current = self.next_sibling(); + while let Some(mut node) = current { + f(&mut node); + current = node.into_next_sibling().ok(); + } + } + + /// Apply function for each previout sibling mutable node reference. + pub fn for_each_prev_sibling<'b, F>(&'b mut self, mut f: F) + where + F: FnMut(&mut NodeMut<'b, T>), + { + let mut current = self.prev_sibling(); + while let Some(mut node) = current { + f(&mut node); + current = node.into_prev_sibling().ok(); + } + } + + /// Apply function for this node and each sibling mutable node reference. + pub fn for_each_sibling(&mut self, mut f: F) + where + F: for<'b> FnMut(&mut NodeMut<'b, T>), + { + self.for_each_prev_sibling(&mut f); + f(self); + self.for_each_next_sibling(&mut f); + } + + /// Apply function for each children mutable node reference. + pub fn for_each_child(&mut self, mut f: F) + where + F: for<'b> FnMut(&mut NodeMut<'b, T>), + { + let Some(mut first_child) = self.first_child() else { + return; + }; + f(&mut first_child); + first_child.for_each_next_sibling(f); + } + + /// Apply function for this node and each descendant mutable node reference. + pub fn for_each_descendant(&mut self, mut f: F) + where + F: FnMut(&mut NodeMut<'_, T>), + { + let id = self.id(); + + f(self); + + // Start at our first child, if any. + let Some(mut node) = self.first_child() else { + return; + }; + + loop { + f(&mut node); + + // Try to go deeper into its first child. + match node.into_first_child() { + Ok(child) => { + node = child; + continue; + } + Err(n) => { + node = n; + } + } + + // No deeper child, so climb until we find a next sibling or hit self. + loop { + match node.into_next_sibling() { + Ok(sib) => { + node = sib; + break; + } + Err(n) => { + node = n; + } + } + + // No sibling, so climb up. + let Ok(parent) = node.into_parent() else { + unreachable!(); + }; + if parent.id() == id { + return; + } + node = parent; + } + } + } + /// Appends a new child to this node. pub fn append(&mut self, value: T) -> NodeMut { let id = self.tree.orphan(value).id; diff --git a/tests/node_mut.rs b/tests/node_mut.rs index 8208bf8..a93d30b 100644 --- a/tests/node_mut.rs +++ b/tests/node_mut.rs @@ -74,6 +74,77 @@ fn has_children() { assert!(!tree.root_mut().first_child().unwrap().has_children()); } +#[test] +fn for_each_next_sibling() { + let mut tree = tree!(1 => { 2, 3, 4, 5, 6 }); + let mut root = tree.root_mut(); + let mut c = root.first_child().unwrap(); + + c.for_each_next_sibling(|n| { + *n.value() += 1; + }); + + let res = tree!(1 => { 2, 4, 5, 6, 7 }); + + assert_eq!(tree, res); +} + +#[test] +fn for_each_prev_sibling() { + let mut tree = tree!(1 => { 2, 3, 4, 5, 6 }); + let mut root = tree.root_mut(); + let mut c = root.last_child().unwrap(); + + c.for_each_prev_sibling(|n| { + *n.value() += 1; + }); + + let res = tree!(1 => { 3, 4, 5, 6, 6 }); + + assert_eq!(tree, res); +} + +#[test] +fn for_each_sibling() { + let rt = 2; + let mut tree = tree!(rt => { 2, 3, 4, 5, 6 }); + let mut root = tree.root_mut(); + let mut c = root.last_child().unwrap(); + + c.for_each_sibling(|n| { + let v = n.parent().map(|mut p| *p.value()).unwrap(); + *n.value() += v; + }); + + let res = tree!(rt => { 4, 5, 6, 7, 8 }); + + assert_eq!(tree, res); +} + +#[test] +fn for_each_child() { + let mut tree = tree!(1 => { 2, 3, 4, 5, 6 }); + let mut root = tree.root_mut(); + root.for_each_child(|n| *n.value() += 1); + + assert_eq!(*root.value(), 1); + + let res = tree!(1 => { 3, 4, 5, 6, 7 }); + + assert_eq!(tree, res); +} + +#[test] +fn for_each_descendant() { + let mut tree = tree!(1 => { 2 => {3, 4, 5, 6}, 3, 4 => {0, 1}, 5, 6 => {1, 2} }); + let mut root = tree.root_mut(); + root.for_each_descendant(|n| *n.value() += 1); + + let tree2 = tree!(2 => { 3 => {4, 5, 6, 7}, 4, 5 => {1, 2}, 6, 7 => {2, 3} }); + + assert_eq!(tree, tree2); +} + #[test] fn append_1() { let mut tree = tree!('a');