prima_core/
opt.rs

1//! Optimization pipeline primitives (spec §10.2): constant folding + common-subexpression
2//! elimination over the hash-consed `ExprDAG`.
3
4use crate::builtins::BuiltinSymbols;
5use crate::expr_pool::{ExprId, ExprPool};
6
7/// Constant folding (spec §10.2 item 1): runs the simplify engine over the DAG and returns the
8/// canonical folded expression. Level-2 rules (0*x, 1*x, constant arithmetic, math constants)
9/// already live in `crate::simplify::simplify`.
10pub fn const_fold(pool: &ExprPool, builtins: &BuiltinSymbols, id: ExprId) -> ExprId {
11    crate::simplify::simplify(pool, builtins, id)
12}
13
14/// CSE (spec §10.2 item 3): returns the canonical `ExprId` for a subexpression. Because the
15/// `ExprPool` hash-conses identical nodes, repeated subexpressions already share one `ExprId`;
16/// this function documents and asserts that invariant (returns `id` unchanged). It exists as the
17/// stable entry point the JIT pipeline calls before codegen.
18pub fn cse(_pool: &ExprPool, id: ExprId) -> ExprId {
19    id
20}
21
22/// Full local optimization run (spec §10.2): `const_fold` then `cse`. Used by the JIT compiler
23/// before translating the DAG to bytecode.
24pub fn optimize(pool: &ExprPool, builtins: &BuiltinSymbols, id: ExprId) -> ExprId {
25    let folded = const_fold(pool, builtins, id);
26    cse(pool, folded)
27}
28
29#[cfg(test)]
30mod tests {
31    use crate::BuiltinSymbols;
32    use crate::expr_pool::{ExprData, ExprPool};
33    use crate::number::Number;
34    use crate::opt::{self, const_fold, cse, optimize};
35    use crate::render::render_latex;
36    use crate::symbol::SymbolTable;
37
38    #[test]
39    fn const_fold_merges_constant_arithmetic() {
40        let pool = ExprPool::global();
41        let builtins = BuiltinSymbols::global();
42        let symbols = SymbolTable::global();
43        let x = pool.symbol(symbols.intern("x"));
44        // Raw, unsimplified `2*3`: `Mul` interned as-is (no level-0/1 folding at build time).
45        let prod = pool.mul(&[pool.integer(2), pool.integer(3)]);
46        assert!(matches!(pool.get(prod), Some(ExprData::Mul(_))));
47        let expr = pool.add(&[prod, x]);
48
49        let folded = const_fold(pool, builtins, expr);
50        assert_ne!(folded, prod);
51        match pool.get(folded) {
52            Some(ExprData::Add(items)) => {
53                assert!(
54                    items.iter().any(
55                        |&it| matches!(pool.const_number(it), Some(n) if n == Number::from(6))
56                    ),
57                    "constant 6 must survive as an `Add` child: {:?}",
58                    items
59                );
60            }
61            other => panic!("expected `Add`, got {:?}", other),
62        }
63        assert_eq!(render_latex(pool, symbols, folded), "x + 6");
64    }
65
66    #[test]
67    fn const_fold_folds_math_functions() {
68        let pool = ExprPool::global();
69        let builtins = BuiltinSymbols::global();
70
71        let sin0 = pool.apply(pool.symbol(builtins.sin), &[pool.integer(0)]);
72        assert_eq!(const_fold(pool, builtins, sin0), pool.integer(0));
73
74        let sqrt4 = pool.apply(pool.symbol(builtins.sqrt), &[pool.integer(4)]);
75        assert_eq!(const_fold(pool, builtins, sqrt4), pool.integer(2));
76    }
77
78    #[test]
79    fn cse_shares_duplicate_subexpressions() {
80        let pool = ExprPool::global();
81        let symbols = SymbolTable::global();
82        let x = pool.symbol(symbols.intern("x"));
83
84        // `x*x` interned twice must yield the SAME `ExprId`: one shared `Mul` node (spec §8.1
85        // hash-consing is the CSE machinery for the DAG).
86        let m1 = pool.mul2(x, x);
87        let m2 = pool.mul2(x, x);
88        assert_eq!(m1, m2);
89        assert_eq!(cse(pool, m1), m1);
90
91        // `x*x + x*x`: the `Add` refers to the same `Mul` id twice — a single shared subexpression.
92        let e = pool.add2(m1, m2);
93        match pool.get(e) {
94            Some(ExprData::Add(items)) => {
95                assert_eq!(items.len(), 2);
96                assert_eq!(items[0], items[1]);
97                assert_eq!(items[0], m1);
98            }
99            other => panic!("expected `Add`, got {:?}", other),
100        }
101    }
102
103    #[test]
104    fn optimize_folds_then_canonicalizes() {
105        let pool = ExprPool::global();
106        let builtins = BuiltinSymbols::global();
107        let symbols = SymbolTable::global();
108        let x = pool.symbol(symbols.intern("x"));
109        let expr = pool.add(&[pool.mul(&[pool.integer(2), pool.integer(3)]), x]);
110
111        let folded = optimize(pool, builtins, expr);
112        // `optimize` is const_fold followed by cse; both must agree on the canonical form.
113        assert_eq!(folded, const_fold(pool, builtins, expr));
114        assert_eq!(folded, opt::optimize(pool, builtins, folded));
115        assert_eq!(render_latex(pool, symbols, folded), "x + 6");
116    }
117}
prima_core/opt.rs

prima_core/
opt.rs
