oximo_expr/
simplify.rs

1use crate::arena::{ExprArena, ExprId, ExprNode};
2
3/// Apply local algebraic simplifications to the subtree rooted at `id`,
4/// returning a (possibly fresh) `ExprId` that is observationally equivalent.
5///
6/// Current rules: constant folding for unary nodes, `Pow`, `Div`, `Atan2`, and
7/// all-constant `Min`/`Max` nodes. The linear fast-path is already canonical,
8/// so we leave `Linear` and n-ary `Add`/`Mul` alone.
9///
10/// TODO: Extend this once we add a CSE pass.
11pub fn simplify(arena: &mut ExprArena, id: ExprId) -> ExprId {
12    let folded = match arena.get(id).clone() {
13        ExprNode::Unary(op, inner) => match arena.get(inner) {
14            ExprNode::Const(c) => Some(ExprNode::Const(op.apply(*c))),
15            _ => None,
16        },
17        ExprNode::Pow(base, exp) => match (arena.get(base), arena.get(exp)) {
18            (ExprNode::Const(b), ExprNode::Const(e)) => Some(ExprNode::Const(b.powf(*e))),
19            _ => None,
20        },
21        ExprNode::Div(num, den) => match (arena.get(num), arena.get(den)) {
22            (ExprNode::Const(n), ExprNode::Const(d)) => Some(ExprNode::Const(n / d)),
23            _ => None,
24        },
25        ExprNode::Atan2(y, x) => match (arena.get(y), arena.get(x)) {
26            (ExprNode::Const(y), ExprNode::Const(x)) => Some(ExprNode::Const(y.atan2(*x))),
27            _ => None,
28        },
29        ExprNode::Min(children) | ExprNode::Max(children) => {
30            let is_min = matches!(arena.get(id), ExprNode::Min(_));
31            let mut values = children.iter().map(|child| match arena.get(*child) {
32                ExprNode::Const(value) => Some(*value),
33                _ => None,
34            });
35            values.next().flatten().and_then(|first| {
36                values
37                    .try_fold(first, |acc, value| {
38                        value.map(|v| if is_min { acc.min(v) } else { acc.max(v) })
39                    })
40                    .map(ExprNode::Const)
41            })
42        }
43        _ => None,
44    };
45    match folded {
46        Some(node) => arena.push(node),
47        None => id,
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    use crate::arena::{ExprArena, ExprNode, UnaryOp};
55
56    #[test]
57    fn folds_abs_of_const() {
58        let mut a = ExprArena::new();
59        let c = a.push(ExprNode::Const(-5.0));
60        let abs = a.push(ExprNode::Unary(UnaryOp::Abs, c));
61        let folded = simplify(&mut a, abs);
62        assert!(matches!(a.get(folded), ExprNode::Const(v) if (*v - 5.0).abs() < f64::EPSILON));
63    }
64}
oximo_expr/simplify.rs

oximo_expr/
simplify.rs
