1use num_bigint::BigInt;
9use num_integer::Integer;
10use num_traits::{One, Signed, ToPrimitive, Zero};
11
12#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub enum SymbolicExpr {
15 Integer(BigInt),
16 Rational(BigInt, BigInt),
18 Sqrt { radicand: BigInt },
20 ScaledSqrt { coeff: (BigInt, BigInt), rad: BigInt },
22 Pi,
23 E,
24 Add(Vec<SymbolicExpr>),
25 Mul(Vec<SymbolicExpr>),
26 Pow { base: Box<SymbolicExpr>, exp: Box<SymbolicExpr> },
27 Sin(Box<SymbolicExpr>),
28 Cos(Box<SymbolicExpr>),
29 Exp(Box<SymbolicExpr>),
30 Ln(Box<SymbolicExpr>),
31}
32
33use SymbolicExpr::*;
34
35impl SymbolicExpr {
36 pub fn int(n: i64) -> Self {
37 Integer(BigInt::from(n))
38 }
39 pub fn rational(p: i64, q: i64) -> Self {
40 Rational(BigInt::from(p), BigInt::from(q))
41 }
42 pub fn sqrt(n: i64) -> Self {
43 Sqrt { radicand: BigInt::from(n) }
44 }
45 pub fn add(terms: Vec<SymbolicExpr>) -> Self {
46 Add(terms)
47 }
48 pub fn mul(factors: Vec<SymbolicExpr>) -> Self {
49 Mul(factors)
50 }
51 pub fn sin(x: SymbolicExpr) -> Self {
52 Sin(Box::new(x))
53 }
54 pub fn cos(x: SymbolicExpr) -> Self {
55 Cos(Box::new(x))
56 }
57 pub fn exp(x: SymbolicExpr) -> Self {
58 Exp(Box::new(x))
59 }
60 pub fn ln(x: SymbolicExpr) -> Self {
61 Ln(Box::new(x))
62 }
63
64 fn as_rational(&self) -> Option<(BigInt, BigInt)> {
66 match self {
67 Integer(n) => Some((n.clone(), BigInt::one())),
68 Rational(p, q) => Some((p.clone(), q.clone())),
69 _ => None,
70 }
71 }
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum TowerLevel {
77 Integer,
78 Rational,
79 Algebraic,
80 Symbolic,
81 Transcendental,
82}
83
84pub fn tower_level(expr: &SymbolicExpr) -> TowerLevel {
86 match expr {
87 Integer(_) => TowerLevel::Integer,
88 Rational(_, _) => TowerLevel::Rational,
89 Sqrt { .. } | ScaledSqrt { .. } => TowerLevel::Algebraic,
90 Pi | E | Sin(_) | Cos(_) | Exp(_) | Ln(_) => TowerLevel::Transcendental,
91 Add(t) => t.iter().map(tower_level).max_by_key(level_rank).unwrap_or(TowerLevel::Integer),
92 Mul(t) => t.iter().map(tower_level).max_by_key(level_rank).unwrap_or(TowerLevel::Integer),
93 Pow { base, .. } => tower_level(base).max_symbolic(),
94 }
95}
96
97fn level_rank(l: &TowerLevel) -> u8 {
98 match l {
99 TowerLevel::Integer => 0,
100 TowerLevel::Rational => 1,
101 TowerLevel::Algebraic => 2,
102 TowerLevel::Symbolic => 3,
103 TowerLevel::Transcendental => 4,
104 }
105}
106
107impl TowerLevel {
108 fn max_symbolic(self) -> TowerLevel {
109 if level_rank(&self) >= level_rank(&TowerLevel::Symbolic) {
110 self
111 } else {
112 TowerLevel::Symbolic
113 }
114 }
115}
116
117pub struct IdentityGraph;
121
122impl Default for IdentityGraph {
123 fn default() -> Self {
124 Self::standard()
125 }
126}
127
128impl IdentityGraph {
129 pub fn standard() -> Self {
131 IdentityGraph
132 }
133
134 pub fn simplify(&self, expr: SymbolicExpr) -> SymbolicExpr {
136 let mut current = expr;
137 for _ in 0..64 {
138 let next = self.step(current.clone());
139 if next == current {
140 return next;
141 }
142 current = next;
143 }
144 current
145 }
146
147 fn step(&self, expr: SymbolicExpr) -> SymbolicExpr {
149 match expr {
150 Add(terms) => self.simplify_add(terms),
151 Mul(factors) => self.simplify_mul(factors),
152 Sin(x) => self.simplify_sin(self.step(*x)),
153 Cos(x) => self.simplify_cos(self.step(*x)),
154 Exp(x) => self.simplify_exp(self.step(*x)),
155 Ln(x) => self.simplify_ln(self.step(*x)),
156 Pow { base, exp } => Pow {
157 base: Box::new(self.step(*base)),
158 exp: Box::new(self.step(*exp)),
159 },
160 Rational(p, q) => normalize_rational(p, q),
161 other => other,
162 }
163 }
164
165 fn simplify_add(&self, terms: Vec<SymbolicExpr>) -> SymbolicExpr {
166 let mut const_num = BigInt::zero();
167 let mut const_den = BigInt::one();
168 let mut others: Vec<SymbolicExpr> = Vec::new();
169 for t in terms {
170 let t = self.step(t);
171 match &t {
172 Add(inner) => {
173 for it in inner.clone() {
175 self.accumulate_add(it, &mut const_num, &mut const_den, &mut others);
176 }
177 }
178 _ => self.accumulate_add(t, &mut const_num, &mut const_den, &mut others),
179 }
180 }
181 others.sort_by(canonical_cmp);
184 let mut result: Vec<SymbolicExpr> = Vec::new();
185 if !const_num.is_zero() {
186 result.push(normalize_rational(const_num, const_den));
187 }
188 result.append(&mut others);
189 match result.len() {
190 0 => SymbolicExpr::int(0),
191 1 => result.into_iter().next().unwrap(),
192 _ => Add(result),
193 }
194 }
195
196 fn accumulate_add(
197 &self,
198 t: SymbolicExpr,
199 num: &mut BigInt,
200 den: &mut BigInt,
201 others: &mut Vec<SymbolicExpr>,
202 ) {
203 if let Some((p, q)) = t.as_rational() {
204 *num = &*num * &q + &p * &*den;
206 *den = &*den * &q;
207 } else {
208 others.push(t);
209 }
210 }
211
212 fn simplify_mul(&self, factors: Vec<SymbolicExpr>) -> SymbolicExpr {
213 let mut coeff_num = BigInt::one();
214 let mut coeff_den = BigInt::one();
215 let mut radicand = BigInt::one();
216 let mut others: Vec<SymbolicExpr> = Vec::new();
217 let mut is_zero = false;
218
219 let mut stack: Vec<SymbolicExpr> = factors.into_iter().map(|f| self.step(f)).collect();
220 while let Some(f) = stack.pop() {
221 match f {
222 Mul(inner) => stack.extend(inner.into_iter().map(|f| self.step(f))),
223 Integer(n) => {
224 if n.is_zero() {
225 is_zero = true;
226 }
227 coeff_num *= n;
228 }
229 Rational(p, q) => {
230 if p.is_zero() {
231 is_zero = true;
232 }
233 coeff_num *= p;
234 coeff_den *= q;
235 }
236 Sqrt { radicand: r } => radicand *= r,
237 ScaledSqrt { coeff: (a, b), rad } => {
238 coeff_num *= a;
239 coeff_den *= b;
240 radicand *= rad;
241 }
242 other => others.push(other),
243 }
244 }
245
246 if is_zero {
247 return SymbolicExpr::int(0);
248 }
249
250 if !radicand.is_one() {
252 match simplify_sqrt(radicand) {
253 Integer(k) => coeff_num *= k,
254 ScaledSqrt { coeff: (a, b), rad } => {
255 coeff_num *= a;
256 coeff_den *= b;
257 others.push(Sqrt { radicand: rad });
258 }
259 Sqrt { radicand: r } => others.push(Sqrt { radicand: r }),
260 e => others.push(e),
261 }
262 }
263
264 let g = coeff_num.gcd(&coeff_den);
266 if !g.is_zero() {
267 coeff_num /= &g;
268 coeff_den /= &g;
269 }
270 if coeff_den.is_negative() {
271 coeff_num = -coeff_num;
272 coeff_den = -coeff_den;
273 }
274
275 let coeff_is_one = coeff_num.is_one() && coeff_den.is_one();
276
277 if others.len() == 1 {
279 if let Sqrt { radicand: r } = &others[0] {
280 if coeff_is_one {
281 return Sqrt { radicand: r.clone() };
282 }
283 return ScaledSqrt {
284 coeff: (coeff_num, coeff_den),
285 rad: r.clone(),
286 };
287 }
288 }
289
290 others.sort_by(canonical_cmp);
291 let mut result: Vec<SymbolicExpr> = Vec::new();
292 if !coeff_is_one {
293 result.push(normalize_rational(coeff_num, coeff_den));
294 }
295 result.append(&mut others);
296 match result.len() {
297 0 => SymbolicExpr::int(1),
298 1 => result.into_iter().next().unwrap(),
299 _ => Mul(result),
300 }
301 }
302
303 fn simplify_sin(&self, x: SymbolicExpr) -> SymbolicExpr {
304 if let Integer(n) = &x {
305 if n.is_zero() {
306 return SymbolicExpr::int(0);
307 }
308 }
309 if let Some((a, b)) = as_pi_multiple(&x) {
310 let (a, b) = reduce(a, b);
311 if let Some(v) = sin_pi_table(&a, &b) {
313 return v;
314 }
315 }
316 Sin(Box::new(x))
317 }
318
319 fn simplify_cos(&self, x: SymbolicExpr) -> SymbolicExpr {
320 if let Integer(n) = &x {
321 if n.is_zero() {
322 return SymbolicExpr::int(1);
323 }
324 }
325 if let Some((a, b)) = as_pi_multiple(&x) {
326 let (a, b) = reduce(a, b);
327 if let Some(v) = cos_pi_table(&a, &b) {
328 return v;
329 }
330 }
331 Cos(Box::new(x))
332 }
333
334 fn simplify_exp(&self, x: SymbolicExpr) -> SymbolicExpr {
335 if let Integer(n) = &x {
336 if n.is_zero() {
337 return SymbolicExpr::int(1);
338 }
339 }
340 Exp(Box::new(x))
345 }
346
347 fn simplify_ln(&self, x: SymbolicExpr) -> SymbolicExpr {
348 if let Integer(n) = &x {
349 if n.is_one() {
350 return SymbolicExpr::int(0);
351 }
352 }
353 Ln(Box::new(x))
354 }
355}
356
357fn canonical_cmp(a: &SymbolicExpr, b: &SymbolicExpr) -> std::cmp::Ordering {
362 format!("{a:?}").cmp(&format!("{b:?}"))
363}
364
365fn normalize_rational(mut p: BigInt, mut q: BigInt) -> SymbolicExpr {
367 if q.is_zero() {
368 return Rational(p, q); }
370 if q.is_negative() {
371 p = -p;
372 q = -q;
373 }
374 let g = p.gcd(&q);
375 if !g.is_zero() {
376 p /= &g;
377 q /= &g;
378 }
379 if q.is_one() {
380 Integer(p)
381 } else {
382 Rational(p, q)
383 }
384}
385
386fn reduce(mut a: BigInt, mut b: BigInt) -> (BigInt, BigInt) {
388 if b.is_negative() {
389 a = -a;
390 b = -b;
391 }
392 let g = a.gcd(&b);
393 if !g.is_zero() {
394 a /= &g;
395 b /= &g;
396 }
397 (a, b)
398}
399
400fn simplify_sqrt(n: BigInt) -> SymbolicExpr {
402 if n.is_negative() || n.is_zero() {
403 return Sqrt { radicand: n };
404 }
405 let nu = match n.to_u128() {
406 Some(v) => v,
407 None => return Sqrt { radicand: n },
408 };
409 let mut square = 1u128;
410 let mut rad = nu;
411 let mut d = 2u128;
412 while d * d <= rad {
413 while rad % (d * d) == 0 {
414 rad /= d * d;
415 square *= d;
416 }
417 d += 1;
418 }
419 let s = BigInt::from(square);
420 let r = BigInt::from(rad);
421 if rad == 1 {
422 Integer(s)
423 } else if square == 1 {
424 Sqrt { radicand: r }
425 } else {
426 ScaledSqrt { coeff: (s, BigInt::one()), rad: r }
427 }
428}
429
430fn as_pi_multiple(expr: &SymbolicExpr) -> Option<(BigInt, BigInt)> {
432 match expr {
433 Pi => Some((BigInt::one(), BigInt::one())),
434 Mul(factors) => {
435 let mut num = BigInt::one();
436 let mut den = BigInt::one();
437 let mut pi_count = 0;
438 for f in factors {
439 match f {
440 Pi => pi_count += 1,
441 Integer(n) => num *= n,
442 Rational(p, q) => {
443 num *= p;
444 den *= q;
445 }
446 _ => return None,
447 }
448 }
449 if pi_count == 1 {
450 Some((num, den))
451 } else {
452 None
453 }
454 }
455 _ => None,
456 }
457}
458
459fn frac_is(a: &BigInt, b: &BigInt, n: i64, d: i64) -> bool {
460 *a == BigInt::from(n) && *b == BigInt::from(d)
461}
462
463fn sin_pi_table(a: &BigInt, b: &BigInt) -> Option<SymbolicExpr> {
465 if a.is_zero() {
466 return Some(SymbolicExpr::int(0));
467 }
468 if frac_is(a, b, 1, 1) {
469 return Some(SymbolicExpr::int(0)); }
471 if frac_is(a, b, 1, 6) {
472 return Some(SymbolicExpr::rational(1, 2));
473 }
474 if frac_is(a, b, 1, 4) {
475 return Some(ScaledSqrt { coeff: (BigInt::one(), BigInt::from(2)), rad: BigInt::from(2) });
476 }
477 if frac_is(a, b, 1, 3) {
478 return Some(ScaledSqrt { coeff: (BigInt::one(), BigInt::from(2)), rad: BigInt::from(3) });
479 }
480 if frac_is(a, b, 1, 2) {
481 return Some(SymbolicExpr::int(1)); }
483 None
484}
485
486fn cos_pi_table(a: &BigInt, b: &BigInt) -> Option<SymbolicExpr> {
488 if a.is_zero() {
489 return Some(SymbolicExpr::int(1));
490 }
491 if frac_is(a, b, 1, 1) {
492 return Some(SymbolicExpr::int(-1)); }
494 if frac_is(a, b, 1, 2) {
495 return Some(SymbolicExpr::int(0)); }
497 None
498}
499
500#[cfg(test)]
501mod tests {
502 use super::*;
503
504 fn g() -> IdentityGraph {
505 IdentityGraph::standard()
506 }
507
508 #[test]
509 fn sin_pi_is_zero() {
510 assert_eq!(g().simplify(SymbolicExpr::sin(Pi)), SymbolicExpr::int(0));
511 }
512
513 #[test]
514 fn cos_pi_is_minus_one() {
515 assert_eq!(g().simplify(SymbolicExpr::cos(Pi)), SymbolicExpr::int(-1));
516 }
517
518 #[test]
519 fn sin_pi_over_six() {
520 let expr = SymbolicExpr::sin(Mul(vec![SymbolicExpr::rational(1, 6), Pi]));
521 assert_eq!(g().simplify(expr), SymbolicExpr::rational(1, 2));
522 }
523
524 #[test]
525 fn exp_zero_is_one() {
526 assert_eq!(g().simplify(SymbolicExpr::exp(SymbolicExpr::int(0))), SymbolicExpr::int(1));
527 }
528
529 #[test]
530 fn ln_one_is_zero() {
531 assert_eq!(g().simplify(SymbolicExpr::ln(SymbolicExpr::int(1))), SymbolicExpr::int(0));
532 }
533
534 #[test]
535 fn sqrt_times_sqrt() {
536 let expr = Mul(vec![SymbolicExpr::sqrt(2), SymbolicExpr::sqrt(2)]);
537 assert_eq!(g().simplify(expr), SymbolicExpr::int(2));
538 }
539
540 #[test]
541 fn x_times_zero() {
542 let expr = Mul(vec![Pi, SymbolicExpr::int(0)]);
543 assert_eq!(g().simplify(expr), SymbolicExpr::int(0));
544 }
545
546 #[test]
547 fn add_zero_identity() {
548 let expr = Add(vec![Pi, SymbolicExpr::int(0)]);
549 assert_eq!(g().simplify(expr), Pi);
550 }
551
552 #[test]
553 fn mul_one_identity() {
554 let expr = Mul(vec![Pi, SymbolicExpr::int(1)]);
555 assert_eq!(g().simplify(expr), Pi);
556 }
557
558 #[test]
559 fn sqrt_eight_simplifies() {
560 assert_eq!(
562 simplify_sqrt(BigInt::from(8)),
563 ScaledSqrt { coeff: (BigInt::from(2), BigInt::one()), rad: BigInt::from(2) }
564 );
565 }
566
567 #[test]
568 fn classification() {
569 assert_eq!(tower_level(&SymbolicExpr::int(3)), TowerLevel::Integer);
570 assert_eq!(tower_level(&SymbolicExpr::sqrt(2)), TowerLevel::Algebraic);
571 assert_eq!(tower_level(&Pi), TowerLevel::Transcendental);
572 }
573}