polars_plan/plans/visitor/
visitors.rs

1use recursive::recursive;
2
3use super::*;
4
5/// An implementor of this trait decides how and in which order its nodes get traversed
6/// Implemented for [`crate::dsl::Expr`] and [`AexprNode`].
7pub trait TreeWalker: Sized {
8    type Arena;
9    fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
10        &self,
11        op: &mut F,
12        arena: &Self::Arena,
13    ) -> PolarsResult<VisitRecursion>;
14
15    fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
16        self,
17        op: &mut F,
18        arena: &mut Self::Arena,
19    ) -> PolarsResult<Self>;
20
21    /// Walks all nodes in depth-first-order.
22    #[recursive]
23    fn visit<V: Visitor<Node = Self, Arena = Self::Arena>>(
24        &self,
25        visitor: &mut V,
26        arena: &Self::Arena,
27    ) -> PolarsResult<VisitRecursion> {
28        match visitor.pre_visit(self, arena)? {
29            VisitRecursion::Continue => {},
30            // If the recursion should skip, do not apply to its children. And let the recursion continue
31            VisitRecursion::Skip => return Ok(VisitRecursion::Continue),
32            // If the recursion should stop, do not apply to its children
33            VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
34        };
35
36        match self.apply_children(&mut |node, arena| node.visit(visitor, arena), arena)? {
37            // let the recursion continue
38            VisitRecursion::Continue | VisitRecursion::Skip => {},
39            // If the recursion should stop, no further post visit will be performed
40            VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
41        }
42
43        visitor.post_visit(self, arena)
44    }
45
46    #[recursive]
47    fn rewrite<R: RewritingVisitor<Node = Self, Arena = Self::Arena>>(
48        self,
49        rewriter: &mut R,
50        arena: &mut Self::Arena,
51    ) -> PolarsResult<Self> {
52        let mutate_this_node = match rewriter.pre_visit(&self, arena)? {
53            RewriteRecursion::MutateAndStop => return rewriter.mutate(self, arena),
54            RewriteRecursion::Stop => return Ok(self),
55            RewriteRecursion::MutateAndContinue => true,
56            RewriteRecursion::NoMutateAndContinue => false,
57        };
58
59        let after_applied_children =
60            self.map_children(&mut |node, arena| node.rewrite(rewriter, arena), arena)?;
61
62        if mutate_this_node {
63            rewriter.mutate(after_applied_children, arena)
64        } else {
65            Ok(after_applied_children)
66        }
67    }
68}
69
70pub trait Visitor {
71    type Node;
72    type Arena;
73
74    /// Invoked before any children of `node` are visited.
75    fn pre_visit(
76        &mut self,
77        _node: &Self::Node,
78        _arena: &Self::Arena,
79    ) -> PolarsResult<VisitRecursion> {
80        Ok(VisitRecursion::Continue)
81    }
82
83    /// Invoked after all children of `node` are visited. Default
84    /// implementation does nothing.
85    fn post_visit(
86        &mut self,
87        _node: &Self::Node,
88        _arena: &Self::Arena,
89    ) -> PolarsResult<VisitRecursion> {
90        Ok(VisitRecursion::Continue)
91    }
92}
93
94pub trait RewritingVisitor {
95    type Node;
96    type Arena;
97
98    /// Invoked before any children of `node` are visited.
99    fn pre_visit(
100        &mut self,
101        _node: &Self::Node,
102        _arena: &mut Self::Arena,
103    ) -> PolarsResult<RewriteRecursion> {
104        Ok(RewriteRecursion::MutateAndContinue)
105    }
106
107    fn mutate(&mut self, node: Self::Node, _arena: &mut Self::Arena) -> PolarsResult<Self::Node>;
108}
polars_plan/plans/visitor/visitors.rs

polars_plan/plans/visitor/
visitors.rs
