Added iterator over mutable children

Author: Quotique

src/iter.rs

@@ -2,7 +2,7 @@ use std::iter::FusedIterator;
 use std::ops::Range;
 use std::{slice, vec};
 
-use crate::{Node, NodeId, NodeRef, Tree};
+use crate::{Node, NodeId, NodeMut, NodeRef, Tree};
 
 /// Iterator that moves out of a tree in insert order.
 #[derive(Debug)]
@@ -216,6 +216,52 @@ impl<'a, T: 'a> DoubleEndedIterator for Children<'a, T> {
     }
 }
 
+/// Iterator over mutable children.
+#[derive(Debug)]
+pub struct ChildrenMut<'a, T: 'a> {
+    tree: &'a mut Tree<T>,
+    front: Option<NodeId>,
+    back: Option<NodeId>,
+}
+impl<'a, T: 'a> FusedIterator for ChildrenMut<'a, T> {}
+impl<'a, T: 'a> Iterator for ChildrenMut<'a, T> {
+    type Item = NodeMut<'a, T>;
+    fn next(&mut self) -> Option<Self::Item> {
+        // lifetime cast here is safe because &mut self living shorter than 'a
+        let tree = unsafe { &mut *(self.tree as *mut Tree<T>) };
+        if self.front == self.back {
+            let node = self.front.take().and_then(|x| tree.get_mut(x));
+            self.back = None;
+            node
+        } else {
+            let mut node = self.front.take().and_then(|x| tree.get_mut(x));
+            self.front = node
+                .as_mut()
+                .and_then(|x: &mut NodeMut<T>| x.next_sibling())
+                .map(|x| x.id());
+            node
+        }
+    }
+}
+impl<'a, T: 'a> DoubleEndedIterator for ChildrenMut<'a, T> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        // lifetime cast here is safe because &mut self living shorter than 'a
+        let tree = unsafe { &mut *(self.tree as *mut Tree<T>) };
+        if self.front == self.back {
+            let node = self.back.take().and_then(|x| tree.get_mut(x));
+            self.front = None;
+            node
+        } else {
+            let mut node = self.back.take().and_then(|x| tree.get_mut(x));
+            self.back = node
+                .as_mut()
+                .and_then(|x: &mut NodeMut<T>| x.prev_sibling())
+                .map(|x| x.id());
+            node
+        }
+    }
+}
+
 /// Open or close edge of a node.
 #[derive(Debug)]
 pub enum Edge<'a, T: 'a> {
@@ -354,3 +400,15 @@ impl<'a, T: 'a> NodeRef<'a, T> {
         Descendants(self.traverse())
     }
 }
+
+impl<'a, T: 'a> NodeMut<'a, T> {
+    /// Returns an iterator over mutable children.
+    pub fn children_mut<'b>(&'b mut self) -> ChildrenMut<'b, T> {
+        let (front, back) = self.node().children.unzip();
+        ChildrenMut {
+            tree: self.tree,
+            front,
+            back,
+        }
+    }
+}

tests/iter.rs

@@ -164,6 +164,55 @@ fn last_children() {
     );
 }
 
+#[test]
+fn children_mut() {
+    let mut tree = tree!( 1 => { 2, 3, 4 } );
+    for (num, mut child) in tree.root_mut().children_mut().enumerate() {
+        *child.value() += 10 + num;
+    }
+    assert_eq!(
+        vec![&12, &14, &16],
+        tree.root()
+            .children()
+            .map(|n| n.value())
+            .collect::<Vec<_>>()
+    );
+}
+
+#[test]
+fn children_mut_fused() {
+    let mut tree = tree!( 1 => { 2, 3, 4 } );
+    let mut root = tree.root_mut();
+    let mut children_mut = root.children_mut();
+    for mut child in children_mut.by_ref() {
+        *child.value() += 10;
+    }
+
+    assert!(children_mut.next().is_none());
+    assert_eq!(
+        vec![&12, &13, &14],
+        tree.root()
+            .children()
+            .map(|n| n.value())
+            .collect::<Vec<_>>()
+    );
+}
+
+#[test]
+fn children_mut_rev() {
+    let mut tree = tree!( 1 => { 2, 3, 4 } );
+    for (num, mut child) in tree.root_mut().children_mut().rev().enumerate() {
+        *child.value() += 10 + num;
+    }
+    assert_eq!(
+        vec![&14, &14, &14],
+        tree.root()
+            .children()
+            .map(|n| n.value())
+            .collect::<Vec<_>>()
+    );
+}
+
 #[test]
 fn traverse() {
     use ego_tree::iter::Edge;