mutable recursive traversal and recursive_set method

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "orx-tree"
-version = "1.4.0"
+version = "1.5.0"
 edition = "2024"
 authors = ["orxfun <[email protected]>"]
 description = "A beautiful tree 🌳 with convenient and efficient growth, mutation and traversal features."

README.md

@@ -375,6 +375,10 @@ let remaining_bfs: Vec<_> = tree.root().walk::<Bfs>().copied().collect();
 assert_eq!(remaining_bfs, [1, 3, 6, 9]);
 ```
 
+### More Examples
+
+* [mutable_recursive_traversal](https://github.com/orxfun/orx-tree/blob/main/examples/mutable_recursive_traversal.rs) demonstrates different approaches to achieve a recursive mutation of all nodes in the tree.
+
 ## Contributing
 
 Contributions are welcome! If you notice an error, have a question or think something could be added or improved, please open an [issue](https://github.com/orxfun/orx-tree/issues/new) or create a PR.

examples/mutable_recursive_traversal.rs

@@ -0,0 +1,238 @@
+// # EXAMPLE DEFINITION
+//
+// cargo run --example mutable_recursive_traversal
+//
+// This example demonstrates a use case where value of a node is defined
+// as a function of the values of its children. Since the value of a child
+// of the node also depends on values of its own children, it follows that
+// the value of a node is a function of values of all of its descendants.
+//
+// The task is to compute and set all values of a tree given the values of
+// the leaves.
+//
+// This is a interesting and common case in terms of requiring mutable
+// recursive traversal over the tree that can be handled with different
+// approaches. Some of these are demonstrated in this example.
+
+use orx_tree::*;
+use std::fmt::Display;
+
+#[derive(Debug, Clone, Copy)]
+enum Instruction {
+    Input(usize),
+    Add,
+    AddI { val: f32 },
+}
+
+impl Display for Instruction {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Input(x) => write!(f, "Input({})", x),
+            Self::Add => write!(f, "Add"),
+            Self::AddI { val } => write!(f, "AddI({})", val),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct InstructionNode {
+    instruction: Instruction,
+    value: f32,
+}
+
+impl InstructionNode {
+    fn new(instruction: Instruction, value: f32) -> Self {
+        Self { instruction, value }
+    }
+}
+
+impl Display for InstructionNode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.instruction {
+            Instruction::Input(x) => write!(f, "Input({}) => {}", x, self.value),
+            Instruction::Add => write!(f, "Add => {}", self.value),
+            Instruction::AddI { val } => write!(f, "AddI({}) => {}", val, self.value),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct Instructions {
+    tree: DynTree<InstructionNode>,
+}
+
+impl Instructions {
+    fn example() -> Self {
+        let mut tree = DynTree::new(InstructionNode::new(Instruction::AddI { val: 100.0 }, 0.0));
+
+        let mut n0 = tree.root_mut();
+        let [n1, n2] = n0.push_children([
+            InstructionNode::new(Instruction::Input(1), 0.0),
+            InstructionNode::new(Instruction::AddI { val: 2.0 }, 0.0),
+        ]);
+        let _n3 = tree
+            .node_mut(&n1)
+            .push_child(InstructionNode::new(Instruction::Input(0), 0.0));
+        let [_n4, _n5] = tree.node_mut(&n2).push_children([
+            InstructionNode::new(Instruction::Add, 0.0),
+            InstructionNode::new(Instruction::AddI { val: 5.0 }, 0.0),
+        ]);
+
+        Self { tree }
+    }
+}
+
+/// Demonstrates manual mutable and recursive traversal over the tree.
+///
+/// Notice that we can freely walk the tree while always having a single
+/// mutable reference to one node. This satisfies the borrow checker rules
+/// and further allows for calling the function recursively.
+///
+/// Note also that, although it is not necessary in this scenario, we are
+/// free to change the shape of the tree during our walk by adding nodes,
+/// moving around or pruning subtrees, etc. In other words, it enables the
+/// greatest freedom while it requires us to make sure that we do not have
+/// errors, such as out-of-bounds errors with the `into_child_mut` call.
+///
+/// * Pros
+///   * Complete freedom to mutate the nodes and the tree structure during
+///     the walk.
+///   * No intermediate allocation is required; borrow checker rules are
+///     satisfied without the need to collect indices.
+/// * Cons
+///   * Implementor is required to define the walk. This example demonstrates
+///     a depth-first walk due to the recursive calls, which is straightforward
+///     to implement.
+///   * Due to lack of tail-call optimization in rust, this function is likely
+///     to encounter stack overflow for very deep trees.
+fn recursive_traversal_over_nodes<'a>(
+    inputs: &[f32],
+    mut node: NodeMut<'a, Dyn<InstructionNode>>,
+) -> (NodeMut<'a, Dyn<InstructionNode>>, f32) {
+    let num_children = node.num_children();
+
+    let mut children_sum = 0.0;
+    for i in 0..num_children {
+        let child = node.into_child_mut(i).unwrap();
+        let (child, child_value) = recursive_traversal_over_nodes(inputs, child);
+        children_sum += child_value;
+        node = child.into_parent_mut().unwrap();
+    }
+
+    let new_value = match node.data().instruction {
+        Instruction::Input(i) => inputs[i],
+        Instruction::Add => children_sum,
+        Instruction::AddI { val } => val + children_sum,
+    };
+
+    (*node.data_mut()).value = new_value;
+
+    (node, new_value)
+}
+
+/// Demonstrates recursive mutable traversal by internally collecting and storing
+/// the child node indices.
+///
+/// This simplifies the borrow relations and allows for the recursive calls only
+/// having a single mutable reference to the tree; however, each recursive call
+/// requires an internal allocation.
+///
+/// * Pros
+///   * Complete freedom to mutate the nodes and the tree structure during
+///     the walk.
+/// * Cons
+///   * Requires to collect indices and results into an internal vector for each
+///     recursive call, requiring additional allocation.
+///   * Implementor is required to define the walk. This example demonstrates
+///     a depth-first walk due to the recursive calls, which is straightforward
+///     to implement.
+///   * Due to lack of tail-call optimization in rust, this function is likely
+///     to encounter stack overflow for very deep trees.
+fn recursive_traversal_over_indices(
+    tree: &mut DynTree<InstructionNode>,
+    inputs: &[f32],
+    node_idx: NodeIdx<Dyn<InstructionNode>>,
+) -> f32 {
+    let node = tree.node(&node_idx);
+
+    let children_ids: Vec<_> = node.children().map(|child| child.idx()).collect();
+    let children: Vec<_> = children_ids
+        .into_iter()
+        .map(|node| recursive_traversal_over_indices(tree, inputs, node))
+        .collect();
+
+    let mut node = tree.node_mut(&node_idx);
+
+    let new_value = match node.data().instruction {
+        Instruction::Input(i) => inputs[i],
+        Instruction::Add => children.into_iter().sum(),
+        Instruction::AddI { val } => children.into_iter().sum::<f32>() + val,
+    };
+    (*node.data_mut()).value = new_value;
+
+    new_value
+}
+
+/// Demonstrates the use of [`recursive_set`] method:
+///
+/// *Recursively sets the data of all nodes belonging to the subtree rooted
+/// at this node using the compute_data function.*
+///
+/// This function fits perfectly to this and similar scenarios where we want
+/// to compute values of all nodes of a tree such that the value of a node
+/// depends on the values of all of its descendants, and hence the name
+/// *recursive*.
+///
+/// * Pros
+///   * More expressive in the sense that the implementor only defines how the
+///     value of a node should be computed given its prior value and values of
+///     its children. Iteration is abstracted away.
+///   * Despite the name, the implementation actually does not require recursive
+///     function calls; and hence, can work with trees of arbitrary depth without
+///     the risk of stack overflow. Instead, it internally uses the [`PostOrder`]
+///     traverser.
+/// * Cons
+///   * It only allows to set the data of the nodes; however, does not allow for
+///     structural mutations.
+///
+/// [`recursive_set`]: orx_tree::NodeMut::recursive_set
+/// [`PostOrder`]: orx_tree::PostOrder
+fn recursive_set(inputs: &[f32], mut node: NodeMut<Dyn<InstructionNode>>) {
+    node.recursive_set(|node_data, children_data| {
+        let instruction = node_data.instruction;
+        let children_sum: f32 = children_data.iter().map(|x| x.value).sum();
+        let value = match node_data.instruction {
+            Instruction::Input(i) => inputs[i],
+            Instruction::Add => children_sum,
+            Instruction::AddI { val } => val + children_sum,
+        };
+
+        InstructionNode { instruction, value }
+    });
+}
+
+fn main() {
+    fn test_implementation(method: &str, f: impl FnOnce(&[f32], &mut Instructions)) {
+        let inputs = [10.0, 20.0];
+        let mut instructions = Instructions::example();
+        println!("\n\n### {}", method);
+        f(&inputs, &mut instructions);
+        println!("\n{}\n", &instructions.tree);
+    }
+
+    test_implementation(
+        "recursive_traversal_over_indices",
+        |inputs, instructions| {
+            let root_idx = instructions.tree.root().idx();
+            recursive_traversal_over_indices(&mut instructions.tree, inputs, root_idx);
+        },
+    );
+
+    test_implementation("recursive_traversal_over_nodes", |inputs, instructions| {
+        recursive_traversal_over_nodes(&inputs, instructions.tree.root_mut());
+    });
+
+    test_implementation("recursive_set", |inputs, instructions| {
+        recursive_set(inputs, instructions.tree.root_mut());
+    });
+}