|
| 1 | +// # EXAMPLE DEFINITION |
| 2 | +// |
| 3 | +// cargo run --example mutable_recursive_traversal |
| 4 | +// |
| 5 | +// This example demonstrates a use case where value of a node is defined |
| 6 | +// as a function of the values of its children. Since the value of a child |
| 7 | +// of the node also depends on values of its own children, it follows that |
| 8 | +// the value of a node is a function of values of all of its descendants. |
| 9 | +// |
| 10 | +// The task is to compute and set all values of a tree given the values of |
| 11 | +// the leaves. |
| 12 | +// |
| 13 | +// This is a interesting and common case in terms of requiring mutable |
| 14 | +// recursive traversal over the tree that can be handled with different |
| 15 | +// approaches. Some of these are demonstrated in this example. |
| 16 | + |
| 17 | +use orx_tree::*; |
| 18 | +use std::fmt::Display; |
| 19 | + |
| 20 | +#[derive(Debug, Clone, Copy)] |
| 21 | +enum Instruction { |
| 22 | + Input(usize), |
| 23 | + Add, |
| 24 | + AddI { val: f32 }, |
| 25 | +} |
| 26 | + |
| 27 | +impl Display for Instruction { |
| 28 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 29 | + match self { |
| 30 | + Self::Input(x) => write!(f, "Input({})", x), |
| 31 | + Self::Add => write!(f, "Add"), |
| 32 | + Self::AddI { val } => write!(f, "AddI({})", val), |
| 33 | + } |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +#[derive(Debug)] |
| 38 | +struct InstructionNode { |
| 39 | + instruction: Instruction, |
| 40 | + value: f32, |
| 41 | +} |
| 42 | + |
| 43 | +impl InstructionNode { |
| 44 | + fn new(instruction: Instruction, value: f32) -> Self { |
| 45 | + Self { instruction, value } |
| 46 | + } |
| 47 | +} |
| 48 | + |
| 49 | +impl Display for InstructionNode { |
| 50 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 51 | + match &self.instruction { |
| 52 | + Instruction::Input(x) => write!(f, "Input({}) => {}", x, self.value), |
| 53 | + Instruction::Add => write!(f, "Add => {}", self.value), |
| 54 | + Instruction::AddI { val } => write!(f, "AddI({}) => {}", val, self.value), |
| 55 | + } |
| 56 | + } |
| 57 | +} |
| 58 | + |
| 59 | +#[derive(Debug)] |
| 60 | +struct Instructions { |
| 61 | + tree: DynTree<InstructionNode>, |
| 62 | +} |
| 63 | + |
| 64 | +impl Instructions { |
| 65 | + fn example() -> Self { |
| 66 | + let mut tree = DynTree::new(InstructionNode::new(Instruction::AddI { val: 100.0 }, 0.0)); |
| 67 | + |
| 68 | + let mut n0 = tree.root_mut(); |
| 69 | + let [n1, n2] = n0.push_children([ |
| 70 | + InstructionNode::new(Instruction::Input(1), 0.0), |
| 71 | + InstructionNode::new(Instruction::AddI { val: 2.0 }, 0.0), |
| 72 | + ]); |
| 73 | + let _n3 = tree |
| 74 | + .node_mut(&n1) |
| 75 | + .push_child(InstructionNode::new(Instruction::Input(0), 0.0)); |
| 76 | + let [_n4, _n5] = tree.node_mut(&n2).push_children([ |
| 77 | + InstructionNode::new(Instruction::Add, 0.0), |
| 78 | + InstructionNode::new(Instruction::AddI { val: 5.0 }, 0.0), |
| 79 | + ]); |
| 80 | + |
| 81 | + Self { tree } |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +/// Demonstrates manual mutable and recursive traversal over the tree. |
| 86 | +/// |
| 87 | +/// Notice that we can freely walk the tree while always having a single |
| 88 | +/// mutable reference to one node. This satisfies the borrow checker rules |
| 89 | +/// and further allows for calling the function recursively. |
| 90 | +/// |
| 91 | +/// Note also that, although it is not necessary in this scenario, we are |
| 92 | +/// free to change the shape of the tree during our walk by adding nodes, |
| 93 | +/// moving around or pruning subtrees, etc. In other words, it enables the |
| 94 | +/// greatest freedom while it requires us to make sure that we do not have |
| 95 | +/// errors, such as out-of-bounds errors with the `into_child_mut` call. |
| 96 | +/// |
| 97 | +/// * Pros |
| 98 | +/// * Complete freedom to mutate the nodes and the tree structure during |
| 99 | +/// the walk. |
| 100 | +/// * No intermediate allocation is required; borrow checker rules are |
| 101 | +/// satisfied without the need to collect indices. |
| 102 | +/// * Cons |
| 103 | +/// * Implementor is required to define the walk. This example demonstrates |
| 104 | +/// a depth-first walk due to the recursive calls, which is straightforward |
| 105 | +/// to implement. |
| 106 | +/// * Due to lack of tail-call optimization in rust, this function is likely |
| 107 | +/// to encounter stack overflow for very deep trees. |
| 108 | +fn recursive_traversal_over_nodes<'a>( |
| 109 | + inputs: &[f32], |
| 110 | + mut node: NodeMut<'a, Dyn<InstructionNode>>, |
| 111 | +) -> (NodeMut<'a, Dyn<InstructionNode>>, f32) { |
| 112 | + let num_children = node.num_children(); |
| 113 | + |
| 114 | + let mut children_sum = 0.0; |
| 115 | + for i in 0..num_children { |
| 116 | + let child = node.into_child_mut(i).unwrap(); |
| 117 | + let (child, child_value) = recursive_traversal_over_nodes(inputs, child); |
| 118 | + children_sum += child_value; |
| 119 | + node = child.into_parent_mut().unwrap(); |
| 120 | + } |
| 121 | + |
| 122 | + let new_value = match node.data().instruction { |
| 123 | + Instruction::Input(i) => inputs[i], |
| 124 | + Instruction::Add => children_sum, |
| 125 | + Instruction::AddI { val } => val + children_sum, |
| 126 | + }; |
| 127 | + |
| 128 | + (*node.data_mut()).value = new_value; |
| 129 | + |
| 130 | + (node, new_value) |
| 131 | +} |
| 132 | + |
| 133 | +/// Demonstrates recursive mutable traversal by internally collecting and storing |
| 134 | +/// the child node indices. |
| 135 | +/// |
| 136 | +/// This simplifies the borrow relations and allows for the recursive calls only |
| 137 | +/// having a single mutable reference to the tree; however, each recursive call |
| 138 | +/// requires an internal allocation. |
| 139 | +/// |
| 140 | +/// * Pros |
| 141 | +/// * Complete freedom to mutate the nodes and the tree structure during |
| 142 | +/// the walk. |
| 143 | +/// * Cons |
| 144 | +/// * Requires to collect indices and results into an internal vector for each |
| 145 | +/// recursive call, requiring additional allocation. |
| 146 | +/// * Implementor is required to define the walk. This example demonstrates |
| 147 | +/// a depth-first walk due to the recursive calls, which is straightforward |
| 148 | +/// to implement. |
| 149 | +/// * Due to lack of tail-call optimization in rust, this function is likely |
| 150 | +/// to encounter stack overflow for very deep trees. |
| 151 | +fn recursive_traversal_over_indices( |
| 152 | + tree: &mut DynTree<InstructionNode>, |
| 153 | + inputs: &[f32], |
| 154 | + node_idx: NodeIdx<Dyn<InstructionNode>>, |
| 155 | +) -> f32 { |
| 156 | + let node = tree.node(&node_idx); |
| 157 | + |
| 158 | + let children_ids: Vec<_> = node.children().map(|child| child.idx()).collect(); |
| 159 | + let children: Vec<_> = children_ids |
| 160 | + .into_iter() |
| 161 | + .map(|node| recursive_traversal_over_indices(tree, inputs, node)) |
| 162 | + .collect(); |
| 163 | + |
| 164 | + let mut node = tree.node_mut(&node_idx); |
| 165 | + |
| 166 | + let new_value = match node.data().instruction { |
| 167 | + Instruction::Input(i) => inputs[i], |
| 168 | + Instruction::Add => children.into_iter().sum(), |
| 169 | + Instruction::AddI { val } => children.into_iter().sum::<f32>() + val, |
| 170 | + }; |
| 171 | + (*node.data_mut()).value = new_value; |
| 172 | + |
| 173 | + new_value |
| 174 | +} |
| 175 | + |
| 176 | +/// Demonstrates the use of [`recursive_set`] method: |
| 177 | +/// |
| 178 | +/// *Recursively sets the data of all nodes belonging to the subtree rooted |
| 179 | +/// at this node using the compute_data function.* |
| 180 | +/// |
| 181 | +/// This function fits perfectly to this and similar scenarios where we want |
| 182 | +/// to compute values of all nodes of a tree such that the value of a node |
| 183 | +/// depends on the values of all of its descendants, and hence the name |
| 184 | +/// *recursive*. |
| 185 | +/// |
| 186 | +/// * Pros |
| 187 | +/// * More expressive in the sense that the implementor only defines how the |
| 188 | +/// value of a node should be computed given its prior value and values of |
| 189 | +/// its children. Iteration is abstracted away. |
| 190 | +/// * Despite the name, the implementation actually does not require recursive |
| 191 | +/// function calls; and hence, can work with trees of arbitrary depth without |
| 192 | +/// the risk of stack overflow. Instead, it internally uses the [`PostOrder`] |
| 193 | +/// traverser. |
| 194 | +/// * Cons |
| 195 | +/// * It only allows to set the data of the nodes; however, does not allow for |
| 196 | +/// structural mutations. |
| 197 | +/// |
| 198 | +/// [`recursive_set`]: orx_tree::NodeMut::recursive_set |
| 199 | +/// [`PostOrder`]: orx_tree::PostOrder |
| 200 | +fn recursive_set(inputs: &[f32], mut node: NodeMut<Dyn<InstructionNode>>) { |
| 201 | + node.recursive_set(|node_data, children_data| { |
| 202 | + let instruction = node_data.instruction; |
| 203 | + let children_sum: f32 = children_data.iter().map(|x| x.value).sum(); |
| 204 | + let value = match node_data.instruction { |
| 205 | + Instruction::Input(i) => inputs[i], |
| 206 | + Instruction::Add => children_sum, |
| 207 | + Instruction::AddI { val } => val + children_sum, |
| 208 | + }; |
| 209 | + |
| 210 | + InstructionNode { instruction, value } |
| 211 | + }); |
| 212 | +} |
| 213 | + |
| 214 | +fn main() { |
| 215 | + fn test_implementation(method: &str, f: impl FnOnce(&[f32], &mut Instructions)) { |
| 216 | + let inputs = [10.0, 20.0]; |
| 217 | + let mut instructions = Instructions::example(); |
| 218 | + println!("\n\n### {}", method); |
| 219 | + f(&inputs, &mut instructions); |
| 220 | + println!("\n{}\n", &instructions.tree); |
| 221 | + } |
| 222 | + |
| 223 | + test_implementation( |
| 224 | + "recursive_traversal_over_indices", |
| 225 | + |inputs, instructions| { |
| 226 | + let root_idx = instructions.tree.root().idx(); |
| 227 | + recursive_traversal_over_indices(&mut instructions.tree, inputs, root_idx); |
| 228 | + }, |
| 229 | + ); |
| 230 | + |
| 231 | + test_implementation("recursive_traversal_over_nodes", |inputs, instructions| { |
| 232 | + recursive_traversal_over_nodes(&inputs, instructions.tree.root_mut()); |
| 233 | + }); |
| 234 | + |
| 235 | + test_implementation("recursive_set", |inputs, instructions| { |
| 236 | + recursive_set(inputs, instructions.tree.root_mut()); |
| 237 | + }); |
| 238 | +} |
0 commit comments