diff --git a/datafusion/common/src/cse.rs b/datafusion/common/src/cse.rs index ab02915858cd..2c2b032f7087 100644 --- a/datafusion/common/src/cse.rs +++ b/datafusion/common/src/cse.rs @@ -597,7 +597,8 @@ impl> CSE mod test { use crate::alias::AliasGenerator; use crate::cse::{CSEController, HashNode, IdArray, Identifier, NodeStats, CSE}; - use crate::tree_node::tests::TestTreeNode; + use crate::tree_node::tests::test_tree_node::TestTreeNode; + use crate::tree_node::tests::TestTree; use crate::Result; use std::collections::HashSet; use std::hash::{Hash, Hasher}; @@ -670,21 +671,22 @@ mod test { TestTreeNodeMask::Normal, )); - let a_plus_1 = TestTreeNode::new( + let a_plus_1 = TestTreeNode::new_with_children( vec![ TestTreeNode::new_leaf("a".to_string()), TestTreeNode::new_leaf("1".to_string()), ], "+".to_string(), ); - let avg_c = TestTreeNode::new( + let avg_c = TestTreeNode::new_with_children( vec![TestTreeNode::new_leaf("c".to_string())], "avg".to_string(), ); - let sum_a_plus_1 = TestTreeNode::new(vec![a_plus_1], "sum".to_string()); + let sum_a_plus_1 = + TestTreeNode::new_with_children(vec![a_plus_1], "sum".to_string()); let sum_a_plus_1_minus_avg_c = - TestTreeNode::new(vec![sum_a_plus_1, avg_c], "-".to_string()); - let root = TestTreeNode::new( + TestTreeNode::new_with_children(vec![sum_a_plus_1, avg_c], "-".to_string()); + let root = TestTreeNode::new_with_children( vec![ sum_a_plus_1_minus_avg_c, TestTreeNode::new_leaf("2".to_string()), diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index a0ad1e80be9b..3d2169a597ab 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -17,9 +17,9 @@ //! [`TreeNode`] for visiting and rewriting expression and plan trees -use std::sync::Arc; - use crate::Result; +use std::marker::PhantomData; +use std::sync::Arc; /// These macros are used to determine continuation during transforming traversals. macro_rules! handle_transform_recursion { @@ -760,6 +760,12 @@ impl Transformed { TreeNodeRecursion::Jump | TreeNodeRecursion::Stop => Ok(self), } } + + /// Replaces recursion state with the given one + pub fn with_tnr(mut self, tnr: TreeNodeRecursion) -> Self { + self.tnr = tnr; + self + } } /// Transformation helper to process a sequence of iterable tree nodes that are siblings. @@ -921,67 +927,84 @@ impl TransformedResult for Result> { /// Helper trait for implementing [`TreeNode`] that have children stored as /// `Arc`s. If some trait object, such as `dyn T`, implements this trait, -/// its related `Arc` will automatically implement [`TreeNode`]. +/// its related `Arc` will automatically implement [`TreeNode`] and [`ConcreteTreeNode`]. pub trait DynTreeNode { /// Returns all children of the specified `TreeNode`. fn arc_children(&self) -> Vec<&Arc>; /// Constructs a new node with the specified children. fn with_new_arc_children( - &self, - arc_self: Arc, + self: Arc, new_children: Vec>, ) -> Result>; } -/// Blanket implementation for any `Arc` where `T` implements [`DynTreeNode`] -/// (such as [`Arc`]). -impl TreeNode for Arc { - fn apply_children<'n, F: FnMut(&'n Self) -> Result>( - &'n self, - f: F, - ) -> Result { - self.arc_children().into_iter().apply_until_stop(f) +/// Note that this implementation doesn't actually take children from DynTreeNode, since they are stored +/// inside shared references, but instead produces new links to them +impl ConcreteTreeNode for Arc { + fn children(&self) -> Vec<&Self> { + self.arc_children() } - fn map_children Result>>( - self, - f: F, - ) -> Result> { - let children = self.arc_children(); - if !children.is_empty() { - let new_children = children - .into_iter() - .cloned() - .map_until_stop_and_collect(f)?; - // Propagate up `new_children.transformed` and `new_children.tnr` - // along with the node containing transformed children. - if new_children.transformed { - let arc_self = Arc::clone(&self); - new_children.map_data(|new_children| { - self.with_new_arc_children(arc_self, new_children) - }) - } else { - Ok(Transformed::new(self, false, new_children.tnr)) - } - } else { - Ok(Transformed::no(self)) + fn take_children(self) -> (Self, Vec) { + let children = self.children().into_iter().cloned().collect(); + (self, children) + } + + fn with_new_children(self, children: Vec) -> Result { + self.with_new_arc_children(children) + } +} + +/// Adapter from the old function-based rewriter to the new Transformer one +struct FuncRewriter< + FD: FnMut(Node) -> Result>, + FU: FnMut(Node) -> Result>, + Node: TreeNode, +> { + f_down_func: FD, + f_up_func: FU, + _node: PhantomData, +} + +impl< + FD: FnMut(Node) -> Result>, + FU: FnMut(Node) -> Result>, + Node: TreeNode, + > FuncRewriter +{ + pub fn new(f_down_func: FD, f_up_func: FU) -> Self { + Self { + f_down_func, + f_up_func, + _node: PhantomData, } } } +impl< + FD: FnMut(Node) -> Result>, + FU: FnMut(Node) -> Result>, + Node: TreeNode, + > TreeNodeRewriter for FuncRewriter +{ + type Node = Node; -/// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for -/// trees that contain nodes with payloads. This approach ensures safe execution of algorithms -/// involving payloads, by enforcing rules for detaching and reattaching child nodes. -pub trait ConcreteTreeNode: Sized { - /// Provides read-only access to child nodes. - fn children(&self) -> &[Self]; + fn f_down(&mut self, node: Self::Node) -> Result> { + (self.f_down_func)(node) + } - /// Detaches the node from its children, returning the node itself and its detached children. - fn take_children(self) -> (Self, Vec); + fn f_up(&mut self, node: Self::Node) -> Result> { + (self.f_up_func)(node) + } +} - /// Reattaches updated child nodes to the node, returning the updated node. - fn with_new_children(self, children: Vec) -> Result; +impl Transformed { + fn children(mut self) -> (Self, Vec) { + let (data, children) = self.data.take_children(); + self.data = data; + + (self, children) + } } impl TreeNode for T { @@ -989,7 +1012,7 @@ impl TreeNode for T { &'n self, f: F, ) -> Result { - self.children().iter().apply_until_stop(f) + self.children().into_iter().apply_until_stop(f) } fn map_children Result>>( @@ -1006,574 +1029,242 @@ impl TreeNode for T { Ok(Transformed::no(new_self)) } } -} - -#[cfg(test)] -pub(crate) mod tests { - use std::collections::HashMap; - use std::fmt::Display; - - use crate::tree_node::{ - Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, TreeNodeRewriter, - TreeNodeVisitor, - }; - use crate::Result; - - #[derive(Debug, Eq, Hash, PartialEq, Clone)] - pub struct TestTreeNode { - pub(crate) children: Vec>, - pub(crate) data: T, - } - - impl TestTreeNode { - pub(crate) fn new(children: Vec>, data: T) -> Self { - Self { children, data } - } - - pub(crate) fn new_leaf(data: T) -> Self { - Self { - children: vec![], - data, - } - } - - pub(crate) fn is_leaf(&self) -> bool { - self.children.is_empty() - } - } - - impl TreeNode for TestTreeNode { - fn apply_children<'n, F: FnMut(&'n Self) -> Result>( - &'n self, - f: F, - ) -> Result { - self.children.iter().apply_until_stop(f) - } - - fn map_children Result>>( - self, - f: F, - ) -> Result> { - Ok(self - .children - .into_iter() - .map_until_stop_and_collect(f)? - .update_data(|new_children| Self { - children: new_children, - ..self - })) - } - } - - // J - // | - // I - // | - // F - // / \ - // E G - // | | - // C H - // / \ - // B D - // | - // A - fn test_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "e".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "i".to_string()); - TestTreeNode::new(vec![node_i], "j".to_string()) - } - - // Continue on all nodes - // Expected visits in a combined traversal - fn all_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - "f_up(d)", - "f_up(c)", - "f_up(e)", - "f_down(g)", - "f_down(h)", - "f_up(h)", - "f_up(g)", - "f_up(f)", - "f_up(i)", - "f_up(j)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - // Expected transformed tree after a combined traversal - fn transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(f_down(a))".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(f_down(b))".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_up(f_down(d))".to_string()); - let node_c = - TestTreeNode::new(vec![node_b, node_d], "f_up(f_down(c))".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string()); - let node_h = TestTreeNode::new_leaf("f_up(f_down(h))".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string()); - let node_f = - TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string()); - TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string()) - } - - // Expected transformed tree after a top-down traversal - fn transformed_down_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_down(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_down(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("f_down(h)".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - // Expected transformed tree after a bottom-up traversal - fn transformed_up_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_up(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_up(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_up(e)".to_string()); - let node_h = TestTreeNode::new_leaf("f_up(h)".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_up(g)".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_up(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_up(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_up(j)".to_string()) - } - - // f_down Jump on A node - fn f_down_jump_on_a_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - "f_up(d)", - "f_up(c)", - "f_up(e)", - "f_down(g)", - "f_down(h)", - "f_up(h)", - "f_up(g)", - "f_up(f)", - "f_up(i)", - "f_up(j)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_down_jump_on_a_transformed_down_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_down(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_down(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("f_down(h)".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - // f_down Jump on E node - fn f_down_jump_on_e_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_up(e)", - "f_down(g)", - "f_down(h)", - "f_up(h)", - "f_up(g)", - "f_up(f)", - "f_up(i)", - "f_up(j)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_down_jump_on_e_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string()); - let node_h = TestTreeNode::new_leaf("f_up(f_down(h))".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string()); - let node_f = - TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string()); - TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string()) - } - - fn f_down_jump_on_e_transformed_down_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("f_down(h)".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - // f_up Jump on A node - fn f_up_jump_on_a_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - "f_down(g)", - "f_down(h)", - "f_up(h)", - "f_up(g)", - "f_up(f)", - "f_up(i)", - "f_up(j)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_up_jump_on_a_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(f_down(a))".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(f_down(b))".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("f_up(f_down(h))".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string()); - let node_f = - TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string()); - TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string()) - } - - fn f_up_jump_on_a_transformed_up_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "e".to_string()); - let node_h = TestTreeNode::new_leaf("f_up(h)".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "f_up(g)".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_up(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_up(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_up(j)".to_string()) - } - // f_up Jump on E node - fn f_up_jump_on_e_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - "f_up(d)", - "f_up(c)", - "f_up(e)", - "f_down(g)", - "f_down(h)", - "f_up(h)", - "f_up(g)", - "f_up(f)", - "f_up(i)", - "f_up(j)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_up_jump_on_e_transformed_tree() -> TestTreeNode { - transformed_tree() - } - - fn f_up_jump_on_e_transformed_up_tree() -> TestTreeNode { - transformed_up_tree() - } - - // f_down Stop on A node - - fn f_down_stop_on_a_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_down_stop_on_a_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_down(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(f_down(b))".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - fn f_down_stop_on_a_transformed_down_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_down(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_down(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - // f_down Stop on E node - fn f_down_stop_on_e_visits() -> Vec { - vec!["f_down(j)", "f_down(i)", "f_down(f)", "f_down(e)"] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_down_stop_on_e_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - fn f_down_stop_on_e_transformed_down_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - // f_up Stop on A node - fn f_up_stop_on_a_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_up_stop_on_a_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(f_down(a))".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(f_down(b))".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) - } - - fn f_up_stop_on_a_transformed_up_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "e".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "i".to_string()); - TestTreeNode::new(vec![node_i], "j".to_string()) - } - - // f_up Stop on E node - fn f_up_stop_on_e_visits() -> Vec { - vec![ - "f_down(j)", - "f_down(i)", - "f_down(f)", - "f_down(e)", - "f_down(c)", - "f_down(b)", - "f_up(b)", - "f_down(d)", - "f_down(a)", - "f_up(a)", - "f_up(d)", - "f_up(c)", - "f_up(e)", - ] - .into_iter() - .map(|s| s.to_string()) - .collect() - } - - fn f_up_stop_on_e_transformed_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(f_down(a))".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(f_down(b))".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_up(f_down(d))".to_string()); - let node_c = - TestTreeNode::new(vec![node_b, node_d], "f_up(f_down(c))".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string()); - TestTreeNode::new(vec![node_i], "f_down(j)".to_string()) + fn transform_down_up< + FD: FnMut(Self) -> Result>, + FU: FnMut(Self) -> Result>, + >( + self, + f_down: FD, + f_up: FU, + ) -> Result> { + self.rewrite(&mut FuncRewriter::new(f_down, f_up)) } - fn f_up_stop_on_e_transformed_up_tree() -> TestTreeNode { - let node_a = TestTreeNode::new_leaf("f_up(a)".to_string()); - let node_b = TestTreeNode::new_leaf("f_up(b)".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "f_up(d)".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "f_up(c)".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "f_up(e)".to_string()); - let node_h = TestTreeNode::new_leaf("h".to_string()); - let node_g = TestTreeNode::new(vec![node_h], "g".to_string()); - let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string()); - let node_i = TestTreeNode::new(vec![node_f], "i".to_string()); - TestTreeNode::new(vec![node_i], "j".to_string()) + fn transform_down Result>>( + self, + f: F, + ) -> Result> { + self.rewrite(&mut FuncRewriter::new(f, |node| Ok(Transformed::no(node)))) } - fn down_visits(visits: Vec) -> Vec { - visits - .into_iter() - .filter(|v| v.starts_with("f_down")) - .collect() + fn transform_up Result>>( + self, + f: F, + ) -> Result> { + self.rewrite(&mut FuncRewriter::new(|node| Ok(Transformed::no(node)), f)) } + fn rewrite>( + self, + rewriter: &mut R, + ) -> Result> { + let mut stack = vec![TransformingState::NotStarted(self)]; + + while let Some(item) = stack.pop() { + match item { + TransformingState::NotStarted(node) => { + let node = rewriter.f_down(node)?; + + stack.push(match node.tnr { + TreeNodeRecursion::Continue => { + let (node, mut non_processed_children) = node.children(); + non_processed_children.reverse(); + + TransformingState::ProcessingChildren { + processed_children: Vec::with_capacity( + non_processed_children.len(), + ), + non_processed_children, + item: node, + } + } + TreeNodeRecursion::Jump => { + TransformingState::ProcessedAllChildren( + // No need to process children, we can just skip that stage + node.with_tnr(TreeNodeRecursion::Continue), + ) + } + TreeNodeRecursion::Stop => { + TransformingState::ProcessedAllChildren(node) + } + }) + } + TransformingState::ProcessingChildren { + mut item, + mut non_processed_children, + mut processed_children, + } => match item.tnr { + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { + if let Some(non_processed_item) = non_processed_children.pop() { + stack.extend([ + // This node still has children, so put it back on the stack + TransformingState::ProcessingChildren { + item, + non_processed_children, + processed_children, + }, + // Also put the child which will be processed first + TransformingState::NotStarted(non_processed_item), + ]); + } else { + item.data = + item.data.with_new_children(processed_children)?; + stack.push(TransformingState::ProcessedAllChildren(item)) + } + } + TreeNodeRecursion::Stop => { + // At this point, we might have some children we haven't yet processed + processed_children + .extend(non_processed_children.into_iter().rev()); + item.data = item.data.with_new_children(processed_children)?; + stack.push(TransformingState::ProcessedAllChildren(item)); + } + }, + TransformingState::ProcessedAllChildren(node) => { + let node = node.transform_parent(|n| rewriter.f_up(n))?; + + if let Some(TransformingState::ProcessingChildren { + item: ref mut parent_node, + ref mut processed_children, + .. + }) = stack.last_mut() + { + // We need to use the returned iteration state when processing the remaining children + parent_node.tnr = node.tnr; + parent_node.transformed |= node.transformed; + processed_children.push(node.data); + } else { + debug_assert!(stack.is_empty()); + return Ok(node); + } + } + } + } - type TestVisitorF = Box) -> Result>; - - struct TestVisitor { - visits: Vec, - f_down: TestVisitorF, - f_up: TestVisitorF, + unreachable!(); } - impl TestVisitor { - fn new(f_down: TestVisitorF, f_up: TestVisitorF) -> Self { - Self { - visits: vec![], - f_down, - f_up, + fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>( + &'n self, + visitor: &mut V, + ) -> Result { + let mut stack = vec![VisitingState::NotStarted(self)]; + + while let Some(node) = stack.last_mut() { + match node { + VisitingState::NotStarted(item) => { + let tnr = visitor.f_down(item)?; + *node = match tnr { + TreeNodeRecursion::Continue => { + let mut non_processed_children = item.children(); + non_processed_children.reverse(); + + VisitingState::VisitingChildren { + non_processed_children, + item, + tnr, + } + } + TreeNodeRecursion::Jump => VisitingState::VisitedAllChildren { + item, + tnr: TreeNodeRecursion::Continue, + }, + TreeNodeRecursion::Stop => return Ok(tnr), + }; + } + VisitingState::VisitingChildren { + item, + ref mut non_processed_children, + tnr, + } => match tnr { + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { + if let Some(non_processed_item) = non_processed_children.pop() { + stack.push(VisitingState::NotStarted(non_processed_item)); + } else { + *node = VisitingState::VisitedAllChildren { item, tnr: *tnr } + } + } + TreeNodeRecursion::Stop => { + return Ok(*tnr); + } + }, + VisitingState::VisitedAllChildren { item, tnr } => { + let tnr = tnr.visit_parent(|| visitor.f_up(item))?; + stack.pop(); + + if let Some(VisitingState::VisitingChildren { + item, + non_processed_children, + .. // we don't care about the parent recursion state, because it will be replaced with the current state anyway + }) = stack.pop() + { + stack.push(VisitingState::VisitingChildren { + item, + non_processed_children, + tnr, + }); + } else { + debug_assert!(stack.is_empty()); + return Ok(tnr); + } + } } } + + unreachable!() } +} - impl<'n, T: Display> TreeNodeVisitor<'n> for TestVisitor { - type Node = TestTreeNode; +/// Node iterative transformation state. Each node on the stack is visited several times, state determines the operation that is going to run next +enum TransformingState { + /// Node was just created. When executed, f_down will be called. + NotStarted(T), + /// DFS over node's children + ProcessingChildren { + item: Transformed, + non_processed_children: Vec, + processed_children: Vec, + }, + /// All children are processed (or jumped through). When executed, f_up may be called + ProcessedAllChildren(Transformed), +} - fn f_down(&mut self, node: &'n Self::Node) -> Result { - self.visits.push(format!("f_down({})", node.data)); - (*self.f_down)(node) - } +/// Node iterative visit state. Each node on the stack is visited several times, state determines the operation that is going to run next +enum VisitingState<'a, T> { + /// Node was just created. When executed, f_down will be called. + NotStarted(&'a T), + /// DFS over node's children. During processing, reference to children are removed from the inner stack + VisitingChildren { + item: &'a T, + non_processed_children: Vec<&'a T>, + tnr: TreeNodeRecursion, + }, + /// All children are processed (or jumped through). When executed, f_up may be called + VisitedAllChildren { item: &'a T, tnr: TreeNodeRecursion }, +} +/// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for +/// trees that contain nodes with payloads. This approach ensures safe execution of algorithms +/// involving payloads, by enforcing rules for detaching and reattaching child nodes. +pub trait ConcreteTreeNode: Sized { + /// Provides read-only access to child nodes. + fn children(&self) -> Vec<&Self>; - fn f_up(&mut self, node: &'n Self::Node) -> Result { - self.visits.push(format!("f_up({})", node.data)); - (*self.f_up)(node) - } - } + /// Detaches the node from its children, returning the node itself and its detached children. + fn take_children(self) -> (Self, Vec); - fn visit_continue(_: &TestTreeNode) -> Result { - Ok(TreeNodeRecursion::Continue) - } + /// Reattaches updated child nodes to the node, returning the updated node. + fn with_new_children(self, children: Vec) -> Result; +} - fn visit_event_on>( - data: D, - event: TreeNodeRecursion, - ) -> impl FnMut(&TestTreeNode) -> Result { - let d = data.into(); - move |node| { - Ok(if node.data == d { - event - } else { - TreeNodeRecursion::Continue - }) - } - } +#[cfg(test)] +pub(crate) mod tests { + use crate::tree_node::{ + Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, TreeNodeRewriter, + TreeNodeVisitor, + }; + use crate::Result; + use std::collections::HashMap; + use std::fmt::Display; macro_rules! visit_test { ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_VISITS:expr) => { @@ -1606,66 +1297,6 @@ pub(crate) mod tests { }; } - type TestRewriterF = - Box) -> Result>>>; - - struct TestRewriter { - f_down: TestRewriterF, - f_up: TestRewriterF, - } - - impl TestRewriter { - fn new(f_down: TestRewriterF, f_up: TestRewriterF) -> Self { - Self { f_down, f_up } - } - } - - impl TreeNodeRewriter for TestRewriter { - type Node = TestTreeNode; - - fn f_down(&mut self, node: Self::Node) -> Result> { - (*self.f_down)(node) - } - - fn f_up(&mut self, node: Self::Node) -> Result> { - (*self.f_up)(node) - } - } - - fn transform_yes>( - transformation_name: N, - ) -> impl FnMut(TestTreeNode) -> Result>> { - move |node| { - Ok(Transformed::yes(TestTreeNode::new( - node.children, - format!("{}({})", transformation_name, node.data).into(), - ))) - } - } - - fn transform_and_event_on< - N: Display, - T: PartialEq + Display + From, - D: Into, - >( - transformation_name: N, - data: D, - event: TreeNodeRecursion, - ) -> impl FnMut(TestTreeNode) -> Result>> { - let d = data.into(); - move |node| { - let new_node = TestTreeNode::new( - node.children, - format!("{}({})", transformation_name, node.data).into(), - ); - Ok(if node.data == d { - Transformed::new(new_node, true, event) - } else { - Transformed::yes(new_node) - }) - } - } - macro_rules! rewrite_test { ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => { #[test] @@ -1715,368 +1346,1179 @@ pub(crate) mod tests { }; } - visit_test!(test_visit, visit_continue, visit_continue, all_visits()); - visit_test!( - test_visit_f_down_jump_on_a, - visit_event_on("a", TreeNodeRecursion::Jump), - visit_continue, - f_down_jump_on_a_visits() - ); - visit_test!( - test_visit_f_down_jump_on_e, - visit_event_on("e", TreeNodeRecursion::Jump), - visit_continue, - f_down_jump_on_e_visits() - ); - visit_test!( - test_visit_f_up_jump_on_a, - visit_continue, - visit_event_on("a", TreeNodeRecursion::Jump), - f_up_jump_on_a_visits() - ); - visit_test!( - test_visit_f_up_jump_on_e, - visit_continue, - visit_event_on("e", TreeNodeRecursion::Jump), - f_up_jump_on_e_visits() - ); - visit_test!( - test_visit_f_down_stop_on_a, - visit_event_on("a", TreeNodeRecursion::Stop), - visit_continue, - f_down_stop_on_a_visits() - ); - visit_test!( - test_visit_f_down_stop_on_e, - visit_event_on("e", TreeNodeRecursion::Stop), - visit_continue, - f_down_stop_on_e_visits() - ); - visit_test!( - test_visit_f_up_stop_on_a, - visit_continue, - visit_event_on("a", TreeNodeRecursion::Stop), - f_up_stop_on_a_visits() - ); - visit_test!( - test_visit_f_up_stop_on_e, - visit_continue, - visit_event_on("e", TreeNodeRecursion::Stop), - f_up_stop_on_e_visits() - ); - - test_apply!(test_apply, visit_continue, down_visits(all_visits())); - test_apply!( - test_apply_f_down_jump_on_a, - visit_event_on("a", TreeNodeRecursion::Jump), - down_visits(f_down_jump_on_a_visits()) - ); - test_apply!( - test_apply_f_down_jump_on_e, - visit_event_on("e", TreeNodeRecursion::Jump), - down_visits(f_down_jump_on_e_visits()) - ); - test_apply!( - test_apply_f_down_stop_on_a, - visit_event_on("a", TreeNodeRecursion::Stop), - down_visits(f_down_stop_on_a_visits()) - ); - test_apply!( - test_apply_f_down_stop_on_e, - visit_event_on("e", TreeNodeRecursion::Stop), - down_visits(f_down_stop_on_e_visits()) - ); - - rewrite_test!( - test_rewrite, - transform_yes("f_down"), - transform_yes("f_up"), - Transformed::yes(transformed_tree()) - ); - rewrite_test!( - test_rewrite_f_down_jump_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), - transform_yes("f_up"), - Transformed::yes(transformed_tree()) - ); - rewrite_test!( - test_rewrite_f_down_jump_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), - transform_yes("f_up"), - Transformed::yes(f_down_jump_on_e_transformed_tree()) - ); - rewrite_test!( - test_rewrite_f_up_jump_on_a, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_a_transformed_tree()) - ); - rewrite_test!( - test_rewrite_f_up_jump_on_e, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_e_transformed_tree()) - ); - rewrite_test!( - test_rewrite_f_down_stop_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), - transform_yes("f_up"), - Transformed::new( - f_down_stop_on_a_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - rewrite_test!( - test_rewrite_f_down_stop_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), - transform_yes("f_up"), - Transformed::new( - f_down_stop_on_e_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - rewrite_test!( - test_rewrite_f_up_stop_on_a, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_a_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - rewrite_test!( - test_rewrite_f_up_stop_on_e, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_e_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - - transform_test!( - test_transform, - transform_yes("f_down"), - transform_yes("f_up"), - Transformed::yes(transformed_tree()) - ); - transform_test!( - test_transform_f_down_jump_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), - transform_yes("f_up"), - Transformed::yes(transformed_tree()) - ); - transform_test!( - test_transform_f_down_jump_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), - transform_yes("f_up"), - Transformed::yes(f_down_jump_on_e_transformed_tree()) - ); - transform_test!( - test_transform_f_up_jump_on_a, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_a_transformed_tree()) - ); - transform_test!( - test_transform_f_up_jump_on_e, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_e_transformed_tree()) - ); - transform_test!( - test_transform_f_down_stop_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), - transform_yes("f_up"), - Transformed::new( - f_down_stop_on_a_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - transform_test!( - test_transform_f_down_stop_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), - transform_yes("f_up"), - Transformed::new( - f_down_stop_on_e_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - transform_test!( - test_transform_f_up_stop_on_a, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_a_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - transform_test!( - test_transform_f_up_stop_on_e, - transform_yes("f_down"), - transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_e_transformed_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - - transform_down_test!( - test_transform_down, - transform_yes("f_down"), - Transformed::yes(transformed_down_tree()) - ); - transform_down_test!( - test_transform_down_f_down_jump_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), - Transformed::yes(f_down_jump_on_a_transformed_down_tree()) - ); - transform_down_test!( - test_transform_down_f_down_jump_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), - Transformed::yes(f_down_jump_on_e_transformed_down_tree()) - ); - transform_down_test!( - test_transform_down_f_down_stop_on_a, - transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), - Transformed::new( - f_down_stop_on_a_transformed_down_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - transform_down_test!( - test_transform_down_f_down_stop_on_e, - transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), - Transformed::new( - f_down_stop_on_e_transformed_down_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - - transform_up_test!( - test_transform_up, - transform_yes("f_up"), - Transformed::yes(transformed_up_tree()) - ); - transform_up_test!( - test_transform_up_f_up_jump_on_a, - transform_and_event_on("f_up", "a", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_a_transformed_up_tree()) - ); - transform_up_test!( - test_transform_up_f_up_jump_on_e, - transform_and_event_on("f_up", "e", TreeNodeRecursion::Jump), - Transformed::yes(f_up_jump_on_e_transformed_up_tree()) - ); - transform_up_test!( - test_transform_up_f_up_stop_on_a, - transform_and_event_on("f_up", "a", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_a_transformed_up_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - transform_up_test!( - test_transform_up_f_up_stop_on_e, - transform_and_event_on("f_up", "e", TreeNodeRecursion::Stop), - Transformed::new( - f_up_stop_on_e_transformed_up_tree(), - true, - TreeNodeRecursion::Stop - ) - ); - - // F - // / | \ - // / | \ - // E C A - // | / \ - // C B D - // / \ | - // B D A - // | - // A - #[test] - fn test_apply_and_visit_references() -> Result<()> { - let node_a = TestTreeNode::new_leaf("a".to_string()); - let node_b = TestTreeNode::new_leaf("b".to_string()); - let node_d = TestTreeNode::new(vec![node_a], "d".to_string()); - let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string()); - let node_e = TestTreeNode::new(vec![node_c], "e".to_string()); - let node_a_2 = TestTreeNode::new_leaf("a".to_string()); - let node_b_2 = TestTreeNode::new_leaf("b".to_string()); - let node_d_2 = TestTreeNode::new(vec![node_a_2], "d".to_string()); - let node_c_2 = TestTreeNode::new(vec![node_b_2, node_d_2], "c".to_string()); - let node_a_3 = TestTreeNode::new_leaf("a".to_string()); - let tree = TestTreeNode::new(vec![node_e, node_c_2, node_a_3], "f".to_string()); - - let node_f_ref = &tree; - let node_e_ref = &node_f_ref.children[0]; - let node_c_ref = &node_e_ref.children[0]; - let node_b_ref = &node_c_ref.children[0]; - let node_d_ref = &node_c_ref.children[1]; - let node_a_ref = &node_d_ref.children[0]; - - let mut m: HashMap<&TestTreeNode, usize> = HashMap::new(); - tree.apply(|e| { - *m.entry(e).or_insert(0) += 1; - Ok(TreeNodeRecursion::Continue) - })?; - - let expected = HashMap::from([ - (node_f_ref, 1), - (node_e_ref, 1), - (node_c_ref, 2), - (node_d_ref, 2), - (node_b_ref, 2), - (node_a_ref, 3), - ]); - assert_eq!(m, expected); - - struct TestVisitor<'n> { - m: HashMap<&'n TestTreeNode, (usize, usize)>, - } + fn down_visits(visits: Vec) -> Vec { + visits + .into_iter() + .filter(|v| v.starts_with("f_down")) + .collect() + } + + /// [`node_tests`] uses methods when generating tests + pub(crate) trait TestTree { + fn new_with_children(children: Vec, data: T) -> Self + where + Self: Sized; - impl<'n> TreeNodeVisitor<'n> for TestVisitor<'n> { - type Node = TestTreeNode; + fn new_leaf(data: T) -> Self; - fn f_down(&mut self, node: &'n Self::Node) -> Result { - let (down_count, _) = self.m.entry(node).or_insert((0, 0)); - *down_count += 1; + fn is_leaf(&self) -> bool; + + fn eq_data(&self, other: &T) -> bool; + + fn with_children_from(data: T, other: Self) -> Self; + } + + macro_rules! node_tests { + ($TYPE:ident) => { + fn visit_continue(_: &$TYPE) -> Result { Ok(TreeNodeRecursion::Continue) } - fn f_up(&mut self, node: &'n Self::Node) -> Result { - let (_, up_count) = self.m.entry(node).or_insert((0, 0)); - *up_count += 1; - Ok(TreeNodeRecursion::Continue) + type TestVisitorF = Box) -> Result>; + + struct TestVisitor { + visits: Vec, + f_down: TestVisitorF, + f_up: TestVisitorF, + } + + impl TestVisitor { + fn new(f_down: TestVisitorF, f_up: TestVisitorF) -> Self { + Self { + visits: vec![], + f_down, + f_up, + } + } + } + + impl<'n, T: Display + Eq> TreeNodeVisitor<'n> for TestVisitor { + type Node = $TYPE; + + fn f_down(&mut self, node: &'n Self::Node) -> Result { + self.visits.push(format!("f_down({})", node.data)); + (*self.f_down)(node) + } + + fn f_up(&mut self, node: &'n Self::Node) -> Result { + self.visits.push(format!("f_up({})", node.data)); + (*self.f_up)(node) + } + } + + fn visit_event_on>( + data: D, + event: TreeNodeRecursion, + ) -> impl FnMut(&$TYPE) -> Result { + let d = data.into(); + move |node| { + Ok(if node.eq_data(&d) { + event + } else { + TreeNodeRecursion::Continue + }) + } + } + + type TestRewriterF = + Box) -> Result>>>; + + struct TestRewriter { + f_down: TestRewriterF, + f_up: TestRewriterF, + } + + impl TestRewriter { + fn new(f_down: TestRewriterF, f_up: TestRewriterF) -> Self { + Self { f_down, f_up } + } + } + + impl TreeNodeRewriter for TestRewriter { + type Node = $TYPE; + + fn f_down( + &mut self, + node: Self::Node, + ) -> Result> { + (*self.f_down)(node) + } + + fn f_up(&mut self, node: Self::Node) -> Result> { + (*self.f_up)(node) + } + } + + fn transform_yes + PartialEq>( + transformation_name: N, + ) -> impl FnMut($TYPE) -> Result>> { + move |node| { + Ok(Transformed::yes($TYPE::with_children_from( + format!("{}({})", transformation_name, node.data).into(), + node, + ))) + } + } + + fn transform_and_event_on< + N: Display, + T: PartialEq + Display + From, + D: Into, + >( + transformation_name: N, + data: D, + event: TreeNodeRecursion, + ) -> impl FnMut($TYPE) -> Result>> { + let d = data.into(); + move |node| { + let node_eq = node.eq_data(&d); + let new_node = $TYPE::with_children_from( + format!("{}({})", transformation_name, node.data).into(), + node, + ); + Ok(if node_eq { + Transformed::new(new_node, true, event) + } else { + Transformed::yes(new_node) + }) + } + } + + // J + // | + // I + // | + // F + // / \ + // E G + // | | + // C H + // / \ + // B D + // | + // A + fn test_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = $TYPE::new_with_children(vec![node_c], "e".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = + $TYPE::new_with_children(vec![node_e, node_g], "f".to_string()); + let node_i = $TYPE::new_with_children(vec![node_f], "i".to_string()); + $TYPE::new_with_children(vec![node_i], "j".to_string()) + } + + // Continue on all nodes + // Expected visits in a combined traversal + fn all_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + "f_up(d)", + "f_up(c)", + "f_up(e)", + "f_down(g)", + "f_down(h)", + "f_up(h)", + "f_up(g)", + "f_up(f)", + "f_up(i)", + "f_up(j)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + // Expected transformed tree after a combined traversal + fn transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(f_down(a))".to_string()); + let node_b = $TYPE::new_leaf("f_up(f_down(b))".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_up(f_down(d))".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_up(f_down(c))".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_up(f_down(e))".to_string()); + let node_h = $TYPE::new_leaf("f_up(f_down(h))".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_up(f_down(g))".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_up(f_down(f))".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_up(f_down(i))".to_string()); + $TYPE::new_with_children(vec![node_i], "f_up(f_down(j))".to_string()) + } + + // Expected transformed tree after a top-down traversal + fn transformed_down_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_down(a)".to_string()); + let node_b = $TYPE::new_leaf("f_down(b)".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("f_down(h)".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_down(g)".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + // Expected transformed tree after a bottom-up traversal + fn transformed_up_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(a)".to_string()); + let node_b = $TYPE::new_leaf("f_up(b)".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_up(d)".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "f_up(c)".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_up(e)".to_string()); + let node_h = $TYPE::new_leaf("f_up(h)".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_up(g)".to_string()); + let node_f = + $TYPE::new_with_children(vec![node_e, node_g], "f_up(f)".to_string()); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_up(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_up(j)".to_string()) + } + + // f_down Jump on A node + fn f_down_jump_on_a_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + "f_up(d)", + "f_up(c)", + "f_up(e)", + "f_down(g)", + "f_down(h)", + "f_up(h)", + "f_up(g)", + "f_up(f)", + "f_up(i)", + "f_up(j)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_down_jump_on_a_transformed_down_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_down(a)".to_string()); + let node_b = $TYPE::new_leaf("f_down(b)".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("f_down(h)".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_down(g)".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + // f_down Jump on E node + fn f_down_jump_on_e_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_up(e)", + "f_down(g)", + "f_down(h)", + "f_up(h)", + "f_up(g)", + "f_up(f)", + "f_up(i)", + "f_up(j)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_down_jump_on_e_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_up(f_down(e))".to_string()); + let node_h = $TYPE::new_leaf("f_up(f_down(h))".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_up(f_down(g))".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_up(f_down(f))".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_up(f_down(i))".to_string()); + $TYPE::new_with_children(vec![node_i], "f_up(f_down(j))".to_string()) + } + + fn f_down_jump_on_e_transformed_down_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("f_down(h)".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_down(g)".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + // f_up Jump on A node + fn f_up_jump_on_a_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + "f_down(g)", + "f_down(h)", + "f_up(h)", + "f_up(g)", + "f_up(f)", + "f_up(i)", + "f_up(j)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_up_jump_on_a_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(f_down(a))".to_string()); + let node_b = $TYPE::new_leaf("f_up(f_down(b))".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("f_up(f_down(h))".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_up(f_down(g))".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_up(f_down(f))".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_up(f_down(i))".to_string()); + $TYPE::new_with_children(vec![node_i], "f_up(f_down(j))".to_string()) + } + + fn f_up_jump_on_a_transformed_up_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(a)".to_string()); + let node_b = $TYPE::new_leaf("f_up(b)".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = $TYPE::new_with_children(vec![node_c], "e".to_string()); + let node_h = $TYPE::new_leaf("f_up(h)".to_string()); + let node_g = + $TYPE::new_with_children(vec![node_h], "f_up(g)".to_string()); + let node_f = + $TYPE::new_with_children(vec![node_e, node_g], "f_up(f)".to_string()); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_up(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_up(j)".to_string()) + } + + // f_up Jump on E node + fn f_up_jump_on_e_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + "f_up(d)", + "f_up(c)", + "f_up(e)", + "f_down(g)", + "f_down(h)", + "f_up(h)", + "f_up(g)", + "f_up(f)", + "f_up(i)", + "f_up(j)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_up_jump_on_e_transformed_tree() -> $TYPE { + transformed_tree() + } + + fn f_up_jump_on_e_transformed_up_tree() -> $TYPE { + transformed_up_tree() + } + + // f_down Stop on A node + + fn f_down_stop_on_a_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_down_stop_on_a_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_down(a)".to_string()); + let node_b = $TYPE::new_leaf("f_up(f_down(b))".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + fn f_down_stop_on_a_transformed_down_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_down(a)".to_string()); + let node_b = $TYPE::new_leaf("f_down(b)".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + // f_down Stop on E node + fn f_down_stop_on_e_visits() -> Vec { + vec!["f_down(j)", "f_down(i)", "f_down(f)", "f_down(e)"] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_down_stop_on_e_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + fn f_down_stop_on_e_transformed_down_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + // f_up Stop on A node + fn f_up_stop_on_a_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_up_stop_on_a_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(f_down(a))".to_string()); + let node_b = $TYPE::new_leaf("f_up(f_down(b))".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_down(d)".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_down(c)".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_down(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + fn f_up_stop_on_a_transformed_up_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(a)".to_string()); + let node_b = $TYPE::new_leaf("f_up(b)".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = $TYPE::new_with_children(vec![node_c], "e".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = + $TYPE::new_with_children(vec![node_e, node_g], "f".to_string()); + let node_i = $TYPE::new_with_children(vec![node_f], "i".to_string()); + $TYPE::new_with_children(vec![node_i], "j".to_string()) + } + + // f_up Stop on E node + fn f_up_stop_on_e_visits() -> Vec { + vec![ + "f_down(j)", + "f_down(i)", + "f_down(f)", + "f_down(e)", + "f_down(c)", + "f_down(b)", + "f_up(b)", + "f_down(d)", + "f_down(a)", + "f_up(a)", + "f_up(d)", + "f_up(c)", + "f_up(e)", + ] + .into_iter() + .map(|s| s.to_string()) + .collect() + } + + fn f_up_stop_on_e_transformed_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(f_down(a))".to_string()); + let node_b = $TYPE::new_leaf("f_up(f_down(b))".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_up(f_down(d))".to_string()); + let node_c = $TYPE::new_with_children( + vec![node_b, node_d], + "f_up(f_down(c))".to_string(), + ); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_up(f_down(e))".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = $TYPE::new_with_children( + vec![node_e, node_g], + "f_down(f)".to_string(), + ); + let node_i = + $TYPE::new_with_children(vec![node_f], "f_down(i)".to_string()); + $TYPE::new_with_children(vec![node_i], "f_down(j)".to_string()) + } + + fn f_up_stop_on_e_transformed_up_tree() -> $TYPE { + let node_a = $TYPE::new_leaf("f_up(a)".to_string()); + let node_b = $TYPE::new_leaf("f_up(b)".to_string()); + let node_d = + $TYPE::new_with_children(vec![node_a], "f_up(d)".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "f_up(c)".to_string()); + let node_e = + $TYPE::new_with_children(vec![node_c], "f_up(e)".to_string()); + let node_h = $TYPE::new_leaf("h".to_string()); + let node_g = $TYPE::new_with_children(vec![node_h], "g".to_string()); + let node_f = + $TYPE::new_with_children(vec![node_e, node_g], "f".to_string()); + let node_i = $TYPE::new_with_children(vec![node_f], "i".to_string()); + $TYPE::new_with_children(vec![node_i], "j".to_string()) + } + + visit_test!(test_visit, visit_continue, visit_continue, all_visits()); + visit_test!( + test_visit_f_down_jump_on_a, + visit_event_on("a", TreeNodeRecursion::Jump), + visit_continue, + f_down_jump_on_a_visits() + ); + visit_test!( + test_visit_f_down_jump_on_e, + visit_event_on("e", TreeNodeRecursion::Jump), + visit_continue, + f_down_jump_on_e_visits() + ); + visit_test!( + test_visit_f_up_jump_on_a, + visit_continue, + visit_event_on("a", TreeNodeRecursion::Jump), + f_up_jump_on_a_visits() + ); + visit_test!( + test_visit_f_up_jump_on_e, + visit_continue, + visit_event_on("e", TreeNodeRecursion::Jump), + f_up_jump_on_e_visits() + ); + visit_test!( + test_visit_f_down_stop_on_a, + visit_event_on("a", TreeNodeRecursion::Stop), + visit_continue, + f_down_stop_on_a_visits() + ); + visit_test!( + test_visit_f_down_stop_on_e, + visit_event_on("e", TreeNodeRecursion::Stop), + visit_continue, + f_down_stop_on_e_visits() + ); + visit_test!( + test_visit_f_up_stop_on_a, + visit_continue, + visit_event_on("a", TreeNodeRecursion::Stop), + f_up_stop_on_a_visits() + ); + visit_test!( + test_visit_f_up_stop_on_e, + visit_continue, + visit_event_on("e", TreeNodeRecursion::Stop), + f_up_stop_on_e_visits() + ); + + test_apply!(test_apply, visit_continue, down_visits(all_visits())); + test_apply!( + test_apply_f_down_jump_on_a, + visit_event_on("a", TreeNodeRecursion::Jump), + down_visits(f_down_jump_on_a_visits()) + ); + test_apply!( + test_apply_f_down_jump_on_e, + visit_event_on("e", TreeNodeRecursion::Jump), + down_visits(f_down_jump_on_e_visits()) + ); + test_apply!( + test_apply_f_down_stop_on_a, + visit_event_on("a", TreeNodeRecursion::Stop), + down_visits(f_down_stop_on_a_visits()) + ); + test_apply!( + test_apply_f_down_stop_on_e, + visit_event_on("e", TreeNodeRecursion::Stop), + down_visits(f_down_stop_on_e_visits()) + ); + + rewrite_test!( + test_rewrite, + transform_yes("f_down"), + transform_yes("f_up"), + Transformed::yes(transformed_tree()) + ); + rewrite_test!( + test_rewrite_f_down_jump_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), + transform_yes("f_up"), + Transformed::yes(transformed_tree()) + ); + rewrite_test!( + test_rewrite_f_down_jump_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), + transform_yes("f_up"), + Transformed::yes(f_down_jump_on_e_transformed_tree()) + ); + rewrite_test!( + test_rewrite_f_up_jump_on_a, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_a_transformed_tree()) + ); + rewrite_test!( + test_rewrite_f_up_jump_on_e, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_e_transformed_tree()) + ); + rewrite_test!( + test_rewrite_f_down_stop_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), + transform_yes("f_up"), + Transformed::new( + f_down_stop_on_a_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + rewrite_test!( + test_rewrite_f_down_stop_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), + transform_yes("f_up"), + Transformed::new( + f_down_stop_on_e_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + rewrite_test!( + test_rewrite_f_up_stop_on_a, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_a_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + rewrite_test!( + test_rewrite_f_up_stop_on_e, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_e_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + + transform_test!( + test_transform, + transform_yes("f_down"), + transform_yes("f_up"), + Transformed::yes(transformed_tree()) + ); + transform_test!( + test_transform_f_down_jump_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), + transform_yes("f_up"), + Transformed::yes(transformed_tree()) + ); + transform_test!( + test_transform_f_down_jump_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), + transform_yes("f_up"), + Transformed::yes(f_down_jump_on_e_transformed_tree()) + ); + transform_test!( + test_transform_f_up_jump_on_a, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_a_transformed_tree()) + ); + transform_test!( + test_transform_f_up_jump_on_e, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_e_transformed_tree()) + ); + transform_test!( + test_transform_f_down_stop_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), + transform_yes("f_up"), + Transformed::new( + f_down_stop_on_a_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + transform_test!( + test_transform_f_down_stop_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), + transform_yes("f_up"), + Transformed::new( + f_down_stop_on_e_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + transform_test!( + test_transform_f_up_stop_on_a, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_a_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + transform_test!( + test_transform_f_up_stop_on_e, + transform_yes("f_down"), + transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_e_transformed_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + + transform_down_test!( + test_transform_down, + transform_yes("f_down"), + Transformed::yes(transformed_down_tree()) + ); + transform_down_test!( + test_transform_down_f_down_jump_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump), + Transformed::yes(f_down_jump_on_a_transformed_down_tree()) + ); + transform_down_test!( + test_transform_down_f_down_jump_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump), + Transformed::yes(f_down_jump_on_e_transformed_down_tree()) + ); + transform_down_test!( + test_transform_down_f_down_stop_on_a, + transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop), + Transformed::new( + f_down_stop_on_a_transformed_down_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + transform_down_test!( + test_transform_down_f_down_stop_on_e, + transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop), + Transformed::new( + f_down_stop_on_e_transformed_down_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + + transform_up_test!( + test_transform_up, + transform_yes("f_up"), + Transformed::yes(transformed_up_tree()) + ); + transform_up_test!( + test_transform_up_f_up_jump_on_a, + transform_and_event_on("f_up", "a", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_a_transformed_up_tree()) + ); + transform_up_test!( + test_transform_up_f_up_jump_on_e, + transform_and_event_on("f_up", "e", TreeNodeRecursion::Jump), + Transformed::yes(f_up_jump_on_e_transformed_up_tree()) + ); + transform_up_test!( + test_transform_up_f_up_stop_on_a, + transform_and_event_on("f_up", "a", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_a_transformed_up_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + transform_up_test!( + test_transform_up_f_up_stop_on_e, + transform_and_event_on("f_up", "e", TreeNodeRecursion::Stop), + Transformed::new( + f_up_stop_on_e_transformed_up_tree(), + true, + TreeNodeRecursion::Stop + ) + ); + + // F + // / | \ + // / | \ + // E C A + // | / \ + // C B D + // / \ | + // B D A + // | + // A + #[test] + fn test_apply_and_visit_references() -> Result<()> { + let node_a = $TYPE::new_leaf("a".to_string()); + let node_b = $TYPE::new_leaf("b".to_string()); + let node_d = $TYPE::new_with_children(vec![node_a], "d".to_string()); + let node_c = + $TYPE::new_with_children(vec![node_b, node_d], "c".to_string()); + let node_e = $TYPE::new_with_children(vec![node_c], "e".to_string()); + let node_a_2 = $TYPE::new_leaf("a".to_string()); + let node_b_2 = $TYPE::new_leaf("b".to_string()); + let node_d_2 = $TYPE::new_with_children(vec![node_a_2], "d".to_string()); + let node_c_2 = + $TYPE::new_with_children(vec![node_b_2, node_d_2], "c".to_string()); + let node_a_3 = $TYPE::new_leaf("a".to_string()); + let tree = $TYPE::new_with_children( + vec![node_e, node_c_2, node_a_3], + "f".to_string(), + ); + + let node_f_ref = &tree; + let node_e_ref = &node_f_ref.children[0]; + let node_c_ref = &node_e_ref.children[0]; + let node_b_ref = &node_c_ref.children[0]; + let node_d_ref = &node_c_ref.children[1]; + let node_a_ref = &node_d_ref.children[0]; + + let mut m: HashMap<&$TYPE, usize> = HashMap::new(); + tree.apply(|e| { + *m.entry(e).or_insert(0) += 1; + Ok(TreeNodeRecursion::Continue) + })?; + + let expected = HashMap::from([ + (node_f_ref, 1), + (node_e_ref, 1), + (node_c_ref, 2), + (node_d_ref, 2), + (node_b_ref, 2), + (node_a_ref, 3), + ]); + assert_eq!(m, expected); + + struct TestVisitor<'n> { + m: HashMap<&'n $TYPE, (usize, usize)>, + } + + impl<'n> TreeNodeVisitor<'n> for TestVisitor<'n> { + type Node = $TYPE; + + fn f_down( + &mut self, + node: &'n Self::Node, + ) -> Result { + let (down_count, _) = self.m.entry(node).or_insert((0, 0)); + *down_count += 1; + Ok(TreeNodeRecursion::Continue) + } + + fn f_up( + &mut self, + node: &'n Self::Node, + ) -> Result { + let (_, up_count) = self.m.entry(node).or_insert((0, 0)); + *up_count += 1; + Ok(TreeNodeRecursion::Continue) + } + } + + let mut visitor = TestVisitor { m: HashMap::new() }; + tree.visit(&mut visitor)?; + + let expected = HashMap::from([ + (node_f_ref, (1, 1)), + (node_e_ref, (1, 1)), + (node_c_ref, (2, 2)), + (node_d_ref, (2, 2)), + (node_b_ref, (2, 2)), + (node_a_ref, (3, 3)), + ]); + assert_eq!(visitor.m, expected); + + Ok(()) + } + }; + } + + macro_rules! test_node { + ($NAME: ident) => { + #[derive(Debug, Eq, Hash, PartialEq, Clone)] + pub struct $NAME { + pub(crate) children: Vec, + pub(crate) data: T, + } + + impl TestTree for $NAME { + fn new_with_children(children: Vec, data: T) -> Self { + Self { children, data } + } + + fn new_leaf(data: T) -> Self { + Self { + children: vec![], + data, + } + } + + fn is_leaf(&self) -> bool { + self.children.is_empty() + } + + fn eq_data(&self, other: &T) -> bool { + self.data == *other + } + + fn with_children_from(data: T, other: Self) -> Self { + Self::new_with_children(other.children, data) + } + } + }; + } + + pub mod test_tree_node { + use super::*; + + test_node!(TestTreeNode); + + impl TreeNode for TestTreeNode { + fn apply_children<'n, F: FnMut(&'n Self) -> Result>( + &'n self, + f: F, + ) -> Result { + self.children.iter().apply_until_stop(f) + } + + fn map_children Result>>( + self, + f: F, + ) -> Result> { + Ok(self + .children + .into_iter() + .map_until_stop_and_collect(f)? + .update_data(|new_children| Self { + children: new_children, + ..self + })) } } - let mut visitor = TestVisitor { m: HashMap::new() }; - tree.visit(&mut visitor)?; + node_tests!(TestTreeNode); + } + + mod test_concrete_tree_node { + use super::*; + use crate::tree_node::ConcreteTreeNode; + use std::mem; + + test_node!(ConcreteTestTreeNode); - let expected = HashMap::from([ - (node_f_ref, (1, 1)), - (node_e_ref, (1, 1)), - (node_c_ref, (2, 2)), - (node_d_ref, (2, 2)), - (node_b_ref, (2, 2)), - (node_a_ref, (3, 3)), - ]); - assert_eq!(visitor.m, expected); + impl ConcreteTreeNode for ConcreteTestTreeNode { + fn children(&self) -> Vec<&Self> { + self.children.iter().collect() + } + + fn take_children(mut self) -> (Self, Vec) { + let children = mem::take(&mut self.children); + (self, children) + } + + fn with_new_children(mut self, children: Vec) -> Result { + self.children = children; + Ok(self) + } + } - Ok(()) + node_tests!(ConcreteTestTreeNode); } } diff --git a/datafusion/physical-expr-common/src/tree_node.rs b/datafusion/physical-expr-common/src/tree_node.rs index c37e67575bf0..790e805230e1 100644 --- a/datafusion/physical-expr-common/src/tree_node.rs +++ b/datafusion/physical-expr-common/src/tree_node.rs @@ -31,11 +31,10 @@ impl DynTreeNode for dyn PhysicalExpr { } fn with_new_arc_children( - &self, - arc_self: Arc, + self: Arc, new_children: Vec>, ) -> Result> { - with_new_children_if_necessary(arc_self, new_children) + with_new_children_if_necessary(self, new_children) } } @@ -89,8 +88,8 @@ impl Display for ExprContext { } impl ConcreteTreeNode for ExprContext { - fn children(&self) -> &[Self] { - &self.children + fn children(&self) -> Vec<&Self> { + self.children.iter().collect() } fn take_children(mut self) -> (Self, Vec) { diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 96bd0de3d37c..14aa0961889a 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -31,11 +31,10 @@ impl DynTreeNode for dyn ExecutionPlan { } fn with_new_arc_children( - &self, - arc_self: Arc, + self: Arc, new_children: Vec>, ) -> Result> { - with_new_children_if_necessary(arc_self, new_children) + with_new_children_if_necessary(self, new_children) } } @@ -91,8 +90,8 @@ impl Display for PlanContext { } impl ConcreteTreeNode for PlanContext { - fn children(&self) -> &[Self] { - &self.children + fn children(&self) -> Vec<&Self> { + self.children.iter().collect() } fn take_children(mut self) -> (Self, Vec) {