use crate::arena::{ExprArena, ExprId, ExprNode};
pub fn simplify(arena: &mut ExprArena, id: ExprId) -> ExprId {
let folded = match arena.get(id).clone() {
ExprNode::Neg(inner) => match arena.get(inner) {
ExprNode::Const(c) => Some(ExprNode::Const(-*c)),
_ => None,
},
ExprNode::Pow(base, exp) => match (arena.get(base), arena.get(exp)) {
(ExprNode::Const(b), ExprNode::Const(e)) => Some(ExprNode::Const(b.powf(*e))),
_ => None,
},
ExprNode::Sin(inner)
| ExprNode::Cos(inner)
| ExprNode::Exp(inner)
| ExprNode::Log(inner) => {
let node = arena.get(id).clone();
if let ExprNode::Const(c) = arena.get(inner) {
Some(ExprNode::Const(match node {
ExprNode::Sin(_) => c.sin(),
ExprNode::Cos(_) => c.cos(),
ExprNode::Exp(_) => c.exp(),
ExprNode::Log(_) => c.ln(),
_ => unreachable!(),
}))
} else {
None
}
}
_ => None,
};
match folded {
Some(node) => arena.push(node),
None => id,
}
}