example refactoring and documentation for mutable traversal

Author: orxfun

examples/mutable_traversal.rs

@@ -1,8 +1,20 @@
+// # EXAMPLE DEFINITION
+//
+// 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;
 
-const INPUTS_COUNT: usize = 2;
-
 #[derive(Debug, Clone, Copy)]
 enum Instruction {
     Input(usize),
@@ -43,11 +55,11 @@ impl Display for InstructionNode {
 }
 
 #[derive(Debug)]
-struct MyTree {
+struct Instructions {
     tree: DynTree<InstructionNode>,
 }
 
-impl MyTree {
+impl Instructions {
     fn example() -> Self {
         let mut tree = DynTree::new(InstructionNode::new(Instruction::AddI { val: 100.0 }, 0.0));
 
@@ -66,32 +78,32 @@ impl MyTree {
 
         Self { tree }
     }
-
-    fn execute_rec(&mut self, inputs: &[f32], node_idx: NodeIdx<Dyn<InstructionNode>>) -> f32 {
-        let node = self.tree.node(&node_idx);
-
-        let children_ids = node.children().map(|child| child.idx()).collect::<Vec<_>>();
-        let mut children = vec![];
-
-        for node in children_ids {
-            let value = self.execute_rec(inputs, node);
-            children.push(value);
-        }
-
-        let mut node = self.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
-    }
 }
 
-fn execute<'a>(
+/// 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) {
@@ -100,7 +112,7 @@ fn execute<'a>(
     let mut children_sum = 0.0;
     for i in 0..num_children {
         let child = node.into_child_mut(i).unwrap();
-        let (child, child_value) = execute(inputs, child);
+        let (child, child_value) = recursive_traversal_over_nodes(inputs, child);
         children_sum += child_value;
         node = child.into_parent_mut().unwrap();
     }
@@ -116,43 +128,109 @@ fn execute<'a>(
     (node, new_value)
 }
 
-fn test_implementation(method: &str, f: impl FnOnce(&[f32], &mut MyTree)) {
-    let inputs = [10.0, 20.0];
+/// 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 tree = MyTree::example();
+    let mut node = tree.node_mut(&node_idx);
 
-    println!("\n\n# {}", method);
-    println!("\ninputs = {:?}\n", &inputs);
-    println!("Before execute:\n{}\n", &tree.tree);
-    f(&inputs, &mut tree);
-    println!("After execute:\n{}\n", &tree.tree);
+    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
 }
 
-fn main() {
-    test_implementation("IMPL OVER CHILDREN INDICES", |inputs, tree| {
-        tree.execute_rec(inputs, tree.tree.root().idx());
+/// 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(
-        "IMPL WITH INTO_CHILD_MUT & INTO_PARENT_MUT",
-        |inputs, tree| {
-            execute(&inputs, tree.tree.root_mut());
+        "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("IMPL recursive_set", |inputs, tree| {
-        tree.tree
-            .root_mut()
-            .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 }
-            });
+    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());
     });
 }