alamb commented on code in PR #8891:
URL: https://github.com/apache/arrow-datafusion/pull/8891#discussion_r1460885187


##########
datafusion/common/src/tree_node.rs:
##########
@@ -248,69 +275,58 @@ pub trait TreeNode: Sized + Clone {
 /// If an [`Err`] result is returned, recursion is stopped
 /// immediately.
 ///
-/// If [`VisitRecursion::Stop`] is returned on a call to pre_visit, no
+/// If [`TreeNodeRecursion::Stop`] is returned on a call to pre_visit, no
 /// children of that tree node are visited, nor is post_visit
 /// called on that tree node
 ///
-/// If [`VisitRecursion::Stop`] is returned on a call to post_visit, no
+/// If [`TreeNodeRecursion::Stop`] is returned on a call to post_visit, no
 /// siblings of that tree node are visited, nor is post_visit
 /// called on its parent tree node
 ///
-/// If [`VisitRecursion::Skip`] is returned on a call to pre_visit, no
+/// If [`TreeNodeRecursion::Skip`] is returned on a call to pre_visit, no
 /// children of that tree node are visited.
 pub trait TreeNodeVisitor: Sized {
     /// The node type which is visitable.
     type N: TreeNode;
 
     /// Invoked before any children of `node` are visited.
-    fn pre_visit(&mut self, node: &Self::N) -> Result<VisitRecursion>;
+    fn pre_visit(&mut self, node: &Self::N) -> Result<TreeNodeRecursion>;
 
     /// Invoked after all children of `node` are visited. Default
     /// implementation does nothing.
-    fn post_visit(&mut self, _node: &Self::N) -> Result<VisitRecursion> {
-        Ok(VisitRecursion::Continue)
+    fn post_visit(&mut self, _node: &Self::N) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
     }
 }
 
-/// Trait for potentially recursively transform an [`TreeNode`] node
-/// tree. When passed to `TreeNode::rewrite`, `TreeNodeRewriter::mutate` is
-/// invoked recursively on all nodes of a tree.
+/// Trait for potentially recursively transform a [`TreeNode`] node tree.
 pub trait TreeNodeRewriter: Sized {
     /// The node type which is rewritable.
-    type N: TreeNode;
+    type Node: TreeNode;
 
-    /// Invoked before (Preorder) any children of `node` are rewritten /
-    /// visited. Default implementation returns `Ok(Recursion::Continue)`
-    fn pre_visit(&mut self, _node: &Self::N) -> Result<RewriteRecursion> {
-        Ok(RewriteRecursion::Continue)
+    /// Invoked while traversing down the tree before any children are 
rewritten /
+    /// visited.
+    /// Default implementation returns the node unmodified and continues 
recursion.
+    fn f_down(&mut self, node: Self::Node) -> Result<(Self::Node, 
TreeNodeRecursion)> {
+        Ok((node, TreeNodeRecursion::Continue))
     }
 
-    /// Invoked after (Postorder) all children of `node` have been mutated and
-    /// returns a potentially modified node.
-    fn mutate(&mut self, node: Self::N) -> Result<Self::N>;
-}
-
-/// Controls how the [`TreeNode`] recursion should proceed for 
[`TreeNode::rewrite`].
-#[derive(Debug)]
-pub enum RewriteRecursion {
-    /// Continue rewrite this node tree.
-    Continue,
-    /// Call 'op' immediately and return.
-    Mutate,
-    /// Do not rewrite the children of this node.
-    Stop,
-    /// Keep recursive but skip apply op on this node
-    Skip,
+    /// Invoked while traversing up the tree after all children have been 
rewritten /
+    /// visited.
+    /// Default implementation returns the node unmodified.
+    fn f_up(&mut self, node: Self::Node) -> Result<Self::Node> {
+        Ok(node)
+    }
 }
 
-/// Controls how the [`TreeNode`] recursion should proceed for 
[`TreeNode::visit`].
+/// Controls how [`TreeNode`] recursions should proceed.
 #[derive(Debug)]
-pub enum VisitRecursion {
-    /// Continue the visit to this node tree.
+pub enum TreeNodeRecursion {

Review Comment:
   Combining these enums seems like an improvement to me



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to