Skip to main content

ocas_eval/
tree.rs

1//! Owned intermediate representation for expression tree compilation.
2//!
3//! [`EvalTree`] is an arena-free, owned expression tree that can be
4//! constructed from an [`Atom`](ocas_atom::Atom) and then optimized
5//! before instruction generation. It decouples the compilation pipeline
6//! from the arena lifetime.
7
8use ocas_atom::{Atom, AtomNode};
9
10/// An owned intermediate representation of a symbolic expression.
11///
12/// Unlike [`Atom`](ocas_atom::Atom), `EvalTree` owns all its data and
13/// does not depend on an arena. This makes it suitable for multi-pass
14/// compilation and optimization.
15#[derive(Debug, Clone, PartialEq)]
16pub enum EvalTree {
17    /// A numeric constant.
18    Num(f64),
19    /// A variable reference.
20    Var(String),
21    /// A named function applied to arguments.
22    Fun(String, Vec<EvalTree>),
23    /// A sum of terms.
24    Add(Vec<EvalTree>),
25    /// A product of factors.
26    Mul(Vec<EvalTree>),
27    /// A power expression: base^exponent.
28    Pow(Box<EvalTree>, Box<EvalTree>),
29}
30
31impl EvalTree {
32    /// Convert an [`Atom`] into an owned `EvalTree`.
33    ///
34    /// This dissociates the expression from the arena, allowing the
35    /// compilation pipeline to work without lifetime constraints.
36    ///
37    /// # Example
38    ///
39    /// ```ignore
40    /// use ocas_atom::AtomArena;
41    /// use ocas_eval::EvalTree;
42    ///
43    /// let arena = ocas_core::arena::Arena::new();
44    /// let ctx = AtomArena::new(&arena);
45    /// let atom = ctx.add(&[ctx.var("x"), ctx.num(2)]);
46    /// let tree = EvalTree::from_atom(atom);
47    /// ```
48    pub fn from_atom(atom: Atom<'_>) -> Self {
49        match atom.node() {
50            AtomNode::Num(n) => EvalTree::Num(*n as f64),
51            AtomNode::Var(s) => EvalTree::Var(s.as_str().to_string()),
52            AtomNode::Fun(s, args) => {
53                let converted: Vec<EvalTree> =
54                    args.iter().map(|a| EvalTree::from_atom(*a)).collect();
55                EvalTree::Fun(s.as_str().to_string(), converted)
56            }
57            AtomNode::Add(terms) => {
58                let converted: Vec<EvalTree> =
59                    terms.iter().map(|a| EvalTree::from_atom(*a)).collect();
60                EvalTree::Add(converted)
61            }
62            AtomNode::Mul(factors) => {
63                let converted: Vec<EvalTree> =
64                    factors.iter().map(|a| EvalTree::from_atom(*a)).collect();
65                EvalTree::Mul(converted)
66            }
67            AtomNode::Pow(base, exp) => EvalTree::Pow(
68                Box::new(EvalTree::from_atom(*base)),
69                Box::new(EvalTree::from_atom(*exp)),
70            ),
71        }
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use ocas_atom::AtomArena;
79    use ocas_core::arena::Arena;
80
81    #[test]
82    fn from_atom_num() {
83        let arena = Arena::new();
84        let ctx = AtomArena::new(&arena);
85        let atom = ctx.num(42);
86        let tree = EvalTree::from_atom(atom);
87        assert_eq!(tree, EvalTree::Num(42.0));
88    }
89
90    #[test]
91    fn from_atom_var() {
92        let arena = Arena::new();
93        let ctx = AtomArena::new(&arena);
94        let atom = ctx.var("x");
95        let tree = EvalTree::from_atom(atom);
96        assert_eq!(tree, EvalTree::Var("x".into()));
97    }
98
99    #[test]
100    fn from_atom_add() {
101        let arena = Arena::new();
102        let ctx = AtomArena::new(&arena);
103        let atom = ctx.add(&[ctx.var("x"), ctx.num(2)]);
104        let tree = EvalTree::from_atom(atom);
105        assert_eq!(
106            tree,
107            EvalTree::Add(vec![EvalTree::Var("x".into()), EvalTree::Num(2.0)])
108        );
109    }
110
111    #[test]
112    fn from_atom_pow() {
113        let arena = Arena::new();
114        let ctx = AtomArena::new(&arena);
115        let atom = ctx.pow(ctx.var("x"), ctx.num(2));
116        let tree = EvalTree::from_atom(atom);
117        assert_eq!(
118            tree,
119            EvalTree::Pow(
120                Box::new(EvalTree::Var("x".into())),
121                Box::new(EvalTree::Num(2.0))
122            )
123        );
124    }
125}