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
8pub fn simplify(pool: &ExprPool, builtins: &BuiltinSymbols, id: ExprId) -> ExprId {
12 simplify_at(pool, builtins, id, 3)
13}
14
15pub 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
132fn 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
202fn 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}