1use super::*;
4
5pub(super) fn polynomial_identity(left: &SymExpr, right: &SymExpr) -> bool {
6 if !polynomial_normalization_can_help(left) && !polynomial_normalization_can_help(right) {
7 return false;
8 }
9 matches!(
10 (Polynomial::from_expr(left), Polynomial::from_expr(right)),
11 (Some(left), Some(right)) if left == right
12 )
13}
14
15fn polynomial_normalization_can_help(expr: &SymExpr) -> bool {
16 let crosses_sum_product_boundary = match expr.kind() {
17 SymExprKind::BinOp(SymBinOp::Mul, left, right) => {
18 matches!(left.kind(), SymExprKind::BinOp(SymBinOp::Add | SymBinOp::Sub, ..))
19 || matches!(right.kind(), SymExprKind::BinOp(SymBinOp::Add | SymBinOp::Sub, ..))
20 }
21 SymExprKind::BinOp(SymBinOp::Add | SymBinOp::Sub, left, right) => {
22 matches!(left.kind(), SymExprKind::BinOp(SymBinOp::Mul, ..))
23 || matches!(right.kind(), SymExprKind::BinOp(SymBinOp::Mul, ..))
24 || matches!(
25 left.kind(),
26 SymExprKind::BinOp(SymBinOp::Shl, _, shift)
27 if shift.as_const().is_some_and(|shift| shift < U256::from(256))
28 )
29 || matches!(
30 right.kind(),
31 SymExprKind::BinOp(SymBinOp::Shl, _, shift)
32 if shift.as_const().is_some_and(|shift| shift < U256::from(256))
33 )
34 }
35 _ => false,
36 };
37 if !crosses_sum_product_boundary {
38 return false;
39 }
40
41 fn ring_shape(
42 expr: &SymExpr,
43 shapes: &mut HashMap<SymExpr, Option<(usize, usize)>>,
44 remaining: &mut usize,
45 ) -> Option<(usize, usize)> {
46 if let Some(shape) = shapes.get(expr) {
47 return *shape;
48 }
49 if *remaining == 0 {
50 return None;
51 }
52 *remaining -= 1;
53 let shape = (|| match expr.kind() {
54 SymExprKind::Const(_) | SymExprKind::Var(_) => Some((0, 0)),
55 SymExprKind::BinOp(
56 op @ (SymBinOp::Add | SymBinOp::Sub | SymBinOp::Mul),
57 left,
58 right,
59 ) => {
60 let left = ring_shape(left, shapes, remaining)?;
61 let right = ring_shape(right, shapes, remaining)?;
62 let operations = left.0.saturating_add(right.0).saturating_add(1);
63 let multiplications = left
64 .1
65 .saturating_add(right.1)
66 .saturating_add(usize::from(*op == SymBinOp::Mul));
67 Some((operations, multiplications))
68 }
69 SymExprKind::BinOp(SymBinOp::Shl, value, shift)
70 if shift.as_const().is_some_and(|shift| shift < U256::from(256)) =>
71 {
72 let shape = ring_shape(value, shapes, remaining)?;
73 Some((shape.0.saturating_add(1), shape.1.saturating_add(1)))
74 }
75 _ => None,
76 })();
77 shapes.insert(expr.clone(), shape);
78 shape
79 }
80
81 let mut shapes = HashMap::default();
82 let mut remaining = MAX_LOCAL_ANALYSIS_NODES;
83 ring_shape(expr, &mut shapes, &mut remaining)
84 .is_some_and(|(operations, multiplications)| operations > 1 && multiplications > 0)
85}
86
87const MAX_POLYNOMIAL_TERMS: usize = 32;
91const MAX_MONOMIAL_FACTORS: usize = 8;
92const MAX_POLYNOMIAL_PRODUCTS: usize = 256;
93
94type Monomial = Vec<SymExpr>;
95
96#[derive(Clone, PartialEq, Eq)]
102struct Polynomial {
103 terms: HashMap<Monomial, U256>,
104}
105
106impl Polynomial {
107 fn from_expr(expr: &SymExpr) -> Option<Self> {
108 let mut remaining = MAX_LOCAL_ANALYSIS_NODES;
109 Self::from_expr_cached(expr, &mut HashMap::default(), &mut remaining)
110 }
111
112 fn from_expr_cached(
113 expr: &SymExpr,
114 polynomials: &mut HashMap<SymExpr, Option<Self>>,
115 remaining: &mut usize,
116 ) -> Option<Self> {
117 if let Some(polynomial) = polynomials.get(expr) {
118 return polynomial.clone();
119 }
120 if *remaining == 0 {
121 polynomials.insert(expr.clone(), None);
122 return None;
123 }
124 *remaining -= 1;
125 let polynomial =
126 (|| match expr.kind() {
127 SymExprKind::Const(value) => Some(Self::constant(*value)),
128 SymExprKind::BinOp(SymBinOp::Add, left, right) => {
129 Self::from_expr_cached(left, polynomials, remaining)?
130 .add(Self::from_expr_cached(right, polynomials, remaining)?)
131 }
132 SymExprKind::BinOp(SymBinOp::Sub, left, right) => {
133 Self::from_expr_cached(left, polynomials, remaining)?
134 .sub(Self::from_expr_cached(right, polynomials, remaining)?)
135 }
136 SymExprKind::BinOp(SymBinOp::Mul, left, right) => {
137 Self::from_expr_cached(left, polynomials, remaining)?
138 .mul(Self::from_expr_cached(right, polynomials, remaining)?)
139 }
140 SymExprKind::BinOp(SymBinOp::Shl, value, shift)
141 if let Some(shift) = shift.as_const()
142 && shift < U256::from(256) =>
143 {
144 let coefficient = U256::ONE << usize::try_from(shift).ok()?;
145 Self::from_expr_cached(value, polynomials, remaining)?
146 .mul(Self::constant(coefficient))
147 }
148 _ => {
149 let terms = HashMap::from_iter([(vec![expr.clone()], U256::ONE)]);
150 Some(Self { terms })
151 }
152 })();
153 polynomials.insert(expr.clone(), polynomial.clone());
154 polynomial
155 }
156
157 fn constant(value: U256) -> Self {
158 let mut terms = HashMap::default();
159 if !value.is_zero() {
160 terms.insert(Vec::new(), value);
161 }
162 Self { terms }
163 }
164
165 fn add(mut self, right: Self) -> Option<Self> {
166 for (monomial, coefficient) in right.terms {
167 self.add_term(monomial, coefficient);
168 if self.terms.len() > MAX_POLYNOMIAL_TERMS {
169 return None;
170 }
171 }
172 Some(self)
173 }
174
175 fn sub(mut self, right: Self) -> Option<Self> {
176 for (monomial, coefficient) in right.terms {
177 self.add_term(monomial, U256::ZERO.wrapping_sub(coefficient));
178 if self.terms.len() > MAX_POLYNOMIAL_TERMS {
179 return None;
180 }
181 }
182 Some(self)
183 }
184
185 fn mul(self, right: Self) -> Option<Self> {
186 let products = self.terms.len().checked_mul(right.terms.len())?;
187 if products > MAX_POLYNOMIAL_PRODUCTS {
188 return None;
189 }
190
191 let mut out = Self { terms: HashMap::default() };
192 for (left_monomial, left_coefficient) in &self.terms {
193 for (right_monomial, right_coefficient) in &right.terms {
194 let factor_count = left_monomial.len().checked_add(right_monomial.len())?;
195 if factor_count > MAX_MONOMIAL_FACTORS {
196 return None;
197 }
198 let mut monomial = Vec::with_capacity(factor_count);
199 monomial.extend(left_monomial.iter().cloned());
200 monomial.extend(right_monomial.iter().cloned());
201 SymExpr::sort_interned_factors(&mut monomial);
202 out.add_term(monomial, left_coefficient.wrapping_mul(*right_coefficient));
203 if out.terms.len() > MAX_POLYNOMIAL_TERMS {
204 return None;
205 }
206 }
207 }
208 Some(out)
209 }
210
211 fn add_term(&mut self, monomial: Monomial, coefficient: U256) {
212 if coefficient.is_zero() {
213 return;
214 }
215 let coefficient =
216 self.terms.get(&monomial).copied().unwrap_or_default().wrapping_add(coefficient);
217 if coefficient.is_zero() {
218 self.terms.remove(&monomial);
219 } else {
220 self.terms.insert(monomial, coefficient);
221 }
222 }
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228
229 #[test]
230 fn polynomial_identity_handles_shared_dag() {
231 let mut cx = SymCx::new();
232 let shared_atom = SymExpr::var(&mut cx, "shared");
233 let mut shared = shared_atom.clone();
234 for _ in 0..64 {
235 shared = SymExpr::binop(&mut cx, SymBinOp::Add, shared.clone(), shared);
236 }
237 let factor = SymExpr::var(&mut cx, "factor");
238 let expression = SymExpr::binop(&mut cx, SymBinOp::Mul, shared, factor.clone());
239 let product = SymExpr::binop(&mut cx, SymBinOp::Mul, shared_atom, factor);
240 let shift = SymExpr::constant(&mut cx, U256::from(64));
241 let expected = SymExpr::binop(&mut cx, SymBinOp::Shl, product, shift);
242
243 assert!(polynomial_identity(&expression, &expected));
244 }
245
246 #[test]
247 fn polynomial_factors_use_interned_identity_order() {
248 let mut cx = SymCx::new();
249 let left = SymExpr::var(&mut cx, "left");
250 let right = SymExpr::var(&mut cx, "right");
251 let left_right = SymExpr::from_kind(
252 &mut cx,
253 SymExprKind::BinOp(SymBinOp::Mul, left.clone(), right.clone()),
254 );
255 let right_left =
256 SymExpr::from_kind(&mut cx, SymExprKind::BinOp(SymBinOp::Mul, right, left));
257
258 assert!(Polynomial::from_expr(&left_right) == Polynomial::from_expr(&right_left));
259 }
260
261 #[test]
262 fn polynomial_identity_stops_at_factor_limit() {
263 let mut cx = SymCx::new();
264 let mut prefix = SymExpr::one(&mut cx);
265 for index in 0..MAX_MONOMIAL_FACTORS - 1 {
266 let factor = SymExpr::var(&mut cx, &format!("x_{index}"));
267 prefix = SymExpr::binop(&mut cx, SymBinOp::Mul, prefix, factor);
268 }
269 let left = SymExpr::var(&mut cx, "left");
270 let right = SymExpr::var(&mut cx, "right");
271 let sum = SymExpr::binop(&mut cx, SymBinOp::Add, left.clone(), right.clone());
272 let factored = SymExpr::binop(&mut cx, SymBinOp::Mul, prefix.clone(), sum);
273 let left_product = SymExpr::binop(&mut cx, SymBinOp::Mul, prefix.clone(), left);
274 let right_product = SymExpr::binop(&mut cx, SymBinOp::Mul, prefix, right);
275 let expanded = SymExpr::binop(&mut cx, SymBinOp::Add, left_product, right_product);
276
277 assert!(polynomial_identity(&factored, &expanded));
278
279 let extra = SymExpr::var(&mut cx, "extra");
280 let over_limit = SymExpr::binop(&mut cx, SymBinOp::Mul, factored, extra);
281
282 assert!(!polynomial_identity(&over_limit, &over_limit));
283 }
284
285 #[test]
286 fn polynomial_identity_stops_at_term_limit() {
287 let mut cx = SymCx::new();
288 let mut expression = SymExpr::zero(&mut cx);
289 for index in 0..33 {
290 let term = SymExpr::var(&mut cx, &format!("x_{index}"));
291 expression = SymExpr::binop(&mut cx, SymBinOp::Add, expression, term);
292 }
293 let factor = SymExpr::var(&mut cx, "factor");
294 expression = SymExpr::binop(&mut cx, SymBinOp::Mul, expression, factor);
295
296 assert!(!polynomial_identity(&expression, &expression));
297 }
298
299 #[test]
300 fn polynomial_identity_stops_at_product_limit() {
301 let mut cx = SymCx::new();
302 let mut left = SymExpr::zero(&mut cx);
303 for index in 0..17 {
304 let term = SymExpr::var(&mut cx, &format!("left_{index}"));
305 left = SymExpr::binop(&mut cx, SymBinOp::Add, left, term);
306 }
307 let mut right = SymExpr::zero(&mut cx);
308 for index in 0..16 {
309 let term = SymExpr::var(&mut cx, &format!("right_{index}"));
310 right = SymExpr::binop(&mut cx, SymBinOp::Add, right, term);
311 }
312 let expression = SymExpr::binop(&mut cx, SymBinOp::Mul, left, right);
313
314 assert!(!polynomial_identity(&expression, &expression));
315 }
316
317 #[test]
318 fn polynomial_identity_skips_irrelevant_and_unsupported_shapes() {
319 let mut cx = SymCx::new();
320 let x = SymExpr::var(&mut cx, "x");
321 let y = SymExpr::var(&mut cx, "y");
322 let single_product = SymExpr::binop(&mut cx, SymBinOp::Mul, x.clone(), y.clone());
323 assert!(!polynomial_identity(&single_product, &single_product));
324
325 let denominator = SymExpr::var(&mut cx, "denominator");
326 let quotient = SymExpr::binop(&mut cx, SymBinOp::UDiv, x.clone(), denominator);
327 let sum = SymExpr::binop(&mut cx, SymBinOp::Add, quotient, y);
328 let unsupported = SymExpr::binop(&mut cx, SymBinOp::Mul, sum, x);
329 assert!(!polynomial_identity(&unsupported, &unsupported));
330 }
331
332 #[test]
333 fn polynomial_analysis_stops_at_input_node_limit() {
334 let mut cx = SymCx::new();
335 let one = SymExpr::one(&mut cx);
336 let mut expression = SymExpr::var(&mut cx, "source");
337 for _ in 0..MAX_LOCAL_ANALYSIS_NODES {
338 expression = SymExpr::from_kind(
339 &mut cx,
340 SymExprKind::BinOp(SymBinOp::Add, expression, one.clone()),
341 );
342 }
343 let factor = SymExpr::var(&mut cx, "factor");
344 let product = SymExpr::from_kind(
345 &mut cx,
346 SymExprKind::BinOp(SymBinOp::Mul, expression.clone(), factor),
347 );
348
349 assert!(!polynomial_normalization_can_help(&product));
350 assert!(Polynomial::from_expr(&expression).is_none());
351 }
352}