Add traverse functions for NodeMut

Author: remimimimimi

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<F>(&mut self, f: F) -> Option<NodeMut<T>>
     where
         F: FnOnce(&mut Node<T>) -> Option<NodeId>,
@@ -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<F>(&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<F>(&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<F>(&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<T> {
         let id = self.tree.orphan(value).id;

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');