powdr_expression/
visitors.rs

1use std::{iter, ops::ControlFlow};
2
3use crate::AlgebraicExpression;
4
5/// Generic trait that allows to iterate over sub-structures.
6///
7/// It is only meant to iterate non-recursively over the direct children.
8/// Self and O do not have to be the same type and we can also have
9/// Children<O1> and Children<O2> implemented for the same type,
10/// if the goal is to iterate over sub-structures of different kinds.
11pub trait Children<O> {
12    /// Returns an iterator over all direct children of kind O in this object.
13    fn children(&self) -> Box<dyn Iterator<Item = &O> + '_>;
14    /// Returns an iterator over all direct children of kind Q in this object.
15    fn children_mut(&mut self) -> Box<dyn Iterator<Item = &mut O> + '_>;
16}
17
18pub trait AllChildren<O> {
19    /// Returns an iterator over all direct and indirect children of kind `O` in this object.
20    /// If `O` and `Self` are the same type, also includes `self`.
21    /// Pre-order visitor.
22    fn all_children(&self) -> Box<dyn Iterator<Item = &O> + '_>;
23}
24
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum VisitOrder {
27    Pre,
28    Post,
29}
30
31/// A trait to be implemented by an AST node.
32///
33/// The idea is that it calls a callback function on each of the sub-nodes
34/// that are expressions.
35/// The difference to the Children<Expr> trait is that ExpressionVisitable
36/// visits recursively.
37/// If a node implements Children<Expr>, it also implements ExpressionVisitable<Expr>.
38pub trait ExpressionVisitable<Expr> {
39    /// Traverses the AST and calls `f` on each Expression in pre-order.
40    fn pre_visit_expressions_mut<F>(&mut self, f: &mut F)
41    where
42        F: FnMut(&mut Expr),
43    {
44        let _ = self.visit_expressions_mut(
45            &mut move |e| {
46                f(e);
47                ControlFlow::Continue::<()>(())
48            },
49            VisitOrder::Pre,
50        );
51    }
52
53    /// Traverses the AST and calls `f` on each Expression in post-order.
54    fn post_visit_expressions_mut<F>(&mut self, f: &mut F)
55    where
56        F: FnMut(&mut Expr),
57    {
58        let _ = self.visit_expressions_mut(
59            &mut move |e| {
60                f(e);
61                ControlFlow::Continue::<()>(())
62            },
63            VisitOrder::Post,
64        );
65    }
66
67    fn visit_expressions<F, B>(&self, f: &mut F, order: VisitOrder) -> ControlFlow<B>
68    where
69        F: FnMut(&Expr) -> ControlFlow<B>;
70
71    fn visit_expressions_mut<F, B>(&mut self, f: &mut F, order: VisitOrder) -> ControlFlow<B>
72    where
73        F: FnMut(&mut Expr) -> ControlFlow<B>;
74}
75
76impl<Expr: ExpressionVisitable<Expr>, C: Children<Expr>> ExpressionVisitable<Expr> for C {
77    fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
78    where
79        F: FnMut(&mut Expr) -> ControlFlow<B>,
80    {
81        self.children_mut()
82            .try_for_each(|child| child.visit_expressions_mut(f, o))
83    }
84
85    fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
86    where
87        F: FnMut(&Expr) -> ControlFlow<B>,
88    {
89        self.children()
90            .try_for_each(|child| child.visit_expressions(f, o))
91    }
92}
93
94impl<Expr: AllChildren<Expr>, C: Children<Expr>> AllChildren<Expr> for C {
95    fn all_children(&self) -> Box<dyn Iterator<Item = &Expr> + '_> {
96        Box::new(self.children().flat_map(|e| e.all_children()))
97    }
98}
99
100impl<T, R> ExpressionVisitable<AlgebraicExpression<T, R>> for AlgebraicExpression<T, R> {
101    fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
102    where
103        F: FnMut(&mut AlgebraicExpression<T, R>) -> ControlFlow<B>,
104    {
105        if o == VisitOrder::Pre {
106            f(self)?;
107        }
108        self.children_mut()
109            .try_for_each(|e| e.visit_expressions_mut(f, o))?;
110        if o == VisitOrder::Post {
111            f(self)?;
112        }
113        ControlFlow::Continue(())
114    }
115
116    fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
117    where
118        F: FnMut(&AlgebraicExpression<T, R>) -> ControlFlow<B>,
119    {
120        if o == VisitOrder::Pre {
121            f(self)?;
122        }
123        self.children()
124            .try_for_each(|e| e.visit_expressions(f, o))?;
125        if o == VisitOrder::Post {
126            f(self)?;
127        }
128        ControlFlow::Continue(())
129    }
130}
131
132impl<T, R> AllChildren<AlgebraicExpression<T, R>> for AlgebraicExpression<T, R> {
133    fn all_children(&self) -> Box<dyn Iterator<Item = &AlgebraicExpression<T, R>> + '_> {
134        Box::new(iter::once(self).chain(self.children().flat_map(|e| e.all_children())))
135    }
136}
powdr_expression/visitors.rs

powdr_expression/
visitors.rs
