From c26e07f9b87f5dc1e1aa0852b11e2d08908f3241 Mon Sep 17 00:00:00 2001 From: blaginin Date: Tue, 29 Oct 2024 20:22:53 +0000 Subject: [PATCH 01/13] Move `DynTreeNode::transform` and ``DynTreeNode::visit` to iterative move --- datafusion/common/src/cse.rs | 14 +- datafusion/common/src/tree_node.rs | 2486 ++++++++++------- .../physical-expr-common/src/tree_node.rs | 5 +- datafusion/physical-plan/src/tree_node.rs | 5 +- 4 files changed, 1528 insertions(+), 982 deletions(-) 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 563f1fa85614..be91b0a5e325 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -17,9 +17,8 @@ //! [`TreeNode`] for visiting and rewriting expression and plan trees -use std::sync::Arc; - use crate::Result; +use std::sync::Arc; /// These macros are used to determine continuation during transforming traversals. macro_rules! handle_transform_recursion { @@ -745,6 +744,16 @@ impl Transformed { Ok(self) } + pub fn on_transform_children(mut self) -> Transformed { + self.tnr = match self.tnr { + TreeNodeRecursion::Continue => TreeNodeRecursion::Jump, + TreeNodeRecursion::Jump => TreeNodeRecursion::Continue, + TreeNodeRecursion::Stop => TreeNodeRecursion::Stop, + }; + + self + } + /// Maps the [`Transformed`] object to the result of the given `f` depending on the /// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current /// node's sibling. @@ -778,6 +787,11 @@ impl Transformed { TreeNodeRecursion::Jump | TreeNodeRecursion::Stop => Ok(self), } } + + 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. @@ -946,8 +960,7 @@ pub trait DynTreeNode { /// Constructs a new node with the specified children. fn with_new_arc_children( - &self, - arc_self: Arc, + self: Arc, new_children: Vec>, ) -> Result>; } @@ -976,9 +989,8 @@ impl TreeNode for Arc { // 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) - }) + new_children + .map_data(|new_children| arc_self.with_new_arc_children(new_children)) } else { Ok(Transformed::new(self, false, new_children.tnr)) } @@ -986,6 +998,210 @@ impl TreeNode for Arc { Ok(Transformed::no(self)) } } + + fn rewrite>( + self, + rewriter: &mut R, + ) -> Result> { + let mut queue = vec![ProcessingState::NotStarted(self)]; + + while let Some(item) = queue.pop() { + match item { + ProcessingState::NotStarted(node) => { + let node = rewriter.f_down(node)?; + + queue.push(match node.tnr { + TreeNodeRecursion::Continue => { + ProcessingState::ProcessingChildren { + non_processed_children: node + .data + .arc_children() + .into_iter() + .cloned() + .rev() + .collect(), + item: node, + processed_children: vec![], + } + } + TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( + node.with_tnr(TreeNodeRecursion::Continue), + ), + TreeNodeRecursion::Stop => { + ProcessingState::ProcessedAllChildren(node) + } + }) + } + ProcessingState::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() { + queue.push(ProcessingState::ProcessingChildren { + item, + non_processed_children, + processed_children, + }); + queue.push(ProcessingState::NotStarted(non_processed_item)); + } else { + item.transformed = + processed_children.iter().any(|item| item.transformed); + item.data = item.data.with_new_arc_children( + processed_children.into_iter().map(|c| c.data).collect(), + )?; + queue.push(ProcessingState::ProcessedAllChildren(item)) + } + } + TreeNodeRecursion::Stop => { + processed_children.extend( + non_processed_children + .into_iter() + .rev() + .map(Transformed::no), + ); + item.transformed = + processed_children.iter().any(|item| item.transformed); + item.data = item.data.with_new_arc_children( + processed_children.into_iter().map(|c| c.data).collect(), + )?; + queue.push(ProcessingState::ProcessedAllChildren(item)); + } + }, + ProcessingState::ProcessedAllChildren(node) => { + let node = node.transform_parent(|n| rewriter.f_up(n))?; + + if let Some(ProcessingState::ProcessingChildren { + item: mut parent_node, + non_processed_children, + mut processed_children, + .. + }) = queue.pop() + { + parent_node.tnr = node.tnr; + processed_children.push(node); + + queue.push(ProcessingState::ProcessingChildren { + item: parent_node, + non_processed_children, + processed_children, + }) + } else { + debug_assert_eq!(queue.len(), 0); + return Ok(node); + } + } + } + } + + unreachable!(); + } + + fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>( + &'n self, + visitor: &mut V, + ) -> Result { + let mut queue = vec![VisitingState::NotStarted(self)]; + + while let Some(item) = queue.pop() { + match item { + VisitingState::NotStarted(item) => { + let tnr = visitor.f_down(item)?; + queue.push(match tnr { + TreeNodeRecursion::Continue => VisitingState::VisitingChildren { + non_processed_children: item + .arc_children() + .into_iter() + .rev() + .collect(), + item, + tnr, + }, + TreeNodeRecursion::Jump => VisitingState::VisitedAllChildren { + item, + tnr: TreeNodeRecursion::Continue, + }, + TreeNodeRecursion::Stop => VisitingState::VisitedAllChildren { + item, + tnr: TreeNodeRecursion::Stop, + }, + }); + } + VisitingState::VisitingChildren { + item, + mut non_processed_children, + tnr, + } => match tnr { + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { + if let Some(non_processed_item) = non_processed_children.pop() { + queue.push(VisitingState::VisitingChildren { + item, + non_processed_children, + tnr, + }); + queue.push(VisitingState::NotStarted(non_processed_item)); + } else { + queue.push(VisitingState::VisitedAllChildren { item, tnr }); + } + } + TreeNodeRecursion::Stop => { + return Ok(tnr); + } + }, + VisitingState::VisitedAllChildren { item, tnr } => { + let tnr = tnr.visit_parent(|| visitor.f_up(item))?; + + if let Some(VisitingState::VisitingChildren { + item, + non_processed_children, + .. + }) = queue.pop() + { + queue.push(VisitingState::VisitingChildren { + item, + non_processed_children, + tnr, + }); + } else { + debug_assert_eq!(queue.len(), 0); + return Ok(tnr); + } + } + } + } + + unreachable!() + } +} + +#[derive(Debug)] +enum ProcessingState { + NotStarted(T), + // f_down is called + ProcessingChildren { + item: Transformed, + non_processed_children: Vec, + processed_children: Vec>, + }, + ProcessedAllChildren(Transformed), + // f_up is called +} + +#[derive(Debug)] +enum VisitingState<'a, T> { + NotStarted(&'a T), + // f_down is called + VisitingChildren { + item: &'a T, + non_processed_children: Vec<&'a T>, + tnr: TreeNodeRecursion, + }, + VisitedAllChildren { + item: &'a T, + tnr: TreeNodeRecursion, + }, + // f_up is called } /// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for @@ -1028,1073 +1244,1403 @@ impl TreeNode for T { #[cfg(test)] pub(crate) mod tests { - use std::collections::HashMap; - use std::fmt::Display; - use crate::tree_node::{ - Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, TreeNodeRewriter, - TreeNodeVisitor, + DynTreeNode, Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, + TreeNodeRewriter, TreeNodeVisitor, }; use crate::Result; + use std::collections::HashMap; + use std::fmt::Display; + use std::sync::Arc; + + macro_rules! visit_test { + ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_VISITS:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + let mut visitor = TestVisitor::new(Box::new($F_DOWN), Box::new($F_UP)); + tree.visit(&mut visitor)?; + assert_eq!(visitor.visits, $EXPECTED_VISITS); - #[derive(Debug, Eq, Hash, PartialEq, Clone)] - pub struct TestTreeNode { - pub(crate) children: Vec>, - pub(crate) data: T, + Ok(()) + } + }; } - impl TestTreeNode { - pub(crate) fn new(children: Vec>, data: T) -> Self { - Self { children, data } - } + macro_rules! test_apply { + ($NAME:ident, $F:expr, $EXPECTED_VISITS:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + let mut visits = vec![]; + tree.apply(|node| { + visits.push(format!("f_down({})", node.data)); + $F(node) + })?; + assert_eq!(visits, $EXPECTED_VISITS); - pub(crate) fn new_leaf(data: T) -> Self { - Self { - children: vec![], - data, + Ok(()) } - } - - 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) - } + macro_rules! rewrite_test { + ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + let mut rewriter = TestRewriter::new(Box::new($F_DOWN), Box::new($F_UP)); + assert_eq!(tree.rewrite(&mut rewriter)?, $EXPECTED_TREE); - 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 - })) - } + Ok(()) + } + }; } - // 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()) - } + macro_rules! transform_test { + ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + assert_eq!(tree.transform_down_up($F_DOWN, $F_UP,)?, $EXPECTED_TREE); - // 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() + Ok(()) + } + }; } - // 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()) - } + macro_rules! transform_down_test { + ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + assert_eq!(tree.transform_down($F)?, $EXPECTED_TREE); - // 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()) + Ok(()) + } + }; } - // 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()) - } + macro_rules! transform_up_test { + ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => { + #[test] + fn $NAME() -> Result<()> { + let tree = test_tree(); + assert_eq!(tree.transform_up($F)?, $EXPECTED_TREE); - // 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() + Ok(()) + } + }; } - 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()) + fn down_visits(visits: Vec) -> Vec { + visits + .into_iter() + .filter(|v| v.starts_with("f_down")) + .collect() } - // 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() - } + pub(crate) trait TestTree + where + T: Sized, + { + fn new_with_children(children: Vec, data: T) -> Self + where + Self: Sized; + + fn new_leaf(data: T) -> Self; + + fn is_leaf(&self) -> bool; + } + + macro_rules! gen_tests { + ($TYPE:ident) => { + // 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()) + } - 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()) - } + // 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() + } - 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()) - } + // 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()) + } - // 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() - } + // 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()) + } - 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()) - } + // 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()) + } - 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_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() + } - // 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_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()) + } - fn f_up_jump_on_e_transformed_tree() -> TestTreeNode { - transformed_tree() - } + // 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_up_jump_on_e_transformed_up_tree() -> TestTreeNode { - transformed_up_tree() - } + 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()) + } - // 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_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()) + } - 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()) - } + // 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_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()) - } + 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()) + } - // 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_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()) + } - 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()) - } + // 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_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()) - } + fn f_up_jump_on_e_transformed_tree() -> $TYPE { + transformed_tree() + } - // 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_jump_on_e_transformed_up_tree() -> $TYPE { + transformed_up_tree() + } - 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()) - } + // 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_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()) - } + 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()) + } - // 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_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()) + } - 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()) - } + // 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_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 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 down_visits(visits: Vec) -> Vec { - visits - .into_iter() - .filter(|v| v.starts_with("f_down")) - .collect() - } + 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()) + } - type TestVisitorF = Box) -> Result>; + // 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() + } - struct TestVisitor { - visits: Vec, - f_down: TestVisitorF, - f_up: TestVisitorF, - } + 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()) + } - impl TestVisitor { - fn new(f_down: TestVisitorF, f_up: TestVisitorF) -> Self { - Self { - visits: vec![], - f_down, - f_up, + 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()) } - } - } - impl<'n, T: Display> TreeNodeVisitor<'n> for TestVisitor { - type Node = TestTreeNode; + // 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_down(&mut self, node: &'n Self::Node) -> Result { - self.visits.push(format!("f_down({})", node.data)); - (*self.f_down)(node) - } + 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(&mut self, node: &'n Self::Node) -> Result { - self.visits.push(format!("f_up({})", node.data)); - (*self.f_up)(node) - } - } + 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()) + } - fn visit_continue(_: &TestTreeNode) -> Result { - Ok(TreeNodeRecursion::Continue) - } + 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() + ); - 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 - }) - } - } + 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()) + ); - macro_rules! visit_test { - ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_VISITS:expr) => { + 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 $NAME() -> Result<()> { - let tree = test_tree(); - let mut visitor = TestVisitor::new(Box::new($F_DOWN), Box::new($F_UP)); + 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)?; - assert_eq!(visitor.visits, $EXPECTED_VISITS); + + 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_apply { - ($NAME:ident, $F:expr, $EXPECTED_VISITS:expr) => { - #[test] - fn $NAME() -> Result<()> { - let tree = test_tree(); - let mut visits = vec![]; - tree.apply(|node| { - visits.push(format!("f_down({})", node.data)); - $F(node) - })?; - assert_eq!(visits, $EXPECTED_VISITS); + pub mod test_tree_node { + use super::*; - Ok(()) + #[derive(Debug, Eq, Hash, PartialEq, Clone)] + pub struct TestTreeNode { + pub(crate) children: Vec>, + pub(crate) data: T, + } + + impl TestTree for TestTreeNode { + fn new_with_children(children: Vec>, data: T) -> Self { + Self { children, data } } - }; - } - type TestRewriterF = - Box) -> Result>>>; + fn new_leaf(data: T) -> Self { + Self { + children: vec![], + data, + } + } - struct TestRewriter { - f_down: TestRewriterF, - f_up: TestRewriterF, - } + fn is_leaf(&self) -> bool { + self.children.is_empty() + } + } - impl TestRewriter { - fn new(f_down: TestRewriterF, f_up: TestRewriterF) -> Self { - Self { f_down, f_up } + 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 + })) + } + } + + fn visit_continue(_: &TestTreeNode) -> Result { + Ok(TreeNodeRecursion::Continue) } - } - impl TreeNodeRewriter for TestRewriter { - type Node = TestTreeNode; + type TestVisitorF = + Box) -> Result>; - fn f_down(&mut self, node: Self::Node) -> Result> { - (*self.f_down)(node) + struct TestVisitor { + visits: Vec, + f_down: TestVisitorF, + f_up: TestVisitorF, } - fn f_up(&mut self, node: Self::Node) -> Result> { - (*self.f_up)(node) + impl TestVisitor { + fn new(f_down: TestVisitorF, f_up: TestVisitorF) -> Self { + Self { + visits: vec![], + f_down, + f_up, + } + } } - } - 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(), - ))) + impl<'n, T: Display> TreeNodeVisitor<'n> for TestVisitor { + type Node = TestTreeNode; + + 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 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) - }) + 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 + }) + } } - } - macro_rules! rewrite_test { - ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => { - #[test] - fn $NAME() -> Result<()> { - let tree = test_tree(); - let mut rewriter = TestRewriter::new(Box::new($F_DOWN), Box::new($F_UP)); - assert_eq!(tree.rewrite(&mut rewriter)?, $EXPECTED_TREE); + type TestRewriterF = + Box) -> Result>>>; - Ok(()) + struct TestRewriter { + f_down: TestRewriterF, + f_up: TestRewriterF, + } + + impl TestRewriter { + fn new(f_down: TestRewriterF, f_up: TestRewriterF) -> Self { + Self { f_down, f_up } } - }; - } + } - macro_rules! transform_test { - ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => { - #[test] - fn $NAME() -> Result<()> { - let tree = test_tree(); - assert_eq!(tree.transform_down_up($F_DOWN, $F_UP,)?, $EXPECTED_TREE); + impl TreeNodeRewriter for TestRewriter { + type Node = TestTreeNode; - Ok(()) + fn f_down(&mut self, node: Self::Node) -> Result> { + (*self.f_down)(node) } - }; - } - macro_rules! transform_down_test { - ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => { - #[test] - fn $NAME() -> Result<()> { - let tree = test_tree(); - assert_eq!(tree.transform_down($F)?, $EXPECTED_TREE); + fn f_up(&mut self, node: Self::Node) -> Result> { + (*self.f_up)(node) + } + } - Ok(()) + fn transform_yes>( + transformation_name: N, + ) -> impl FnMut(TestTreeNode) -> Result>> { + move |node| { + Ok(Transformed::yes(TestTreeNode::new_with_children( + 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_with_children( + node.children, + format!("{}({})", transformation_name, node.data).into(), + ); + Ok(if node.data == d { + Transformed::new(new_node, true, event) + } else { + Transformed::yes(new_node) + }) + } + } + + gen_tests!(TestTreeNode); } - macro_rules! transform_up_test { - ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => { - #[test] - fn $NAME() -> Result<()> { - let tree = test_tree(); - assert_eq!(tree.transform_up($F)?, $EXPECTED_TREE); + pub mod arc_test_tree_node { + use super::*; - Ok(()) + #[derive(Debug, Eq, Hash, PartialEq, Clone)] + pub struct DynTestNode { + pub(crate) data: Arc, + pub(crate) children: Vec>>, + } + + impl TestTree for Arc> { + fn new_with_children(children: Vec, data: T) -> Self + where + Self: Sized, + { + Arc::new(DynTestNode { + data: Arc::new(data), + children, + }) } - }; - } - 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; + fn new_leaf(data: T) -> Self { + Arc::new(DynTestNode { + data: Arc::new(data), + children: vec![], + }) + } + + fn is_leaf(&self) -> bool { + self.children.is_empty() + } + } + + impl DynTreeNode for DynTestNode { + fn arc_children(&self) -> Vec<&Arc> { + self.children.iter().collect() + } + + fn with_new_arc_children( + self: Arc, + new_children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + children: new_children, + data: Arc::clone(&self.data), + })) + } + } + + type ArcTestNode = Arc>; + + fn visit_continue(_: &ArcTestNode) -> Result { 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)>, } - impl<'n> TreeNodeVisitor<'n> for TestVisitor<'n> { - type Node = TestTreeNode; + 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> TreeNodeVisitor<'n> for TestVisitor { + type Node = ArcTestNode; 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) + self.visits.push(format!("f_down({})", node.data)); + (*self.f_down)(node) } 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) + self.visits.push(format!("f_up({})", node.data)); + (*self.f_up)(node) + } + } + + fn visit_event_on>( + data: D, + event: TreeNodeRecursion, + ) -> impl FnMut(&ArcTestNode) -> Result { + let d = data.into(); + move |node| { + Ok(if node.data.as_ref() == &d { + event + } else { + TreeNodeRecursion::Continue + }) } } - let mut visitor = TestVisitor { m: HashMap::new() }; - tree.visit(&mut visitor)?; + type TestRewriterF = + Box) -> Result>>>; - 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); + 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 = ArcTestNode; + + 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(ArcTestNode) -> Result>> { + move |node| { + Ok(Transformed::yes(ArcTestNode::new_with_children( + node.children.clone(), + 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(ArcTestNode) -> Result>> { + let d = data.into(); + move |node| { + let new_node = ArcTestNode::new_with_children( + node.children.clone(), + format!("{}({})", transformation_name, node.data).into(), + ); + Ok(if node.data.as_ref() == &d { + Transformed::new(new_node, true, event) + } else { + Transformed::yes(new_node) + }) + } + } - Ok(()) + gen_tests!(ArcTestNode); + + #[test] + fn test_recursion() { + // todo: switch to test gen + + let mut item = ArcTestNode::new_leaf("initial".to_string()); + for i in 0..3000 { + item = + ArcTestNode::new_with_children(vec![item], format!("parent-{}", i)); + } + + let mut visitor = + TestVisitor::new(Box::new(visit_continue), Box::new(visit_continue)); + + item.visit(&mut visitor).unwrap(); + } } } diff --git a/datafusion/physical-expr-common/src/tree_node.rs b/datafusion/physical-expr-common/src/tree_node.rs index d9892ce55509..0060546804c7 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) } } diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index 96bd0de3d37c..96a586209f15 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) } } From 6a4dd2c5eee28a91db0a989be73fe6650f547b99 Mon Sep 17 00:00:00 2001 From: blaginin Date: Tue, 29 Oct 2024 20:59:27 +0000 Subject: [PATCH 02/13] Rewrite all other methods in DynTreeNode to be iterative --- datafusion/common/src/tree_node.rs | 216 ++++++++++++++++------------- 1 file changed, 123 insertions(+), 93 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index be91b0a5e325..09a4ad619a48 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -912,6 +912,104 @@ macro_rules! map_until_stop_and_collect { }} } +macro_rules! rewrite_recursive { + ($START:ident, $NAME:ident, $TRANSFORM_UP:expr, $TRANSFORM_DOWN:expr) => { + let mut queue = vec![ProcessingState::NotStarted($START)]; + + while let Some(item) = queue.pop() { + match item { + ProcessingState::NotStarted($NAME) => { + let node = $TRANSFORM_DOWN?; + + queue.push(match node.tnr { + TreeNodeRecursion::Continue => { + ProcessingState::ProcessingChildren { + non_processed_children: node + .data + .arc_children() + .into_iter() + .cloned() + .rev() + .collect(), + item: node, + processed_children: vec![], + } + } + TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( + node.with_tnr(TreeNodeRecursion::Continue), + ), + TreeNodeRecursion::Stop => { + ProcessingState::ProcessedAllChildren(node) + } + }) + } + ProcessingState::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() { + queue.push(ProcessingState::ProcessingChildren { + item, + non_processed_children, + processed_children, + }); + queue.push(ProcessingState::NotStarted(non_processed_item)); + } else { + item.transformed |= + processed_children.iter().any(|item| item.transformed); + item.data = item.data.with_new_arc_children( + processed_children.into_iter().map(|c| c.data).collect(), + )?; + queue.push(ProcessingState::ProcessedAllChildren(item)) + } + } + TreeNodeRecursion::Stop => { + processed_children.extend( + non_processed_children + .into_iter() + .rev() + .map(Transformed::no), + ); + item.transformed |= + processed_children.iter().any(|item| item.transformed); + item.data = item.data.with_new_arc_children( + processed_children.into_iter().map(|c| c.data).collect(), + )?; + queue.push(ProcessingState::ProcessedAllChildren(item)); + } + }, + ProcessingState::ProcessedAllChildren(node) => { + let node = node.transform_parent(|$NAME| $TRANSFORM_UP)?; + + if let Some(ProcessingState::ProcessingChildren { + item: mut parent_node, + non_processed_children, + mut processed_children, + .. + }) = queue.pop() + { + parent_node.tnr = node.tnr; + processed_children.push(node); + + queue.push(ProcessingState::ProcessingChildren { + item: parent_node, + non_processed_children, + processed_children, + }) + } else { + debug_assert_eq!(queue.len(), 0); + return Ok(node); + } + } + } + } + + unreachable!(); + }; +} + /// Transformation helper to access [`Transformed`] fields in a [`Result`] easily. /// /// # Example @@ -999,103 +1097,35 @@ impl TreeNode for Arc { } } - fn rewrite>( + fn transform_down_up< + FD: FnMut(Self) -> Result>, + FU: FnMut(Self) -> Result>, + >( self, - rewriter: &mut R, + mut f_down: FD, + mut f_up: FU, ) -> Result> { - let mut queue = vec![ProcessingState::NotStarted(self)]; - - while let Some(item) = queue.pop() { - match item { - ProcessingState::NotStarted(node) => { - let node = rewriter.f_down(node)?; - - queue.push(match node.tnr { - TreeNodeRecursion::Continue => { - ProcessingState::ProcessingChildren { - non_processed_children: node - .data - .arc_children() - .into_iter() - .cloned() - .rev() - .collect(), - item: node, - processed_children: vec![], - } - } - TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( - node.with_tnr(TreeNodeRecursion::Continue), - ), - TreeNodeRecursion::Stop => { - ProcessingState::ProcessedAllChildren(node) - } - }) - } - ProcessingState::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() { - queue.push(ProcessingState::ProcessingChildren { - item, - non_processed_children, - processed_children, - }); - queue.push(ProcessingState::NotStarted(non_processed_item)); - } else { - item.transformed = - processed_children.iter().any(|item| item.transformed); - item.data = item.data.with_new_arc_children( - processed_children.into_iter().map(|c| c.data).collect(), - )?; - queue.push(ProcessingState::ProcessedAllChildren(item)) - } - } - TreeNodeRecursion::Stop => { - processed_children.extend( - non_processed_children - .into_iter() - .rev() - .map(Transformed::no), - ); - item.transformed = - processed_children.iter().any(|item| item.transformed); - item.data = item.data.with_new_arc_children( - processed_children.into_iter().map(|c| c.data).collect(), - )?; - queue.push(ProcessingState::ProcessedAllChildren(item)); - } - }, - ProcessingState::ProcessedAllChildren(node) => { - let node = node.transform_parent(|n| rewriter.f_up(n))?; - - if let Some(ProcessingState::ProcessingChildren { - item: mut parent_node, - non_processed_children, - mut processed_children, - .. - }) = queue.pop() - { - parent_node.tnr = node.tnr; - processed_children.push(node); + rewrite_recursive!(self, node, f_up(node), f_down(node)); + } - queue.push(ProcessingState::ProcessingChildren { - item: parent_node, - non_processed_children, - processed_children, - }) - } else { - debug_assert_eq!(queue.len(), 0); - return Ok(node); - } - } - } - } + fn transform_down Result>>( + self, + f: F, + ) -> Result> { + self.transform_down_up(f, |node| Ok(Transformed::no(node))) + } - unreachable!(); + fn transform_up Result>>( + self, + f: F, + ) -> Result> { + self.transform_down_up(|node| Ok(Transformed::no(node)), f) + } + fn rewrite>( + self, + rewriter: &mut R, + ) -> Result> { + rewrite_recursive!(self, node, rewriter.f_up(node), rewriter.f_down(node)); } fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>( From 83d194e84a112d6e1fca41fd8571a6038efbef57 Mon Sep 17 00:00:00 2001 From: blaginin Date: Tue, 29 Oct 2024 21:40:11 +0000 Subject: [PATCH 03/13] Remove macros in favour of `LegacyRewriter` --- datafusion/common/src/tree_node.rs | 283 +++++++++++++++++------------ 1 file changed, 166 insertions(+), 117 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 09a4ad619a48..250d25a8e98d 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -18,6 +18,7 @@ //! [`TreeNode`] for visiting and rewriting expression and plan trees use crate::Result; +use std::marker::PhantomData; use std::sync::Arc; /// These macros are used to determine continuation during transforming traversals. @@ -912,104 +913,6 @@ macro_rules! map_until_stop_and_collect { }} } -macro_rules! rewrite_recursive { - ($START:ident, $NAME:ident, $TRANSFORM_UP:expr, $TRANSFORM_DOWN:expr) => { - let mut queue = vec![ProcessingState::NotStarted($START)]; - - while let Some(item) = queue.pop() { - match item { - ProcessingState::NotStarted($NAME) => { - let node = $TRANSFORM_DOWN?; - - queue.push(match node.tnr { - TreeNodeRecursion::Continue => { - ProcessingState::ProcessingChildren { - non_processed_children: node - .data - .arc_children() - .into_iter() - .cloned() - .rev() - .collect(), - item: node, - processed_children: vec![], - } - } - TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( - node.with_tnr(TreeNodeRecursion::Continue), - ), - TreeNodeRecursion::Stop => { - ProcessingState::ProcessedAllChildren(node) - } - }) - } - ProcessingState::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() { - queue.push(ProcessingState::ProcessingChildren { - item, - non_processed_children, - processed_children, - }); - queue.push(ProcessingState::NotStarted(non_processed_item)); - } else { - item.transformed |= - processed_children.iter().any(|item| item.transformed); - item.data = item.data.with_new_arc_children( - processed_children.into_iter().map(|c| c.data).collect(), - )?; - queue.push(ProcessingState::ProcessedAllChildren(item)) - } - } - TreeNodeRecursion::Stop => { - processed_children.extend( - non_processed_children - .into_iter() - .rev() - .map(Transformed::no), - ); - item.transformed |= - processed_children.iter().any(|item| item.transformed); - item.data = item.data.with_new_arc_children( - processed_children.into_iter().map(|c| c.data).collect(), - )?; - queue.push(ProcessingState::ProcessedAllChildren(item)); - } - }, - ProcessingState::ProcessedAllChildren(node) => { - let node = node.transform_parent(|$NAME| $TRANSFORM_UP)?; - - if let Some(ProcessingState::ProcessingChildren { - item: mut parent_node, - non_processed_children, - mut processed_children, - .. - }) = queue.pop() - { - parent_node.tnr = node.tnr; - processed_children.push(node); - - queue.push(ProcessingState::ProcessingChildren { - item: parent_node, - non_processed_children, - processed_children, - }) - } else { - debug_assert_eq!(queue.len(), 0); - return Ok(node); - } - } - } - } - - unreachable!(); - }; -} - /// Transformation helper to access [`Transformed`] fields in a [`Result`] easily. /// /// # Example @@ -1063,6 +966,59 @@ pub trait DynTreeNode { ) -> Result>; } +pub struct LegacyRewriter< + 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, + > LegacyRewriter +{ + 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 LegacyRewriter +{ + type Node = Node; + + fn f_down(&mut self, node: Self::Node) -> Result> { + (self.f_down_func)(node) + } + + fn f_up(&mut self, node: Self::Node) -> Result> { + (self.f_up_func)(node) + } +} + +macro_rules! update_rec_node { + ($NAME:ident, $CHILDREN:ident) => {{ + $NAME.transformed |= $CHILDREN.iter().any(|item| item.transformed); + + $NAME.data = $NAME + .data + .with_new_arc_children($CHILDREN.into_iter().map(|c| c.data).collect())?; + + $NAME + }}; +} + /// Blanket implementation for any `Arc` where `T` implements [`DynTreeNode`] /// (such as [`Arc`]). impl TreeNode for Arc { @@ -1102,43 +1058,134 @@ impl TreeNode for Arc { FU: FnMut(Self) -> Result>, >( self, - mut f_down: FD, - mut f_up: FU, + f_down: FD, + f_up: FU, ) -> Result> { - rewrite_recursive!(self, node, f_up(node), f_down(node)); + self.rewrite(&mut LegacyRewriter::new(f_down, f_up)) } fn transform_down Result>>( self, f: F, ) -> Result> { - self.transform_down_up(f, |node| Ok(Transformed::no(node))) + self.rewrite(&mut LegacyRewriter::new(f, |node| { + Ok(Transformed::no(node)) + })) } fn transform_up Result>>( self, f: F, ) -> Result> { - self.transform_down_up(|node| Ok(Transformed::no(node)), f) + self.rewrite(&mut LegacyRewriter::new( + |node| Ok(Transformed::no(node)), + f, + )) } fn rewrite>( self, rewriter: &mut R, ) -> Result> { - rewrite_recursive!(self, node, rewriter.f_up(node), rewriter.f_down(node)); + let mut stack = vec![ProcessingState::NotStarted(self)]; + + while let Some(item) = stack.pop() { + match item { + ProcessingState::NotStarted(node) => { + let node = rewriter.f_down(node)?; + + stack.push(match node.tnr { + TreeNodeRecursion::Continue => { + ProcessingState::ProcessingChildren { + non_processed_children: node + .data + .arc_children() + .into_iter() + .cloned() + .rev() + .collect(), + item: node, + processed_children: vec![], + } + } + TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( + node.with_tnr(TreeNodeRecursion::Continue), + ), + TreeNodeRecursion::Stop => { + ProcessingState::ProcessedAllChildren(node) + } + }) + } + ProcessingState::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.push(ProcessingState::ProcessingChildren { + item, + non_processed_children, + processed_children, + }); + stack.push(ProcessingState::NotStarted(non_processed_item)); + } else { + stack.push(ProcessingState::ProcessedAllChildren( + update_rec_node!(item, processed_children), + )) + } + } + TreeNodeRecursion::Stop => { + processed_children.extend( + non_processed_children + .into_iter() + .rev() + .map(Transformed::no), + ); + stack.push(ProcessingState::ProcessedAllChildren( + update_rec_node!(item, processed_children), + )); + } + }, + ProcessingState::ProcessedAllChildren(node) => { + let node = node.transform_parent(|n| rewriter.f_up(n))?; + + if let Some(ProcessingState::ProcessingChildren { + item: mut parent_node, + non_processed_children, + mut processed_children, + .. + }) = stack.pop() + { + parent_node.tnr = node.tnr; + processed_children.push(node); + + stack.push(ProcessingState::ProcessingChildren { + item: parent_node, + non_processed_children, + processed_children, + }) + } else { + debug_assert_eq!(stack.len(), 0); + return Ok(node); + } + } + } + } + + unreachable!(); } fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>( &'n self, visitor: &mut V, ) -> Result { - let mut queue = vec![VisitingState::NotStarted(self)]; + let mut stack = vec![VisitingState::NotStarted(self)]; - while let Some(item) = queue.pop() { + while let Some(item) = stack.pop() { match item { VisitingState::NotStarted(item) => { let tnr = visitor.f_down(item)?; - queue.push(match tnr { + stack.push(match tnr { TreeNodeRecursion::Continue => VisitingState::VisitingChildren { non_processed_children: item .arc_children() @@ -1165,14 +1212,14 @@ impl TreeNode for Arc { } => match tnr { TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { if let Some(non_processed_item) = non_processed_children.pop() { - queue.push(VisitingState::VisitingChildren { + stack.push(VisitingState::VisitingChildren { item, non_processed_children, tnr, }); - queue.push(VisitingState::NotStarted(non_processed_item)); + stack.push(VisitingState::NotStarted(non_processed_item)); } else { - queue.push(VisitingState::VisitedAllChildren { item, tnr }); + stack.push(VisitingState::VisitedAllChildren { item, tnr }); } } TreeNodeRecursion::Stop => { @@ -1186,15 +1233,15 @@ impl TreeNode for Arc { item, non_processed_children, .. - }) = queue.pop() + }) = stack.pop() { - queue.push(VisitingState::VisitingChildren { + stack.push(VisitingState::VisitingChildren { item, non_processed_children, tnr, }); } else { - debug_assert_eq!(queue.len(), 0); + debug_assert_eq!(stack.len(), 0); return Ok(tnr); } } @@ -1208,30 +1255,32 @@ impl TreeNode for Arc { #[derive(Debug)] enum ProcessingState { NotStarted(T), - // f_down is called + // ← at this point, f_down is called ProcessingChildren { item: Transformed, non_processed_children: Vec, processed_children: Vec>, }, + // ← at this point, all children are processed ProcessedAllChildren(Transformed), - // f_up is called + // ← at this point, f_up is called } #[derive(Debug)] enum VisitingState<'a, T> { NotStarted(&'a T), - // f_down is called + // ← at this point, f_down is called VisitingChildren { item: &'a T, non_processed_children: Vec<&'a T>, tnr: TreeNodeRecursion, }, + // ← at this point, all children are visited VisitedAllChildren { item: &'a T, tnr: TreeNodeRecursion, }, - // f_up is called + // ← at this point, f_up is called } /// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for From 482546b66b11b15dcce2999ef569d3ca7b78e010 Mon Sep 17 00:00:00 2001 From: blaginin Date: Wed, 30 Oct 2024 18:49:20 +0000 Subject: [PATCH 04/13] Remove unused method --- datafusion/common/src/tree_node.rs | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 250d25a8e98d..3f1ac3d0cd5d 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -745,16 +745,6 @@ impl Transformed { Ok(self) } - pub fn on_transform_children(mut self) -> Transformed { - self.tnr = match self.tnr { - TreeNodeRecursion::Continue => TreeNodeRecursion::Jump, - TreeNodeRecursion::Jump => TreeNodeRecursion::Continue, - TreeNodeRecursion::Stop => TreeNodeRecursion::Stop, - }; - - self - } - /// Maps the [`Transformed`] object to the result of the given `f` depending on the /// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current /// node's sibling. From f682be3ec4a42205688dd2937c34092d5487e56c Mon Sep 17 00:00:00 2001 From: blaginin Date: Wed, 30 Oct 2024 19:32:55 +0000 Subject: [PATCH 05/13] Remove `TreeNodeVisitor` / `TestRewriter` repetition --- datafusion/common/src/tree_node.rs | 367 +++++++++++------------------ 1 file changed, 140 insertions(+), 227 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 3f1ac3d0cd5d..146fcbf220e0 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -1409,7 +1409,8 @@ pub(crate) mod tests { .collect() } - pub(crate) trait TestTree + /// Implement this train in order for your struct to be supported in parametrisation + pub(crate) trait TestTree where T: Sized, { @@ -1420,10 +1421,128 @@ pub(crate) mod tests { fn new_leaf(data: T) -> Self; fn is_leaf(&self) -> bool; + + fn eq_data(&self, other: &T) -> bool; + + fn with_children_from(data: T, other: Self) -> Self; } macro_rules! gen_tests { ($TYPE:ident) => { + fn visit_continue(_: &$TYPE) -> Result { + 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 @@ -2374,12 +2493,12 @@ pub(crate) mod tests { use super::*; #[derive(Debug, Eq, Hash, PartialEq, Clone)] - pub struct TestTreeNode { + pub struct TestTreeNode { pub(crate) children: Vec>, pub(crate) data: T, } - impl TestTree for TestTreeNode { + impl TestTree for TestTreeNode { fn new_with_children(children: Vec>, data: T) -> Self { Self { children, data } } @@ -2394,9 +2513,17 @@ pub(crate) mod tests { 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) + } } - impl TreeNode for TestTreeNode { + impl TreeNode for TestTreeNode { fn apply_children<'n, F: FnMut(&'n Self) -> Result>( &'n self, f: F, @@ -2419,117 +2546,6 @@ pub(crate) mod tests { } } - fn visit_continue(_: &TestTreeNode) -> Result { - 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> TreeNodeVisitor<'n> for TestVisitor { - type Node = TestTreeNode; - - 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(&TestTreeNode) -> Result { - let d = data.into(); - move |node| { - Ok(if node.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 = 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_with_children( - 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_with_children( - node.children, - format!("{}({})", transformation_name, node.data).into(), - ); - Ok(if node.data == d { - Transformed::new(new_node, true, event) - } else { - Transformed::yes(new_node) - }) - } - } - gen_tests!(TestTreeNode); } @@ -2542,7 +2558,7 @@ pub(crate) mod tests { pub(crate) children: Vec>>, } - impl TestTree for Arc> { + impl TestTree for Arc> { fn new_with_children(children: Vec, data: T) -> Self where Self: Sized, @@ -2563,6 +2579,14 @@ pub(crate) mod tests { fn is_leaf(&self) -> bool { self.children.is_empty() } + + fn eq_data(&self, other: &T) -> bool { + self.data.as_ref().eq(other) + } + + fn with_children_from(data: T, other: Self) -> Self { + Self::new_with_children(other.children.clone(), data) + } } impl DynTreeNode for DynTestNode { @@ -2583,117 +2607,6 @@ pub(crate) mod tests { type ArcTestNode = Arc>; - fn visit_continue(_: &ArcTestNode) -> Result { - 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> TreeNodeVisitor<'n> for TestVisitor { - type Node = ArcTestNode; - - 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(&ArcTestNode) -> Result { - let d = data.into(); - move |node| { - Ok(if node.data.as_ref() == &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 = ArcTestNode; - - 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(ArcTestNode) -> Result>> { - move |node| { - Ok(Transformed::yes(ArcTestNode::new_with_children( - node.children.clone(), - 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(ArcTestNode) -> Result>> { - let d = data.into(); - move |node| { - let new_node = ArcTestNode::new_with_children( - node.children.clone(), - format!("{}({})", transformation_name, node.data).into(), - ); - Ok(if node.data.as_ref() == &d { - Transformed::new(new_node, true, event) - } else { - Transformed::yes(new_node) - }) - } - } - gen_tests!(ArcTestNode); #[test] From 4e1b81e31ce3c901315ad0e620c413fd61cc0ffb Mon Sep 17 00:00:00 2001 From: blaginin Date: Wed, 30 Oct 2024 19:57:52 +0000 Subject: [PATCH 06/13] Rename mods and test --- datafusion/common/src/tree_node.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 146fcbf220e0..1bd36e16afee 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -2549,7 +2549,7 @@ pub(crate) mod tests { gen_tests!(TestTreeNode); } - pub mod arc_test_tree_node { + pub mod test_dyn_tree_node { use super::*; #[derive(Debug, Eq, Hash, PartialEq, Clone)] @@ -2610,9 +2610,7 @@ pub(crate) mod tests { gen_tests!(ArcTestNode); #[test] - fn test_recursion() { - // todo: switch to test gen - + fn test_large_tree() { let mut item = ArcTestNode::new_leaf("initial".to_string()); for i in 0..3000 { item = From e475464892a2262cdab3a2cf52d8ebdb71e36ec0 Mon Sep 17 00:00:00 2001 From: blaginin Date: Wed, 30 Oct 2024 20:40:35 +0000 Subject: [PATCH 07/13] Improve readability --- datafusion/common/src/tree_node.rs | 138 +++++++++++++++-------------- 1 file changed, 70 insertions(+), 68 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 1bd36e16afee..fc322d622da5 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -779,6 +779,7 @@ impl Transformed { } } + /// Replaces recursion state with the given one pub fn with_tnr(mut self, tnr: TreeNodeRecursion) -> Self { self.tnr = tnr; self @@ -956,7 +957,8 @@ pub trait DynTreeNode { ) -> Result>; } -pub struct LegacyRewriter< +/// Adapter from the old function-based rewriter to the new Transformer one +struct FuncRewriter< FD: FnMut(Node) -> Result>, FU: FnMut(Node) -> Result>, Node: TreeNode, @@ -970,7 +972,7 @@ impl< FD: FnMut(Node) -> Result>, FU: FnMut(Node) -> Result>, Node: TreeNode, - > LegacyRewriter + > FuncRewriter { pub fn new(f_down_func: FD, f_up_func: FU) -> Self { Self { @@ -984,7 +986,7 @@ impl< FD: FnMut(Node) -> Result>, FU: FnMut(Node) -> Result>, Node: TreeNode, - > TreeNodeRewriter for LegacyRewriter + > TreeNodeRewriter for FuncRewriter { type Node = Node; @@ -997,7 +999,8 @@ impl< } } -macro_rules! update_rec_node { +/// Replaces node's children and recomputes the state +macro_rules! update_node_after_recursion { ($NAME:ident, $CHILDREN:ident) => {{ $NAME.transformed |= $CHILDREN.iter().any(|item| item.transformed); @@ -1011,6 +1014,7 @@ macro_rules! update_rec_node { /// Blanket implementation for any `Arc` where `T` implements [`DynTreeNode`] /// (such as [`Arc`]). +/// Unlike [`TreeNode`], performs node traversal iteratively rather than recursively to avoid stack overflow impl TreeNode for Arc { fn apply_children<'n, F: FnMut(&'n Self) -> Result>( &'n self, @@ -1051,41 +1055,36 @@ impl TreeNode for Arc { f_down: FD, f_up: FU, ) -> Result> { - self.rewrite(&mut LegacyRewriter::new(f_down, f_up)) + self.rewrite(&mut FuncRewriter::new(f_down, f_up)) } fn transform_down Result>>( self, f: F, ) -> Result> { - self.rewrite(&mut LegacyRewriter::new(f, |node| { - Ok(Transformed::no(node)) - })) + self.rewrite(&mut FuncRewriter::new(f, |node| Ok(Transformed::no(node)))) } fn transform_up Result>>( self, f: F, ) -> Result> { - self.rewrite(&mut LegacyRewriter::new( - |node| Ok(Transformed::no(node)), - f, - )) + self.rewrite(&mut FuncRewriter::new(|node| Ok(Transformed::no(node)), f)) } fn rewrite>( self, rewriter: &mut R, ) -> Result> { - let mut stack = vec![ProcessingState::NotStarted(self)]; + let mut stack = vec![TransformingState::NotStarted(self)]; while let Some(item) = stack.pop() { match item { - ProcessingState::NotStarted(node) => { + TransformingState::NotStarted(node) => { let node = rewriter.f_down(node)?; stack.push(match node.tnr { TreeNodeRecursion::Continue => { - ProcessingState::ProcessingChildren { + TransformingState::ProcessingChildren { non_processed_children: node .data .arc_children() @@ -1097,59 +1096,68 @@ impl TreeNode for Arc { processed_children: vec![], } } - TreeNodeRecursion::Jump => ProcessingState::ProcessedAllChildren( - node.with_tnr(TreeNodeRecursion::Continue), - ), + TreeNodeRecursion::Jump => { + TransformingState::ProcessedAllChildren( + // No need to process children, we can just this stage + node.with_tnr(TreeNodeRecursion::Continue), + ) + } TreeNodeRecursion::Stop => { - ProcessingState::ProcessedAllChildren(node) + TransformingState::ProcessedAllChildren(node) } }) } - ProcessingState::ProcessingChildren { + 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.push(ProcessingState::ProcessingChildren { - item, - non_processed_children, - processed_children, - }); - stack.push(ProcessingState::NotStarted(non_processed_item)); + stack.extend([ + // This node still has children, so put it back in 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 { - stack.push(ProcessingState::ProcessedAllChildren( - update_rec_node!(item, processed_children), + stack.push(TransformingState::ProcessedAllChildren( + update_node_after_recursion!(item, processed_children), )) } } TreeNodeRecursion::Stop => { + // At this point, we might have some children we haven't yet processed processed_children.extend( non_processed_children .into_iter() .rev() .map(Transformed::no), ); - stack.push(ProcessingState::ProcessedAllChildren( - update_rec_node!(item, processed_children), + stack.push(TransformingState::ProcessedAllChildren( + update_node_after_recursion!(item, processed_children), )); } }, - ProcessingState::ProcessedAllChildren(node) => { + TransformingState::ProcessedAllChildren(node) => { let node = node.transform_parent(|n| rewriter.f_up(n))?; - if let Some(ProcessingState::ProcessingChildren { + if let Some(TransformingState::ProcessingChildren { item: mut parent_node, non_processed_children, mut processed_children, .. }) = stack.pop() { + // We need use returned recursion state when processing the remaining children parent_node.tnr = node.tnr; processed_children.push(node); - stack.push(ProcessingState::ProcessingChildren { + stack.push(TransformingState::ProcessingChildren { item: parent_node, non_processed_children, processed_children, @@ -1189,10 +1197,7 @@ impl TreeNode for Arc { item, tnr: TreeNodeRecursion::Continue, }, - TreeNodeRecursion::Stop => VisitingState::VisitedAllChildren { - item, - tnr: TreeNodeRecursion::Stop, - }, + TreeNodeRecursion::Stop => return Ok(tnr), }); } VisitingState::VisitingChildren { @@ -1202,12 +1207,15 @@ impl TreeNode for Arc { } => match tnr { TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { if let Some(non_processed_item) = non_processed_children.pop() { - stack.push(VisitingState::VisitingChildren { - item, - non_processed_children, - tnr, - }); - stack.push(VisitingState::NotStarted(non_processed_item)); + stack.extend([ + // Returning the node on the stack because there are more children to process + VisitingState::VisitingChildren { + item, + non_processed_children, + tnr, + }, + VisitingState::NotStarted(non_processed_item), + ]); } else { stack.push(VisitingState::VisitedAllChildren { item, tnr }); } @@ -1220,10 +1228,10 @@ impl TreeNode for Arc { let tnr = tnr.visit_parent(|| visitor.f_up(item))?; if let Some(VisitingState::VisitingChildren { - item, - non_processed_children, - .. - }) = stack.pop() + 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, @@ -1242,35 +1250,32 @@ impl TreeNode for Arc { } } -#[derive(Debug)] -enum ProcessingState { +/// 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), - // ← at this point, f_down is called + /// DFS over node's children ProcessingChildren { item: Transformed, non_processed_children: Vec, processed_children: Vec>, }, - // ← at this point, all children are processed + /// All children are processed (or jumped through). When executed, f_up may be called ProcessedAllChildren(Transformed), - // ← at this point, f_up is called } -#[derive(Debug)] +/// 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), - // ← at this point, f_down is called + /// 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, }, - // ← at this point, all children are visited - VisitedAllChildren { - item: &'a T, - tnr: TreeNodeRecursion, - }, - // ← at this point, f_up is called + /// 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 @@ -1409,11 +1414,8 @@ pub(crate) mod tests { .collect() } - /// Implement this train in order for your struct to be supported in parametrisation - pub(crate) trait TestTree - where - T: Sized, - { + /// [`node_tests`] uses methods when generating tests + pub(crate) trait TestTree { fn new_with_children(children: Vec, data: T) -> Self where Self: Sized; @@ -1427,7 +1429,7 @@ pub(crate) mod tests { fn with_children_from(data: T, other: Self) -> Self; } - macro_rules! gen_tests { + macro_rules! node_tests { ($TYPE:ident) => { fn visit_continue(_: &$TYPE) -> Result { Ok(TreeNodeRecursion::Continue) @@ -2546,7 +2548,7 @@ pub(crate) mod tests { } } - gen_tests!(TestTreeNode); + node_tests!(TestTreeNode); } pub mod test_dyn_tree_node { @@ -2607,7 +2609,7 @@ pub(crate) mod tests { type ArcTestNode = Arc>; - gen_tests!(ArcTestNode); + node_tests!(ArcTestNode); #[test] fn test_large_tree() { From b70a0f6e32394b322a70d1bf2391b12ee3fdd2fd Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 31 Oct 2024 15:23:58 +0000 Subject: [PATCH 08/13] Small performance optimisations for visit - reuse vec instead of making a new one - use mut ref to the last element instead of taking it on every iteration --- datafusion/common/src/tree_node.rs | 68 +++++++++++++++--------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index fc322d622da5..0ea207da77d3 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -1084,16 +1084,20 @@ impl TreeNode for Arc { stack.push(match node.tnr { TreeNodeRecursion::Continue => { + let non_processed_children: Vec<_> = node + .data + .arc_children() + .into_iter() + .cloned() + .rev() + .collect(); + TransformingState::ProcessingChildren { - non_processed_children: node - .data - .arc_children() - .into_iter() - .cloned() - .rev() - .collect(), + processed_children: Vec::with_capacity( + non_processed_children.len(), + ), + non_processed_children, item: node, - processed_children: vec![], } } TreeNodeRecursion::Jump => { @@ -1163,7 +1167,7 @@ impl TreeNode for Arc { processed_children, }) } else { - debug_assert_eq!(stack.len(), 0); + debug_assert!(stack.is_empty()); return Ok(node); } } @@ -1179,53 +1183,47 @@ impl TreeNode for Arc { ) -> Result { let mut stack = vec![VisitingState::NotStarted(self)]; - while let Some(item) = stack.pop() { - match item { + while let Some(node) = stack.last_mut() { + match node { VisitingState::NotStarted(item) => { let tnr = visitor.f_down(item)?; - stack.push(match tnr { - TreeNodeRecursion::Continue => VisitingState::VisitingChildren { - non_processed_children: item - .arc_children() - .into_iter() - .rev() - .collect(), - item, - tnr, - }, + *node = match tnr { + TreeNodeRecursion::Continue => { + let mut non_processed_children = item.arc_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, - mut non_processed_children, + ref mut non_processed_children, tnr, } => match tnr { TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => { if let Some(non_processed_item) = non_processed_children.pop() { - stack.extend([ - // Returning the node on the stack because there are more children to process - VisitingState::VisitingChildren { - item, - non_processed_children, - tnr, - }, - VisitingState::NotStarted(non_processed_item), - ]); + stack.push(VisitingState::NotStarted(non_processed_item)); } else { - stack.push(VisitingState::VisitedAllChildren { item, tnr }); + *node = VisitingState::VisitedAllChildren { item, tnr: *tnr } } } TreeNodeRecursion::Stop => { - return Ok(tnr); + return Ok(*tnr); } }, VisitingState::VisitedAllChildren { item, tnr } => { let tnr = tnr.visit_parent(|| visitor.f_up(item))?; + stack.pop(); if let Some(VisitingState::VisitingChildren { item, @@ -1239,7 +1237,7 @@ impl TreeNode for Arc { tnr, }); } else { - debug_assert_eq!(stack.len(), 0); + debug_assert!(stack.is_empty()); return Ok(tnr); } } From 7eb7174664ac8a02c662927d98b4f8b963af6f35 Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 1 Nov 2024 17:57:38 +0000 Subject: [PATCH 09/13] Do not take parent from the stack --- datafusion/common/src/tree_node.rs | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 0ea207da77d3..8fb4e7ab97d3 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -1151,21 +1151,14 @@ impl TreeNode for Arc { let node = node.transform_parent(|n| rewriter.f_up(n))?; if let Some(TransformingState::ProcessingChildren { - item: mut parent_node, - non_processed_children, - mut processed_children, + item: ref mut parent_node, + ref mut processed_children, .. - }) = stack.pop() + }) = stack.last_mut() { - // We need use returned recursion state when processing the remaining children + // We need to use the returned recursion state when processing the remaining children parent_node.tnr = node.tnr; processed_children.push(node); - - stack.push(TransformingState::ProcessingChildren { - item: parent_node, - non_processed_children, - processed_children, - }) } else { debug_assert!(stack.is_empty()); return Ok(node); From 810c899c6525b00f2113996d35ceec798d2e5d1b Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 1 Nov 2024 20:17:24 +0000 Subject: [PATCH 10/13] Remove extra Vec collect and update transformed in-place --- datafusion/common/src/tree_node.rs | 38 +++++++++--------------------- 1 file changed, 11 insertions(+), 27 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index 8fb4e7ab97d3..a8ffbbe7fbb7 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -999,19 +999,6 @@ impl< } } -/// Replaces node's children and recomputes the state -macro_rules! update_node_after_recursion { - ($NAME:ident, $CHILDREN:ident) => {{ - $NAME.transformed |= $CHILDREN.iter().any(|item| item.transformed); - - $NAME.data = $NAME - .data - .with_new_arc_children($CHILDREN.into_iter().map(|c| c.data).collect())?; - - $NAME - }}; -} - /// Blanket implementation for any `Arc` where `T` implements [`DynTreeNode`] /// (such as [`Arc`]). /// Unlike [`TreeNode`], performs node traversal iteratively rather than recursively to avoid stack overflow @@ -1129,22 +1116,18 @@ impl TreeNode for Arc { TransformingState::NotStarted(non_processed_item), ]); } else { - stack.push(TransformingState::ProcessedAllChildren( - update_node_after_recursion!(item, processed_children), - )) + item.data = + item.data.with_new_arc_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() - .map(Transformed::no), - ); - stack.push(TransformingState::ProcessedAllChildren( - update_node_after_recursion!(item, processed_children), - )); + processed_children + .extend(non_processed_children.into_iter().rev()); + item.data = + item.data.with_new_arc_children(processed_children)?; + stack.push(TransformingState::ProcessedAllChildren(item)); } }, TransformingState::ProcessedAllChildren(node) => { @@ -1158,7 +1141,8 @@ impl TreeNode for Arc { { // We need to use the returned recursion state when processing the remaining children parent_node.tnr = node.tnr; - processed_children.push(node); + parent_node.transformed |= node.transformed; + processed_children.push(node.data); } else { debug_assert!(stack.is_empty()); return Ok(node); @@ -1249,7 +1233,7 @@ enum TransformingState { ProcessingChildren { item: Transformed, non_processed_children: Vec, - processed_children: Vec>, + processed_children: Vec, }, /// All children are processed (or jumped through). When executed, f_up may be called ProcessedAllChildren(Transformed), From 0587e18a9ed0259e7bfa6e1d42401dcd8390b8a3 Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 1 Nov 2024 21:04:37 +0000 Subject: [PATCH 11/13] Implement iterations for `ConcreteTreeNode` --- datafusion/common/src/tree_node.rs | 193 +++++++++++++++++------------ 1 file changed, 115 insertions(+), 78 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index a8ffbbe7fbb7..f94a517d82b5 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -943,6 +943,24 @@ impl TransformedResult for Result> { } } +pub trait RecursiveNode: Sized { + fn children(&self) -> Vec<&Self>; + + /// Detaches children from the parent node (if possible) + fn take_children(self) -> (Self, Vec); + + 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) + } +} + /// 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`]. @@ -999,38 +1017,42 @@ impl< } } -/// Blanket implementation for any `Arc` where `T` implements [`DynTreeNode`] -/// (such as [`Arc`]). -/// Unlike [`TreeNode`], performs node traversal iteratively rather than recursively to avoid stack overflow -impl TreeNode for Arc { +/// Note that this implementation won't actually take children, but instead clone a reference +impl RecursiveNode for Arc { + fn children(&self) -> Vec<&Self> { + self.arc_children() + } + + fn take_children(self) -> (Self, Vec) { + let children = self.arc_children().into_iter().cloned().collect(); + (self, children) + } + + fn with_new_children(self, children: Vec) -> Result { + self.with_new_arc_children(children) + } +} + +impl TreeNode for T { fn apply_children<'n, F: FnMut(&'n Self) -> Result>( &'n self, f: F, ) -> Result { - self.arc_children().into_iter().apply_until_stop(f) + self.children().into_iter().apply_until_stop(f) } fn map_children Result>>( self, f: F, ) -> Result> { - let children = self.arc_children(); + let (new_self, children) = self.take_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| arc_self.with_new_arc_children(new_children)) - } else { - Ok(Transformed::new(self, false, new_children.tnr)) - } + let new_children = children.into_iter().map_until_stop_and_collect(f)?; + // Propagate up `new_children.transformed` and `new_children.tnr` along with + // the node containing transformed children. + new_children.map_data(|new_children| new_self.with_new_children(new_children)) } else { - Ok(Transformed::no(self)) + Ok(Transformed::no(new_self)) } } @@ -1071,13 +1093,8 @@ impl TreeNode for Arc { stack.push(match node.tnr { TreeNodeRecursion::Continue => { - let non_processed_children: Vec<_> = node - .data - .arc_children() - .into_iter() - .cloned() - .rev() - .collect(); + let (node, mut non_processed_children) = node.children(); + non_processed_children.reverse(); TransformingState::ProcessingChildren { processed_children: Vec::with_capacity( @@ -1117,7 +1134,7 @@ impl TreeNode for Arc { ]); } else { item.data = - item.data.with_new_arc_children(processed_children)?; + item.data.with_new_children(processed_children)?; stack.push(TransformingState::ProcessedAllChildren(item)) } } @@ -1125,8 +1142,7 @@ impl TreeNode for Arc { // 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_arc_children(processed_children)?; + item.data = item.data.with_new_children(processed_children)?; stack.push(TransformingState::ProcessedAllChildren(item)); } }, @@ -1166,7 +1182,7 @@ impl TreeNode for Arc { let tnr = visitor.f_down(item)?; *node = match tnr { TreeNodeRecursion::Continue => { - let mut non_processed_children = item.arc_children(); + let mut non_processed_children = item.children(); non_processed_children.reverse(); VisitingState::VisitingChildren { @@ -1252,7 +1268,6 @@ enum VisitingState<'a, T> { /// 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. @@ -1267,27 +1282,17 @@ pub trait ConcreteTreeNode: Sized { fn with_new_children(self, children: Vec) -> Result; } -impl TreeNode for T { - fn apply_children<'n, F: FnMut(&'n Self) -> Result>( - &'n self, - f: F, - ) -> Result { - self.children().iter().apply_until_stop(f) +impl RecursiveNode for T { + fn children(&self) -> Vec<&Self> { + self.children().iter().collect() } - fn map_children Result>>( - self, - f: F, - ) -> Result> { - let (new_self, children) = self.take_children(); - if !children.is_empty() { - let new_children = children.into_iter().map_until_stop_and_collect(f)?; - // Propagate up `new_children.transformed` and `new_children.tnr` along with - // the node containing transformed children. - new_children.map_data(|new_children| new_self.with_new_children(new_children)) - } else { - Ok(Transformed::no(new_self)) - } + fn take_children(self) -> (Self, Vec) { + self.take_children() + } + + fn with_new_children(self, children: Vec) -> Result { + self.with_new_children(children) } } @@ -2466,39 +2471,45 @@ pub(crate) mod tests { }; } - pub mod test_tree_node { - use super::*; + macro_rules! test_node { + ($NAME: ident) => { + #[derive(Debug, Eq, Hash, PartialEq, Clone)] + pub struct $NAME { + pub(crate) children: Vec, + pub(crate) data: T, + } - #[derive(Debug, Eq, Hash, PartialEq, Clone)] - pub struct TestTreeNode { - pub(crate) children: Vec>, - pub(crate) data: T, - } + impl TestTree for $NAME { + fn new_with_children(children: Vec, data: T) -> Self { + Self { children, data } + } - impl TestTree for TestTreeNode { - fn new_with_children(children: Vec>, data: T) -> Self { - Self { children, data } - } + fn new_leaf(data: T) -> Self { + Self { + children: vec![], + data, + } + } - fn new_leaf(data: T) -> Self { - Self { - children: vec![], - data, + fn is_leaf(&self) -> bool { + self.children.is_empty() } - } - fn is_leaf(&self) -> bool { - self.children.is_empty() - } + fn eq_data(&self, other: &T) -> bool { + self.data == *other + } - 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) + } } + }; + } - 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>( @@ -2526,7 +2537,7 @@ pub(crate) mod tests { node_tests!(TestTreeNode); } - pub mod test_dyn_tree_node { + mod test_dyn_tree_node { use super::*; #[derive(Debug, Eq, Hash, PartialEq, Clone)] @@ -2587,7 +2598,7 @@ pub(crate) mod tests { node_tests!(ArcTestNode); #[test] - fn test_large_tree() { + fn test_large_visit() { let mut item = ArcTestNode::new_leaf("initial".to_string()); for i in 0..3000 { item = @@ -2600,4 +2611,30 @@ pub(crate) mod tests { item.visit(&mut visitor).unwrap(); } } + + mod test_concrete_tree_node { + use super::*; + use crate::tree_node::ConcreteTreeNode; + use std::mem; + + test_node!(ConcreteTestTreeNode); + + impl ConcreteTreeNode for ConcreteTestTreeNode { + fn children(&self) -> &[Self] { + self.children.as_slice() + } + + 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) + } + } + + node_tests!(ConcreteTestTreeNode); + } } From ee4f6636a3b95d0d0281c25ed46da17221aef3e5 Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 1 Nov 2024 23:56:42 +0000 Subject: [PATCH 12/13] More verbose notes --- datafusion/common/src/tree_node.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index f94a517d82b5..eb90aed0242f 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -943,12 +943,16 @@ impl TransformedResult for Result> { } } +/// A node trait which makes all tree traversal iterative and not recursive. pub trait RecursiveNode: Sized { + /// Read-only access to children fn children(&self) -> Vec<&Self>; - /// Detaches children from the parent node (if possible) + /// Detaches children from the parent node (if possible). + /// Unlike [`ConcreteTreeNode`] it doesn't possible that the value will actually be removed from the parent fn take_children(self) -> (Self, Vec); + /// Replaces children with the given one fn with_new_children(self, children: Vec) -> Result; } @@ -1106,7 +1110,7 @@ impl TreeNode for T { } TreeNodeRecursion::Jump => { TransformingState::ProcessedAllChildren( - // No need to process children, we can just this stage + // No need to process children, we can just skip that stage node.with_tnr(TreeNodeRecursion::Continue), ) } @@ -1123,7 +1127,7 @@ impl TreeNode for T { 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 in the stack + // This node still has children, so put it back on the stack TransformingState::ProcessingChildren { item, non_processed_children, @@ -1155,7 +1159,7 @@ impl TreeNode for T { .. }) = stack.last_mut() { - // We need to use the returned recursion state when processing the remaining children + // 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); From 403230aafcad9f75f806b56a2d5c0848b6c3a4c5 Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 8 Nov 2024 16:55:08 +0000 Subject: [PATCH 13/13] Remove `RecursiveNode` --- datafusion/common/src/tree_node.rs | 158 ++++-------------- .../physical-expr-common/src/tree_node.rs | 4 +- datafusion/physical-plan/src/tree_node.rs | 4 +- 3 files changed, 32 insertions(+), 134 deletions(-) diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs index a876a24964f6..3d2169a597ab 100644 --- a/datafusion/common/src/tree_node.rs +++ b/datafusion/common/src/tree_node.rs @@ -925,31 +925,9 @@ impl TransformedResult for Result> { } } -/// A node trait which makes all tree traversal iterative and not recursive. -pub trait RecursiveNode: Sized { - /// Read-only access to children - fn children(&self) -> Vec<&Self>; - - /// Detaches children from the parent node (if possible). - /// Unlike [`ConcreteTreeNode`] it doesn't possible that the value will actually be removed from the parent - fn take_children(self) -> (Self, Vec); - - /// Replaces children with the given one - 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) - } -} - /// 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>; @@ -961,6 +939,23 @@ pub trait DynTreeNode { ) -> Result>; } +/// 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 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>, @@ -1003,23 +998,16 @@ impl< } } -/// Note that this implementation won't actually take children, but instead clone a reference -impl RecursiveNode for Arc { - fn children(&self) -> Vec<&Self> { - self.arc_children() - } +impl Transformed { + fn children(mut self) -> (Self, Vec) { + let (data, children) = self.data.take_children(); + self.data = data; - fn take_children(self) -> (Self, Vec) { - let children = self.arc_children().into_iter().cloned().collect(); (self, children) } - - fn with_new_children(self, children: Vec) -> Result { - self.with_new_arc_children(children) - } } -impl TreeNode for T { +impl TreeNode for T { fn apply_children<'n, F: FnMut(&'n Self) -> Result>( &'n self, f: F, @@ -1259,7 +1247,7 @@ enum VisitingState<'a, T> { /// 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 children(&self) -> Vec<&Self>; /// Detaches the node from its children, returning the node itself and its detached children. fn take_children(self) -> (Self, Vec); @@ -1268,30 +1256,15 @@ pub trait ConcreteTreeNode: Sized { fn with_new_children(self, children: Vec) -> Result; } -impl RecursiveNode for T { - fn children(&self) -> Vec<&Self> { - self.children().iter().collect() - } - - fn take_children(self) -> (Self, Vec) { - self.take_children() - } - - fn with_new_children(self, children: Vec) -> Result { - self.with_new_children(children) - } -} - #[cfg(test)] pub(crate) mod tests { use crate::tree_node::{ - DynTreeNode, Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, - TreeNodeRewriter, TreeNodeVisitor, + Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, TreeNodeRewriter, + TreeNodeVisitor, }; use crate::Result; use std::collections::HashMap; use std::fmt::Display; - use std::sync::Arc; macro_rules! visit_test { ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_VISITS:expr) => { @@ -2523,81 +2496,6 @@ pub(crate) mod tests { node_tests!(TestTreeNode); } - mod test_dyn_tree_node { - use super::*; - - #[derive(Debug, Eq, Hash, PartialEq, Clone)] - pub struct DynTestNode { - pub(crate) data: Arc, - pub(crate) children: Vec>>, - } - - impl TestTree for Arc> { - fn new_with_children(children: Vec, data: T) -> Self - where - Self: Sized, - { - Arc::new(DynTestNode { - data: Arc::new(data), - children, - }) - } - - fn new_leaf(data: T) -> Self { - Arc::new(DynTestNode { - data: Arc::new(data), - children: vec![], - }) - } - - fn is_leaf(&self) -> bool { - self.children.is_empty() - } - - fn eq_data(&self, other: &T) -> bool { - self.data.as_ref().eq(other) - } - - fn with_children_from(data: T, other: Self) -> Self { - Self::new_with_children(other.children.clone(), data) - } - } - - impl DynTreeNode for DynTestNode { - fn arc_children(&self) -> Vec<&Arc> { - self.children.iter().collect() - } - - fn with_new_arc_children( - self: Arc, - new_children: Vec>, - ) -> Result> { - Ok(Arc::new(Self { - children: new_children, - data: Arc::clone(&self.data), - })) - } - } - - type ArcTestNode = Arc>; - - node_tests!(ArcTestNode); - - #[test] - fn test_large_visit() { - let mut item = ArcTestNode::new_leaf("initial".to_string()); - for i in 0..3000 { - item = - ArcTestNode::new_with_children(vec![item], format!("parent-{}", i)); - } - - let mut visitor = - TestVisitor::new(Box::new(visit_continue), Box::new(visit_continue)); - - item.visit(&mut visitor).unwrap(); - } - } - mod test_concrete_tree_node { use super::*; use crate::tree_node::ConcreteTreeNode; @@ -2606,8 +2504,8 @@ pub(crate) mod tests { test_node!(ConcreteTestTreeNode); impl ConcreteTreeNode for ConcreteTestTreeNode { - fn children(&self) -> &[Self] { - self.children.as_slice() + fn children(&self) -> Vec<&Self> { + self.children.iter().collect() } fn take_children(mut self) -> (Self, Vec) { diff --git a/datafusion/physical-expr-common/src/tree_node.rs b/datafusion/physical-expr-common/src/tree_node.rs index d03ac6da2c7b..790e805230e1 100644 --- a/datafusion/physical-expr-common/src/tree_node.rs +++ b/datafusion/physical-expr-common/src/tree_node.rs @@ -88,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 96a586209f15..14aa0961889a 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -90,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) {