refactor: Improve Spartan SNARK polynomial computations and evaluations

This gathers the changes to the pre-processing SNARK excerpted from the Supernova implementation in PR microsoft#283.
The main changes extracted here are the introduction of the masked eq polynomial (and its use fixing [this issue](https://hackmd.io/@adr1anh/Sy08YaVBa)), additional sum-check tooling, removal of two calls to evaluation argument.

Main reference Arecibo PRs:
- argumentcomputer/arecibo#106
- argumentcomputer/arecibo#131
- argumentcomputer/arecibo#174
- argumentcomputer/arecibo#182

- Enhancement of polynomial related code in Spartan protocol including new polynomial types, modified existing types, better evaluation methods, and improved polynomial operations.
- Introduction of `squares` function and change in the generation of `t_pow` in Spartan.
- Addition of a new polynomial type through `MaskedEqPolynomial` with methods for its creation and evaluation.
- Enhancements in `UniPoly` struct by addition of `PartialEq`, and `Eq` traits.
- Improvements in `snark.rs` for proving and verifying batch evaluation claims, leveraging `gamma` and `rho` for random linear combinations, and optimizing various variable computations.
- Updates in `multilinear.rs` with refactoring, optimization, error handling, and supplementing unit tests.
- Refactor in `spartan/mod.rs` with import updates, function overhauls, struct visibility changes, and asynchronous operations for efficient calculations.
- Additions and amendments in `sumcheck.rs` for batch verification, handling of vectors of claims, handling of cubic bounds with additive terms, visibility adjustments, and typo fixes.
- Modifications in `eq.rs` including a debug derive for `EqPolynomial`, enhanced visibility of `r` vector, provision of `evals_from_points` for enhanced evaluation, and addition of `FromIterator` implementation.

Co-authored-by: porcuquine <[email protected]>
Co-authored-by: Matej Penciak <[email protected]>
Co-authored-by: Adrian Hamelink <[email protected]>

src/spartan/mod.rs

@@ -5,24 +5,29 @@
 //! We also provide direct.rs that allows proving a step circuit directly with either of the two SNARKs.
 //!
 //! In polynomial.rs we also provide foundational types and functions for manipulating multilinear polynomials.
+pub mod direct;
 #[macro_use]
 mod macros;
-pub mod direct;
 pub(crate) mod math;
 pub mod polys;
 pub mod ppsnark;
 pub mod snark;
 mod sumcheck;
 
-use crate::{traits::Engine, Commitment};
+use crate::{
+  r1cs::{R1CSShape, SparseMatrix},
+  traits::Engine,
+  Commitment,
+};
 use ff::Field;
 use itertools::Itertools as _;
 use polys::multilinear::SparsePolynomial;
 use rayon::{iter::IntoParallelRefIterator, prelude::*};
 
+// Creates a vector of the first `n` powers of `s`.
 fn powers<E: Engine>(s: &E::Scalar, n: usize) -> Vec<E::Scalar> {
   assert!(n >= 1);
-  let mut powers = Vec::new();
+  let mut powers = Vec::with_capacity(n);
   powers.push(E::Scalar::ONE);
   for i in 1..n {
     powers.push(powers[i - 1] * s);
@@ -31,35 +36,60 @@ fn powers<E: Engine>(s: &E::Scalar, n: usize) -> Vec<E::Scalar> {
 }
 
 /// A type that holds a witness to a polynomial evaluation instance
-pub struct PolyEvalWitness<E: Engine> {
+struct PolyEvalWitness<E: Engine> {
   p: Vec<E::Scalar>, // polynomial
 }
 
 impl<E: Engine> PolyEvalWitness<E> {
-  fn pad(mut W: Vec<PolyEvalWitness<E>>) -> Vec<PolyEvalWitness<E>> {
-    // determine the maximum size
-    if let Some(n) = W.iter().map(|w| w.p.len()).max() {
-      W.iter_mut().for_each(|w| {
-        w.p.resize(n, E::Scalar::ZERO);
-      });
-      W
-    } else {
-      Vec::new()
-    }
-  }
+  /// Given [Pᵢ] and s, compute P = ∑ᵢ sⁱ⋅Pᵢ
+  ///
+  /// # Details
+  ///
+  /// We allow the input polynomials to have different sizes, and interpret smaller ones as
+  /// being padded with 0 to the maximum size of all polynomials.
+  fn batch_diff_size(W: Vec<PolyEvalWitness<E>>, s: E::Scalar) -> PolyEvalWitness<E> {
+    let powers = powers::<E>(&s, W.len());
+
+    let size_max = W.iter().map(|w| w.p.len()).max().unwrap();
+    // Scale the input polynomials by the power of s
+    let p = W
+      .into_par_iter()
+      .zip_eq(powers.par_iter())
+      .map(|(mut w, s)| {
+        if *s != E::Scalar::ONE {
+          w.p.par_iter_mut().for_each(|e| *e *= s);
+        }
+        w.p
+      })
+      .reduce(
+        || vec![E::Scalar::ZERO; size_max],
+        |left, right| {
+          // Sum into the largest polynomial
+          let (mut big, small) = if left.len() > right.len() {
+            (left, right)
+          } else {
+            (right, left)
+          };
+
+          big
+            .par_iter_mut()
+            .zip(small.par_iter())
+            .for_each(|(b, s)| *b += s);
+
+          big
+        },
+      );
 
-  fn weighted_sum(W: &[PolyEvalWitness<E>], s: &[E::Scalar]) -> PolyEvalWitness<E> {
-    assert_eq!(W.len(), s.len());
-    let mut p = vec![E::Scalar::ZERO; W[0].p.len()];
-    for i in 0..W.len() {
-      for j in 0..W[i].p.len() {
-        p[j] += W[i].p[j] * s[i]
-      }
-    }
     PolyEvalWitness { p }
   }
 
-  // This method panics unless all vectors in p_vec are of the same length
+  /// Given a set of polynomials \[Pᵢ\] and a scalar `s`, this method computes the weighted sum
+  /// of the polynomials, where each polynomial Pᵢ is scaled by sⁱ. The method handles polynomials
+  /// of different sizes by padding smaller ones with zeroes up to the size of the largest polynomial.
+  ///
+  /// # Panics
+  ///
+  /// This method panics if the polynomials in `p_vec` are not all of the same length.
   fn batch(p_vec: &[&Vec<E::Scalar>], s: &E::Scalar) -> PolyEvalWitness<E> {
     p_vec
       .iter()
@@ -69,7 +99,7 @@ impl<E: Engine> PolyEvalWitness<E> {
 
     let p = zip_with!(par_iter, (p_vec, powers_of_s), |v, weight| {
       // compute the weighted sum for each vector
-      v.iter().map(|&x| x * weight).collect::<Vec<E::Scalar>>()
+      v.iter().map(|&x| x * *weight).collect::<Vec<E::Scalar>>()
     })
     .reduce(
       || vec![E::Scalar::ZERO; p_vec[0].len()],
@@ -84,25 +114,54 @@ impl<E: Engine> PolyEvalWitness<E> {
 }
 
 /// A type that holds a polynomial evaluation instance
-pub struct PolyEvalInstance<E: Engine> {
+struct PolyEvalInstance<E: Engine> {
   c: Commitment<E>,  // commitment to the polynomial
   x: Vec<E::Scalar>, // evaluation point
   e: E::Scalar,      // claimed evaluation
 }
 
 impl<E: Engine> PolyEvalInstance<E> {
-  fn pad(U: Vec<PolyEvalInstance<E>>) -> Vec<PolyEvalInstance<E>> {
-    // determine the maximum size
-    if let Some(ell) = U.iter().map(|u| u.x.len()).max() {
-      U.into_iter()
-        .map(|mut u| {
-          let mut x = vec![E::Scalar::ZERO; ell - u.x.len()];
-          x.append(&mut u.x);
-          PolyEvalInstance { x, ..u }
-        })
-        .collect()
-    } else {
-      Vec::new()
+  fn batch_diff_size(
+    c_vec: &[Commitment<E>],
+    e_vec: &[E::Scalar],
+    num_vars: &[usize],
+    x: Vec<E::Scalar>,
+    s: E::Scalar,
+  ) -> PolyEvalInstance<E> {
+    let num_instances = num_vars.len();
+    assert_eq!(c_vec.len(), num_instances);
+    assert_eq!(e_vec.len(), num_instances);
+
+    let num_vars_max = x.len();
+    let powers: Vec<E::Scalar> = powers::<E>(&s, num_instances);
+    // Rescale evaluations by the first Lagrange polynomial,
+    // so that we can check its evaluation against x
+    let evals_scaled = zip_with!(iter, (e_vec, num_vars), |eval, num_rounds| {
+      // x_lo = [ x[0]   , ..., x[n-nᵢ-1] ]
+      // x_hi = [ x[n-nᵢ], ..., x[n]      ]
+      let (r_lo, _r_hi) = x.split_at(num_vars_max - num_rounds);
+      // Compute L₀(x_lo)
+      let lagrange_eval = r_lo
+        .iter()
+        .map(|r| E::Scalar::ONE - r)
+        .product::<E::Scalar>();
+
+      // vᵢ = L₀(x_lo)⋅Pᵢ(x_hi)
+      lagrange_eval * eval
+    })
+    .collect::<Vec<_>>();
+
+    // C = ∑ᵢ γⁱ⋅Cᵢ
+    let comm_joint = zip_with!(iter, (c_vec, powers), |c, g_i| *c * *g_i)
+      .fold(Commitment::<E>::default(), |acc, item| acc + item);
+
+    // v = ∑ᵢ γⁱ⋅vᵢ
+    let eval_joint = zip_with!((evals_scaled.into_iter(), powers.iter()), |e, g_i| e * g_i).sum();
+
+    PolyEvalInstance {
+      c: comm_joint,
+      x,
+      e: eval_joint,
     }
   }
 
@@ -112,8 +171,13 @@ impl<E: Engine> PolyEvalInstance<E> {
     e_vec: &[E::Scalar],
     s: &E::Scalar,
   ) -> PolyEvalInstance<E> {
-    let powers_of_s = powers::<E>(s, c_vec.len());
+    let num_instances = c_vec.len();
+    assert_eq!(e_vec.len(), num_instances);
+
+    let powers_of_s = powers::<E>(s, num_instances);
+    // Weighted sum of evaluations
     let e = zip_with!(par_iter, (e_vec, powers_of_s), |e, p| *e * p).sum();
+    // Weighted sum of commitments
     let c = zip_with!(par_iter, (c_vec, powers_of_s), |c, p| *c * *p)
       .reduce(Commitment::<E>::default, |acc, item| acc + item);
 
@@ -124,3 +188,43 @@ impl<E: Engine> PolyEvalInstance<E> {
     }
   }
 }
+
+/// Bounds "row" variables of (A, B, C) matrices viewed as 2d multilinear polynomials
+fn compute_eval_table_sparse<E: Engine>(
+  S: &R1CSShape<E>,
+  rx: &[E::Scalar],
+) -> (Vec<E::Scalar>, Vec<E::Scalar>, Vec<E::Scalar>) {
+  assert_eq!(rx.len(), S.num_cons);
+
+  let inner = |M: &SparseMatrix<E::Scalar>, M_evals: &mut Vec<E::Scalar>| {
+    for (row_idx, ptrs) in M.indptr.windows(2).enumerate() {
+      for (val, col_idx) in M.get_row_unchecked(ptrs.try_into().unwrap()) {
+        M_evals[*col_idx] += rx[row_idx] * val;
+      }
+    }
+  };
+
+  let (A_evals, (B_evals, C_evals)) = rayon::join(
+    || {
+      let mut A_evals: Vec<E::Scalar> = vec![E::Scalar::ZERO; 2 * S.num_vars];
+      inner(&S.A, &mut A_evals);
+      A_evals
+    },
+    || {
+      rayon::join(
+        || {
+          let mut B_evals: Vec<E::Scalar> = vec![E::Scalar::ZERO; 2 * S.num_vars];
+          inner(&S.B, &mut B_evals);
+          B_evals
+        },
+        || {
+          let mut C_evals: Vec<E::Scalar> = vec![E::Scalar::ZERO; 2 * S.num_vars];
+          inner(&S.C, &mut C_evals);
+          C_evals
+        },
+      )
+    },
+  );
+
+  (A_evals, B_evals, C_evals)
+}

src/spartan/polys/eq.rs

@@ -14,8 +14,9 @@ use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, Parall
 /// This polynomial evaluates to 1 if every component $x_i$ equals its corresponding $e_i$, and 0 otherwise.
 ///
 /// For instance, for e = 6 (with a binary representation of 0b110), the vector r would be [1, 1, 0].
+#[derive(Debug)]
 pub struct EqPolynomial<Scalar: PrimeField> {
-  r: Vec<Scalar>,
+  pub(in crate::spartan::polys) r: Vec<Scalar>,
 }
 
 impl<Scalar: PrimeField> EqPolynomial<Scalar> {
@@ -43,12 +44,20 @@ impl<Scalar: PrimeField> EqPolynomial<Scalar> {
   ///
   /// Returns a vector of Scalars, each corresponding to the polynomial evaluation at a specific point.
   pub fn evals(&self) -> Vec<Scalar> {
-    let ell = self.r.len();
+    Self::evals_from_points(&self.r)
+  }
+
+  /// Evaluates the `EqPolynomial` from the `2^|r|` points in its domain, without creating an intermediate polynomial
+  /// representation.
+  ///
+  /// Returns a vector of Scalars, each corresponding to the polynomial evaluation at a specific point.
+  pub fn evals_from_points(r: &[Scalar]) -> Vec<Scalar> {
+    let ell = r.len();
     let mut evals: Vec<Scalar> = vec![Scalar::ZERO; (2_usize).pow(ell as u32)];
     let mut size = 1;
     evals[0] = Scalar::ONE;
 
-    for r in self.r.iter().rev() {
+    for r in r.iter().rev() {
       let (evals_left, evals_right) = evals.split_at_mut(size);
       let (evals_right, _) = evals_right.split_at_mut(size);
 
@@ -64,6 +73,13 @@ impl<Scalar: PrimeField> EqPolynomial<Scalar> {
   }
 }
 
+impl<Scalar: PrimeField> FromIterator<Scalar> for EqPolynomial<Scalar> {
+  fn from_iter<I: IntoIterator<Item = Scalar>>(iter: I) -> Self {
+    let r: Vec<_> = iter.into_iter().collect();
+    EqPolynomial { r }
+  }
+}
+
 #[cfg(test)]
 mod tests {
   use crate::provider;