Skip to main content

foundry_evm_symbolic/runtime/solver/normalize/
polynomial.rs

1//! Bounded sparse-polynomial identity reasoning over EVM words.
2
3use 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
87// Keep distributive expansion predictably bounded. The motivating accounting identity needs two
88// terms with two factors; these limits leave ample room for ordinary identities without allowing
89// adversarial expressions to explode.
90const MAX_POLYNOMIAL_TERMS: usize = 32;
91const MAX_MONOMIAL_FACTORS: usize = 8;
92const MAX_POLYNOMIAL_PRODUCTS: usize = 256;
93
94type Monomial = Vec<SymExpr>;
95
96/// A sparse polynomial over the EVM word ring Z/(2^256).
97///
98/// Addition, subtraction, and multiplication of EVM words obey the ring laws even when they
99/// wrap. Canonicalizing small expressions here lets the solver recognize nonlinear algebraic
100/// identities without replacing bit-vector semantics with unbounded integer arithmetic.
101#[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}