sca: building a gate polynomial from the bottom up does not seem to work

ekiwi · ekiwi · commit 3fdb066a900d · 2026-02-24T15:21:33.000-05:00
diff --git a/patronus-sca/Cargo.toml b/patronus-sca/Cargo.toml
@@ -14,5 +14,5 @@ patronus = { path = "../patronus" }
 baa = {version = "0.17.1", features = ["bigint"]}
 rustc-hash.workspace = true
 smallvec.workspace = true
-polysub = "0.2.4"
+polysub = "0.2.5"
 bit-set = "0.8.0"
diff --git a/patronus-sca/src/forward.rs b/patronus-sca/src/forward.rs
@@ -7,25 +7,73 @@
 
 use crate::{Poly, expr_to_var, extract_bit};
 use baa::{BitVecOps, BitVecValueRef};
-use patronus::expr::{Context, Expr, ExprRef, ExprSet, SparseExprSet, TypeCheck, traversal};
-use polysub::Coef;
-use rustc_hash::FxHashSet;
+use patronus::expr::{
+    Context, Expr, ExprRef, ExprSet, ForEachChild, SerializableIrNode, SparseExprMap,
+    SparseExprSet, TypeCheck, count_expr_uses, traversal,
+};
+use polysub::{Coef, Mod};
+use rustc_hash::{FxHashMap, FxHashSet};
+
+#[derive(Debug, Copy, Clone)]
+pub enum BuildPolyMode {
+    Arithmetic,
+    Gates(Mod),
+}
 
 /// Returns a polynomial representation of the expression + all input expressions if possible.
 /// Returns `None` if the conversion fails.
 pub fn build_bottom_up_poly(
     ctx: &mut Context,
     inputs: &FxHashSet<ExprRef>,
     e: ExprRef,
+    mode: BuildPolyMode,
 ) -> Option<Poly> {
-    traversal::bottom_up_mut(ctx, e, |ctx, e, c| to_poly(ctx, inputs, e, c)).map(|(p, _)| p)
+    // we use a custom traversal that caches polynomials until they are no longer used
+    let mut uses = count_expr_uses(ctx, vec![e]);
+    let mut todo = vec![e];
+    let mut result: FxHashMap<ExprRef, Option<(Poly, bool)>> = FxHashMap::default();
+    let mut child_vec = Vec::with_capacity(4);
+
+    while let Some(e) = todo.pop() {
+        assert!(!result.contains_key(&e));
+
+        let expr = &ctx[e];
+        // find children that are not available yet.
+        debug_assert!(child_vec.is_empty());
+        expr.collect_children(&mut child_vec);
+        let all_available = child_vec.iter().all(|c| result.contains_key(c));
+
+        if all_available {
+            let child_results: Vec<Option<&(Poly, bool)>> =
+                child_vec.iter().map(|c| result[c].as_ref()).collect();
+            let r = to_poly(ctx, inputs, mode, e, &child_results);
+            result.insert(e, r);
+            for child in child_vec.drain(..) {
+                let old_use = uses[child];
+                if old_use == 1 {
+                    result.remove(&child);
+                }
+                uses[child] = old_use - 1;
+            }
+        } else {
+            todo.push(e);
+            for child in child_vec.drain(..) {
+                if !result.contains_key(&child) {
+                    todo.push(child);
+                }
+            }
+        }
+    }
+
+    result[&e].as_ref().map(|(p, _)| p.clone())
 }
 
 fn to_poly(
     ctx: &mut Context,
     inputs: &FxHashSet<ExprRef>,
+    mode: BuildPolyMode,
     e: ExprRef,
-    children: &[Option<(Poly, bool)>],
+    children: &[Option<&(Poly, bool)>],
 ) -> Option<(Poly, bool)> {
     // have we given up yet?
     if children.iter().any(|c| c.is_none()) {
@@ -37,6 +85,89 @@ fn to_poly(
         return Some((poly_for_bv_expr(ctx, e), false));
     }
 
+    match mode {
+        BuildPolyMode::Arithmetic => to_poly_arithmetic(ctx, inputs, e, children),
+        BuildPolyMode::Gates(m) => Some((to_poly_gate(ctx, inputs, m, e, children), false)),
+    }
+}
+
+/// For bit-level polynomials from gates, the normal arithmetic rules and overflow checking do not
+/// apply. Instead, we use the modulo coefficient of the top-level expression.
+fn to_poly_gate(
+    ctx: &mut Context,
+    inputs: &FxHashSet<ExprRef>,
+    m: Mod,
+    e: ExprRef,
+    children: &[Option<&(Poly, bool)>],
+) -> Poly {
+    debug_assert!(
+        children
+            .iter()
+            .all(|c| c.map(|(_, ov)| !*ov).unwrap_or(true))
+    );
+
+    match (ctx[e].clone(), children) {
+        (Expr::BVSymbol { .. }, _) => unreachable!("all symbols should be in inputs"),
+        (Expr::BVLiteral(value), _) => {
+            let mut r = poly_for_bv_literal(value.get(ctx));
+            r.change_mod(m);
+            r
+        }
+        (Expr::BVSlice { e, hi, lo }, _) if hi == lo && inputs.contains(&e) => {
+            // special case: bit_slice of an input
+            let var = expr_to_var(extract_bit(ctx, e, hi));
+            Poly::from_monoms(m, [(Coef::from_i64(1, m), vec![var].into())].into_iter())
+        }
+        (Expr::BVConcat(_, be, w), [Some((a, _)), Some((b, _))]) => {
+            // left shift a
+            let shift_by = be.get_bv_type(ctx).unwrap();
+            let shift_coef = Coef::pow2(shift_by, m);
+            let mut r: Poly = a.clone();
+            r.scale(&shift_coef);
+            r.add_assign(b);
+            r
+        }
+        (Expr::BVOr(_, _, 1), [Some((a, _)), Some((b, _))]) => {
+            // a + b - ab
+            let mut r = a.mul(b);
+            r.scale(&Coef::from_i64(-1, a.get_mod()));
+            r.add_assign(a);
+            r.add_assign(b);
+            r
+        }
+        (Expr::BVXor(_, _, 1), [Some((a, _)), Some((b, _))]) => {
+            // a + b - 2ab
+            let mut r = a.mul(b);
+            let minus_2 = Coef::from_i64(-2, a.get_mod());
+            r.scale(&minus_2);
+            r.add_assign(a);
+            r.add_assign(b);
+            r
+        }
+        (Expr::BVAnd(_, _, 1), [Some((a, _)), Some((b, _))]) => {
+            // ab
+            a.clone().mul(b)
+        }
+        (Expr::BVNot(_, 1), [Some((a, _))]) => {
+            // 1 - a
+            let one = Poly::from_monoms(m, [(Coef::from_i64(1, m), vec![].into())].into_iter());
+            let mut r = a.clone();
+            r.scale(&Coef::from_i64(-1, a.get_mod()));
+            r.add_assign(&one);
+            r
+        }
+        (other, cs) => todo!("{other:?}: {cs:?}"),
+    }
+}
+
+/// When building a polynomial over an arithmetic circuit, we are tracking overflow and
+/// determining the modulo coefficient from the bit-widths.
+fn to_poly_arithmetic(
+    ctx: &mut Context,
+    _inputs: &FxHashSet<ExprRef>,
+    e: ExprRef,
+    children: &[Option<&(Poly, bool)>],
+) -> Option<(Poly, bool)> {
     match (ctx[e].clone(), children) {
         (Expr::BVSymbol { .. }, _) => unreachable!("all symbols should be in inputs"),
         (Expr::BVLiteral(value), _) => Some((poly_for_bv_literal(value.get(ctx)), false)),
diff --git a/patronus-sca/src/lib.rs b/patronus-sca/src/lib.rs
@@ -5,8 +5,8 @@
 mod forward;
 mod rewrite;
 
-use crate::forward::{build_bottom_up_poly, poly_for_bv_expr};
-use crate::rewrite::{backwards_sub, build_gate_polynomial};
+use crate::forward::{BuildPolyMode, build_bottom_up_poly, poly_for_bv_expr};
+use crate::rewrite::backwards_sub;
 use baa::{BitVecOps, BitVecValue, BitVecValueRef};
 use patronus::expr::*;
 use polysub::{Coef, Term, VarIndex};
@@ -32,11 +32,13 @@ pub fn verify_word_level_equality(ctx: &mut Context, p: ScaEqualityProblem) -> S
     let inputs = find_symbols(ctx, p.word_level);
 
     // create a reference polynomial from the word level side
-    let mut word_poly = match build_bottom_up_poly(ctx, &inputs, p.word_level) {
-        None => return ScaVerifyResult::Unknown,
-        Some(p) => p,
-    };
+    let mut word_poly =
+        match build_bottom_up_poly(ctx, &inputs, p.word_level, BuildPolyMode::Arithmetic) {
+            None => return ScaVerifyResult::Unknown,
+            Some(p) => p,
+        };
     println!("word-level polynomial: {word_poly}");
+    println!("word-level bits: {}", word_poly.get_mod().bits());
 
     // collect all (bit-level) input variables
     let input_vars: FxHashSet<VarIndex> = inputs
@@ -50,7 +52,11 @@ pub fn verify_word_level_equality(ctx: &mut Context, p: ScaEqualityProblem) -> S
         })
         .collect();
 
-    //let gate_poly = build_gate_polynomial(ctx, &input_vars, word_poly.get_mod(), p.gate_level);
+    // TODO: how do we calculate a correct polynomial from the gate level?
+    // let gate_poly = build_bottom_up_poly(ctx, &inputs, p.gate_level, BuildPolyMode::Gates(word_poly.get_mod()));
+    // if let Some(gate_poly) = gate_poly {
+    //     println!("GATE POLY: {gate_poly}");
+    // }
 
     // the actual reference polynomial needs to contain the output bits as well
     let output_poly = poly_for_bv_expr(ctx, p.word_level);
@@ -395,7 +401,9 @@ mod tests {
             let cs = find_sca_simplification_candidates(&ctx, e);
             for c in cs {
                 let inputs = find_symbols(&ctx, c.word_level);
-                if let Some(word_poly) = build_bottom_up_poly(&mut ctx, &inputs, c.word_level) {
+                if let Some(word_poly) =
+                    build_bottom_up_poly(&mut ctx, &inputs, c.word_level, BuildPolyMode::Arithmetic)
+                {
                     println!("{filename}:\n{}\n", PrettyPoly::n(&ctx, &word_poly));
                 } else {
                     println!(
diff --git a/patronus-sca/src/rewrite.rs b/patronus-sca/src/rewrite.rs
@@ -317,67 +317,6 @@ fn try_exhaustive(
     min_with_lowest_use
 }
 
-/// tries to build the gate-level polynomial from the bottom up.
-pub fn build_gate_polynomial(
-    ctx: &Context,
-    input_vars: &FxHashSet<VarIndex>,
-    m: Mod,
-    e: ExprRef,
-) -> Poly {
-    traversal::bottom_up(ctx, e, |ctx, gate, children: &[Poly]| {
-        let var = expr_to_var(gate);
-
-        // if this is an input, we just want to return the variable
-        if input_vars.contains(&var) {
-            return Poly::from_monoms(m, [(Coef::from_i64(1, m), vec![var].into())].into_iter());
-        }
-
-        match (ctx[gate].clone(), children) {
-            (Expr::BVOr(_, _, 1), [a, b]) => {
-                // a + b - ab
-                let mut r = a.mul(b);
-                r.scale(&Coef::from_i64(-1, m));
-                r.add_assign(a);
-                r.add_assign(b);
-                r
-            }
-            (Expr::BVXor(_, _, 1), [a, b]) => {
-                // a + b - 2ab
-                let mut r = a.mul(b);
-                r.scale(&Coef::from_i64(-2, m));
-                r.add_assign(a);
-                r.add_assign(b);
-                r
-            }
-            (Expr::BVAnd(_, _, 1), [a, b]) => {
-                // ab
-                a.mul(b)
-            }
-            (Expr::BVNot(_, 1), [a]) => {
-                // 1 - a
-                let one = Poly::from_monoms(m, [(Coef::from_i64(1, m), vec![].into())].into_iter());
-                let mut r = a.clone();
-                r.scale(&Coef::from_i64(-1, m));
-                r.add_assign(&one);
-                r
-            }
-            (Expr::BVSlice { .. }, [_]) => {
-                todo!("should not get here!")
-            }
-            (Expr::BVLiteral(value), _) => {
-                let value = value.get(ctx);
-                debug_assert_eq!(value.width(), 1);
-                if value.is_true() {
-                    Poly::from_monoms(m, [(Coef::from_i64(1, m), vec![].into())].into_iter())
-                } else {
-                    Poly::from_monoms(m, [(Coef::from_i64(0, m), vec![].into())].into_iter())
-                }
-            }
-            other => todo!("add support for {other:?}"),
-        }
-    })
-}
-
 /// Calculates for each expression which root depends on it.
 fn analyze_uses(ctx: &Context, roots: &[ExprRef]) -> impl ExprMap<BitSet> {
     let mut out = DenseExprMetaData::<BitSet>::default();
diff --git a/patronus/src/expr/analysis.rs b/patronus/src/expr/analysis.rs
@@ -9,7 +9,7 @@ use rustc_hash::FxHashSet;
 pub type UseCountInt = u16;
 
 /// Counts how often expressions in the DAGs characterized by the provided roots are used.
-pub fn count_expr_uses(ctx: &Context, roots: Vec<ExprRef>) -> impl ExprMap<UseCountInt> {
+pub fn count_expr_uses(ctx: &Context, roots: Vec<ExprRef>) -> impl ExprMap<UseCountInt> + use<> {
     let mut use_count = SparseExprMap::default();
     let mut todo = roots;
 
