diff --git a/Cargo.lock b/Cargo.lock index bd47eec..5b74bfc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "crossbeam-deque" @@ -29,7 +29,7 @@ checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" [[package]] name = "easy-tree" -version = "0.1.3" +version = "0.2.0" dependencies = [ "rayon", ] diff --git a/Cargo.toml b/Cargo.toml index c00490e..4db3615 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "easy-tree" -version = "0.1.3" +version = "0.2.0" authors = ["Anton Suprunchuk "] edition = "2021" description = "A simple and efficient tree structure library for Rust with recursive traversal" diff --git a/src/lib.rs b/src/lib.rs index f8da7e9..eb12386 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -213,6 +213,8 @@ impl Node { #[derive(Clone)] pub struct Tree { nodes: Vec>, + /// Stack for traverse_mut to avoid allocations + stack: Vec<(usize, bool)>, } impl Default for Tree { @@ -234,7 +236,10 @@ impl Tree { /// let tree: Tree = Tree::new(); /// ``` pub fn new() -> Self { - Self { nodes: Vec::new() } + Self { + nodes: Vec::new(), + stack: Vec::new(), + } } /// Adds a new node to the tree. @@ -523,6 +528,42 @@ impl Tree { } } + /// Walks the tree recursively, applying the given functions before and after processing the + /// children of each node. This version allows for mutable access to the nodes. + pub fn traverse_mut( + &mut self, + mut before_processing_children: impl FnMut(usize, &mut T, &mut S), + mut after_processing_the_subtree: impl FnMut(usize, &mut T, &mut S), + s: &mut S, + ) { + if self.is_empty() { + return; + } + + self.stack.clear(); + self.stack.push((0, false)); + + while let Some((index, children_visited)) = self.stack.pop() { + if children_visited { + // All children are processed, call f2 + let node = &mut self.nodes[index]; + after_processing_the_subtree(index, &mut node.data, s); + } else { + // Call f and mark this node's children for processing + let node = &mut self.nodes[index]; + before_processing_children(index, &mut node.data, s); + + // Re-push the current node with children_visited set to true + self.stack.push((index, true)); + + // Push all children onto the stack + for &child in node.children.iter().rev() { + self.stack.push((child, false)); + } + } + } + } + /// Returns an iterator over the indices and data of the nodes in the tree. pub fn iter(&self) -> impl Iterator { self.nodes @@ -646,13 +687,10 @@ mod tests { let mut result = vec![]; tree.traverse( - |index, node, result| { - result.push(format!("Calling handler for node {}: {}", index, node)) - }, + |index, node, result| result.push(format!("Calling handler for node {index}: {node}")), |index, _node, result| { result.push(format!( - "Finished handling node {} and all it's children", - index + "Finished handling node {index} and all its children" )) }, &mut result, @@ -664,11 +702,11 @@ mod tests { "Calling handler for node 0: 0", "Calling handler for node 1: 1", "Calling handler for node 3: 3", - "Finished handling node 3 and all it's children", - "Finished handling node 1 and all it's children", + "Finished handling node 3 and all its children", + "Finished handling node 1 and all its children", "Calling handler for node 2: 2", - "Finished handling node 2 and all it's children", - "Finished handling node 0 and all it's children", + "Finished handling node 2 and all its children", + "Finished handling node 0 and all its children", ] ); }