Skip to main content

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 and `Pow`. The linear
7/// fast-path is already canonical, so we leave `Linear` and n-ary `Add`/`Mul`
8/// 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::Neg(inner) => match arena.get(inner) {
14            ExprNode::Const(c) => Some(ExprNode::Const(-*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::Sin(inner)
22        | ExprNode::Cos(inner)
23        | ExprNode::Exp(inner)
24        | ExprNode::Log(inner) => {
25            let node = arena.get(id).clone();
26            if let ExprNode::Const(c) = arena.get(inner) {
27                Some(ExprNode::Const(match node {
28                    ExprNode::Sin(_) => c.sin(),
29                    ExprNode::Cos(_) => c.cos(),
30                    ExprNode::Exp(_) => c.exp(),
31                    ExprNode::Log(_) => c.ln(),
32                    _ => unreachable!(),
33                }))
34            } else {
35                None
36            }
37        }
38        _ => None,
39    };
40    match folded {
41        Some(node) => arena.push(node),
42        None => id,
43    }
44}