Skip to main content

prima_core/
simplify.rs

1use num_bigint::BigInt;
2use num_rational::BigRational;
3
4use crate::builtins::BuiltinSymbols;
5use crate::expr_pool::{ExprData, ExprId, ExprPool};
6use crate::number::Number;
7
8/// Full simplification (spec §8.3 levels 2/3): fold recursively, then apply rules on demand —
9/// `Pow(sqrt(x), 2) → x`, Euler's `e^{iθ} → cos + i·sin`, constant folding for `sin/cos/exp/log/ln/abs/sqrt`,
10/// `Pow(x, 1/2) → \sqrt{x}`. Simplification never changes the mathematical value.
11pub fn simplify(pool: &ExprPool, builtins: &BuiltinSymbols, id: ExprId) -> ExprId {
12    simplify_at(pool, builtins, id, 3)
13}
14
15/// Simplification gated by a level (spec §8.3/§13.2): a lower `simplify_level` reduces how many
16/// rules are applied, while intern-time level-0/1 canonicalization (`0*x→0`, `1*x→x`, constant
17/// merging) always holds. Level `>= 2` enables the symbolic rules (builtin constant folding,
18/// `sqrt`-elimination, Euler's formula); level `3` is a superset for future rationalization rules.
19pub fn simplify_at(pool: &ExprPool, builtins: &BuiltinSymbols, id: ExprId, level: u8) -> ExprId {
20    let node = match pool.get(id) {
21        Some(n) => n,
22        None => return id,
23    };
24    match node {
25        ExprData::Add(items) => {
26            if items.is_empty() {
27                return id;
28            }
29            let mut acc = simplify_at(pool, builtins, items[0], level);
30            for &it in &items[1..] {
31                let s = simplify_at(pool, builtins, it, level);
32                acc = pool.add2(acc, s);
33            }
34            acc
35        }
36        ExprData::Mul(items) => {
37            if items.is_empty() {
38                return id;
39            }
40            let mut acc = simplify_at(pool, builtins, items[0], level);
41            for &it in &items[1..] {
42                let s = simplify_at(pool, builtins, it, level);
43                acc = pool.mul2(acc, s);
44            }
45            acc
46        }
47        ExprData::Pow { base, exp } => {
48            let b = simplify_at(pool, builtins, base, level);
49            let e = simplify_at(pool, builtins, exp, level);
50            if level >= 2 {
51                if let Some(ExprData::Apply { f, args }) = pool.get(b)
52                    && f == pool.symbol(builtins.sqrt)
53                    && args.len() == 1
54                    && e == pool.integer(2)
55                {
56                    return simplify_at(pool, builtins, args[0], level);
57                }
58                if b == pool.symbol(builtins.e) && let Some(r) = euler(pool, builtins, e) {
59                    return r;
60                }
61            }
62            pool.pow2(b, e)
63        }
64        ExprData::Apply { f, args } => {
65            let mut new_args = Vec::with_capacity(args.len());
66            for &a in args.iter() {
67                new_args.push(simplify_at(pool, builtins, a, level));
68            }
69            if level >= 2
70                && let Some(r) = apply_rule(pool, builtins, f, &new_args)
71            {
72                return r;
73            }
74            pool.apply(f, &new_args)
75        }
76        _ => id,
77    }
78}
79
80fn apply_rule(pool: &ExprPool, builtins: &BuiltinSymbols, f: ExprId, args: &[ExprId]) -> Option<ExprId> {
81    if args.len() != 1 {
82        return None;
83    }
84    let arg = args[0];
85    if f == pool.symbol(builtins.sqrt) {
86        if let Some(n) = pool.const_number(arg) && let Some(s) = n.sqrt() {
87            return Some(pool.number(&s));
88        }
89        return None;
90    }
91    let sin = pool.symbol(builtins.sin);
92    let cos = pool.symbol(builtins.cos);
93    let tan = pool.symbol(builtins.tan);
94    if f == sin || f == cos || f == tan {
95        if let Some((c, s)) = trig_of_angle(pool, builtins, arg) {
96            if f == sin {
97                return Some(pool.number(&s));
98            }
99            if f == cos {
100                return Some(pool.number(&c));
101            }
102            if !c.is_zero() {
103                return Some(pool.number(&(s / c)));
104            }
105        }
106        return None;
107    }
108    if f == pool.symbol(builtins.exp) {
109        if arg == pool.integer(0) {
110            return Some(pool.integer(1));
111        }
112        return None;
113    }
114    if f == pool.symbol(builtins.log) || f == pool.symbol(builtins.ln) {
115        if arg == pool.integer(1) {
116            return Some(pool.integer(0));
117        }
118        if arg == pool.symbol(builtins.e) {
119            return Some(pool.integer(1));
120        }
121        return None;
122    }
123    if f == pool.symbol(builtins.abs) {
124        if let Some(n) = pool.const_number(arg) {
125            return Some(pool.number(&n.abs()));
126        }
127        return None;
128    }
129    None
130}
131
132/// An angle of the form `k·\pi`: extract the rational coefficient k, then consult the exact trig table (spec §7 built-in symbol simplification).
133fn trig_of_angle(pool: &ExprPool, builtins: &BuiltinSymbols, expr: ExprId) -> Option<(Number, Number)> {
134    let c = rational_pi_coefficient(pool, builtins, expr)?;
135    exact_trig(&c)
136}
137
138fn rational_pi_coefficient(pool: &ExprPool, builtins: &BuiltinSymbols, expr: ExprId) -> Option<BigRational> {
139    let node = pool.get(expr)?;
140    match node {
141        ExprData::Symbol(s) if s == builtins.pi => Some(BigRational::new(BigInt::from(1), BigInt::from(1))),
142        ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Real(_) => {
143            let n = pool.const_number(expr)?;
144            if n.is_zero() {
145                Some(BigRational::new(BigInt::from(0), BigInt::from(1)))
146            } else {
147                None
148            }
149        }
150        ExprData::Mul(items) => {
151            let mut coeff: Option<BigRational> = None;
152            let mut found_pi = false;
153            for &it in items.iter() {
154                match pool.get(it)? {
155                    ExprData::Symbol(s) if s == builtins.pi => found_pi = true,
156                    ExprData::Integer(_) | ExprData::Rational(_) => {
157                        let c = match pool.const_number(it)? {
158                            Number::Integer(i) => BigRational::from_integer(i),
159                            Number::Rational(r) => r,
160                            _ => return None,
161                        };
162                        coeff = Some(match coeff {
163                            Some(acc) => acc * c,
164                            None => c,
165                        });
166                    }
167                    _ => return None,
168                }
169            }
170            if !found_pi {
171                return None;
172            }
173            Some(coeff.unwrap_or_else(|| BigRational::new(BigInt::from(1), BigInt::from(1))))
174        }
175        _ => None,
176    }
177}
178
179fn exact_trig(c: &BigRational) -> Option<(Number, Number)> {
180    let two = BigRational::new(BigInt::from(2), BigInt::from(1));
181    let mut c = c % two.clone();
182    if c < BigRational::new(BigInt::from(0), BigInt::from(1)) {
183        c += two;
184    }
185    let zero = BigRational::new(BigInt::from(0), BigInt::from(1));
186    let half = BigRational::new(BigInt::from(1), BigInt::from(2));
187    let one = BigRational::new(BigInt::from(1), BigInt::from(1));
188    let three_halves = BigRational::new(BigInt::from(3), BigInt::from(2));
189    if c == zero {
190        Some((Number::from(1), Number::from(0)))
191    } else if c == half {
192        Some((Number::from(0), Number::from(1)))
193    } else if c == one {
194        Some((Number::from(-1), Number::from(0)))
195    } else if c == three_halves {
196        Some((Number::from(0), Number::from(-1)))
197    } else {
198        None
199    }
200}
201
202/// Euler's formula (spec §7.4): fold `e^{iθ}` into `cosθ + i·sinθ` using the exact trig values of θ.
203fn euler(pool: &ExprPool, builtins: &BuiltinSymbols, z: ExprId) -> Option<ExprId> {
204    let i = pool.symbol(builtins.i);
205    let theta = match pool.get(z)? {
206        ExprData::Symbol(s) if s == builtins.i => return None,
207        ExprData::Mul(items) => {
208            let mut theta_items = Vec::new();
209            let mut has_i = false;
210            for &it in items.iter() {
211                if it == i {
212                    has_i = true;
213                } else {
214                    theta_items.push(it);
215                }
216            }
217            if !has_i || theta_items.is_empty() {
218                return None;
219            }
220            let mut acc = theta_items[0];
221            for &it in &theta_items[1..] {
222                acc = pool.mul2(acc, it);
223            }
224            acc
225        }
226        _ => return None,
227    };
228    let (c, s) = trig_of_angle(pool, builtins, theta)?;
229    if s == Number::from(0) {
230        Some(pool.number(&c))
231    } else if c == Number::from(0) && s == Number::from(1) {
232        Some(i)
233    } else if c == Number::from(0) && s == Number::from(-1) {
234        Some(pool.mul2(pool.integer(-1), i))
235    } else {
236        None
237    }
238}