oximo_expr/
visit.rs

1use crate::arena::{ExprArena, ExprId, ExprNode};
2
3/// Pre-order visitor over an arena. Backends implement this to translate the
4/// expression tree into solver-specific representations without copying.
5pub trait Visitor {
6    fn visit(&mut self, arena: &ExprArena, id: ExprId, node: &ExprNode);
7}
8
9/// Walk the subtree rooted at `id` in pre-order.
10///
11/// Implemented with an explicit stack so arbitrarily deep expressions are
12/// safe.
13pub fn walk<V: Visitor>(arena: &ExprArena, id: ExprId, visitor: &mut V) {
14    let mut stack = vec![id];
15    while let Some(cur) = stack.pop() {
16        let node = arena.get(cur);
17        visitor.visit(arena, cur, node);
18        match node {
19            ExprNode::Add(children)
20            | ExprNode::Mul(children)
21            | ExprNode::Min(children)
22            | ExprNode::Max(children) => {
23                for &child in children.iter().rev() {
24                    stack.push(child);
25                }
26            }
27            ExprNode::Unary(_, inner) => {
28                stack.push(*inner);
29            }
30            ExprNode::Pow(base, exp) | ExprNode::Atan2(base, exp) => {
31                stack.push(*exp);
32                stack.push(*base);
33            }
34            ExprNode::Div(num, den) => {
35                stack.push(*den);
36                stack.push(*num);
37            }
38            ExprNode::Const(_)
39            | ExprNode::Var(_)
40            | ExprNode::Param(_)
41            | ExprNode::Linear { .. } => {}
42        }
43    }
44}
45
46/// Walk the distinct nodes reachable from `id` in pre-order, visiting each
47/// shared subexpression exactly once.
48pub fn walk_shared<V: Visitor>(arena: &ExprArena, id: ExprId, visitor: &mut V) {
49    let mut visited = rustc_hash::FxHashSet::default();
50    let mut stack = vec![id];
51    while let Some(cur) = stack.pop() {
52        if !visited.insert(cur) {
53            continue;
54        }
55        let node = arena.get(cur);
56        visitor.visit(arena, cur, node);
57        match node {
58            ExprNode::Add(children)
59            | ExprNode::Mul(children)
60            | ExprNode::Min(children)
61            | ExprNode::Max(children) => {
62                for &child in children.iter().rev() {
63                    stack.push(child);
64                }
65            }
66            ExprNode::Unary(_, inner) => {
67                stack.push(*inner);
68            }
69            ExprNode::Pow(base, exp) | ExprNode::Atan2(base, exp) => {
70                stack.push(*exp);
71                stack.push(*base);
72            }
73            ExprNode::Div(num, den) => {
74                stack.push(*den);
75                stack.push(*num);
76            }
77            ExprNode::Const(_)
78            | ExprNode::Var(_)
79            | ExprNode::Param(_)
80            | ExprNode::Linear { .. } => {}
81        }
82    }
83}
oximo_expr/visit.rs

oximo_expr/
visit.rs
