add reference visitor APIs

peter-toth · May 11, 2024 · e844799 · e844799
1 parent 2a15614
commit e844799
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 8 deletions.
diff --git a/datafusion/common/src/tree_node.rs b/datafusion/common/src/tree_node.rs
@@ -133,6 +133,16 @@ pub trait TreeNode: Sized {
             .visit_parent(|| visitor.f_up(self))
     }
 
+    fn visit_ref<'n, V: TreeNodeRefVisitor<'n, Node = Self>>(
+        &'n self,
+        visitor: &mut V,
+    ) -> Result<TreeNodeRecursion> {
+        visitor
+            .f_down(self)?
+            .visit_children(|| self.apply_children_ref(|c| c.visit_ref(visitor)))?
+            .visit_parent(|| visitor.f_up(self))
+    }
+
     /// Rewrite the tree node with a [`TreeNodeRewriter`], performing a
     /// depth-first walk of the node and its children.
     ///
@@ -204,6 +214,24 @@ pub trait TreeNode: Sized {
         apply_impl(self, &mut f)
     }
 
+    fn apply_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
+        mut f: F,
+    ) -> Result<TreeNodeRecursion> {
+        fn apply_ref_impl<
+            'n,
+            N: TreeNode,
+            F: FnMut(&'n N) -> Result<TreeNodeRecursion>,
+        >(
+            node: &'n N,
+            f: &mut F,
+        ) -> Result<TreeNodeRecursion> {
+            f(node)?.visit_children(|| node.apply_children_ref(|c| apply_ref_impl(c, f)))
+        }
+
+        apply_ref_impl(self, &mut f)
+    }
+
     /// Recursively rewrite the node's children and then the node using `f`
     /// (a bottom-up post-order traversal).
     ///
@@ -430,6 +458,13 @@ pub trait TreeNode: Sized {
     fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
         &self,
         f: F,
+    ) -> Result<TreeNodeRecursion> {
+        self.apply_children_ref(f)
+    }
+
+    fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
+        f: F,
     ) -> Result<TreeNodeRecursion>;
 
     /// Low-level API used to implement other APIs.
@@ -483,6 +518,23 @@ pub trait TreeNodeVisitor: Sized {
     }
 }
 
+pub trait TreeNodeRefVisitor<'n>: Sized {
+    /// The node type which is visitable.
+    type Node: TreeNode;
+
+    /// Invoked while traversing down the tree, before any children are visited.
+    /// Default implementation continues the recursion.
+    fn f_down(&mut self, _node: &'n Self::Node) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
+
+    /// Invoked while traversing up the tree after children are visited. Default
+    /// implementation continues the recursion.
+    fn f_up(&mut self, _node: &'n Self::Node) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
+}
+
 /// A [Visitor](https://en.wikipedia.org/wiki/Visitor_pattern) for recursively
 /// rewriting [`TreeNode`]s via [`TreeNode::rewrite`].
 ///
@@ -857,6 +909,10 @@ pub trait DynTreeNode {
     /// Returns all children of the specified `TreeNode`.
     fn arc_children(&self) -> Vec<Arc<Self>>;
 
+    fn children(&self) -> Vec<&Arc<Self>> {
+        panic!("DynTreeNode::children is not implemented yet")
+    }
+
     /// Constructs a new node with the specified children.
     fn with_new_arc_children(
         &self,
@@ -875,6 +931,13 @@ impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
         self.arc_children().iter().apply_until_stop(f)
     }
 
+    fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
+        f: F,
+    ) -> Result<TreeNodeRecursion> {
+        self.children().into_iter().apply_until_stop(f)
+    }
+
     fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
         self,
         f: F,
@@ -913,8 +976,8 @@ pub trait ConcreteTreeNode: Sized {
 }
 
 impl<T: ConcreteTreeNode> TreeNode for T {
-    fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
-        &self,
+    fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
         f: F,
     ) -> Result<TreeNodeRecursion> {
         self.children().into_iter().apply_until_stop(f)
@@ -959,8 +1022,8 @@ mod tests {
     }
 
     impl<T> TreeNode for TestTreeNode<T> {
-        fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
-            &self,
+        fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+            &'n self,
             f: F,
         ) -> Result<TreeNodeRecursion> {
             self.children.iter().apply_until_stop(f)

diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs
@@ -56,8 +56,8 @@ use datafusion_common::{
 };
 
 impl TreeNode for LogicalPlan {
-    fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
-        &self,
+    fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
         f: F,
     ) -> Result<TreeNodeRecursion> {
         self.inputs().into_iter().apply_until_stop(f)

diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs
@@ -30,8 +30,8 @@ use datafusion_common::tree_node::{
 use datafusion_common::{internal_err, map_until_stop_and_collect, Result};
 
 impl TreeNode for Expr {
-    fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
-        &self,
+    fn apply_children_ref<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
+        &'n self,
         f: F,
     ) -> Result<TreeNodeRecursion> {
         let children = match self {