Rework batches in ppsnark

src/provider/util/mod.rs

@@ -68,7 +68,7 @@ pub mod test_utils {
       .collect::<Vec<_>>();
 
     // Calculation evaluation of point over polynomial.
-    let eval = MultilinearPolynomial::evaluate_with(poly.evaluations(), &point);
+    let eval = poly.evaluate(&point);
 
     (poly, point, eval)
   }

src/r1cs/sparse.rs

@@ -135,6 +135,12 @@ impl<F: PrimeField> SparseMatrix<F> {
     name = "SparseMatrix::multiply_vec_unchecked"
   )]
   pub fn multiply_vec_unchecked(&self, vector: &[F]) -> Vec<F> {
+    let mut sink: Vec<F> = Vec::with_capacity(self.indptr.len() - 1);
+    self.multiply_vec_into_unchecked(vector, &mut sink);
+    sink
+  }
+
+  pub fn multiply_vec_into_unchecked(&self, vector: &[F], sink: &mut Vec<F>) {
     self
       .indptr
       .par_windows(2)
@@ -144,7 +150,7 @@ impl<F: PrimeField> SparseMatrix<F> {
           .map(|(val, col_idx)| *val * vector[*col_idx])
           .sum()
       })
-      .collect()
+      .collect_into_vec(sink);
   }
 
   /// Multiply by a witness representing a dense vector; uses rayon to parallelize.

src/spartan/batched.rs

@@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize};
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
 use rayon::prelude::*;
-use std::sync::Arc;
+use std::{iter, sync::Arc};
 
 use super::{
   compute_eval_table_sparse,
@@ -45,7 +45,7 @@ use crate::{
 pub struct BatchedRelaxedR1CSSNARK<E: Engine, EE: EvaluationEngineTrait<E>> {
   sc_proof_outer: SumcheckProof<E>,
   // Claims ([Azᵢ(τᵢ)], [Bzᵢ(τᵢ)], [Czᵢ(τᵢ)])
-  claims_outer: (Vec<E::Scalar>, Vec<E::Scalar>, Vec<E::Scalar>),
+  claims_outer: Vec<(E::Scalar, E::Scalar, E::Scalar)>,
   // [Eᵢ(r_x)]
   evals_E: Vec<E::Scalar>,
   sc_proof_inner: SumcheckProof<E>,
@@ -134,17 +134,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
   ) -> Result<Self, NovaError> {
     let num_instances = U.len();
     // Pad shapes and ensure their sizes are correct
-    let S = S
-      .iter()
-      .map(|s| {
-        let s = s.pad();
-        if s.is_regular_shape() {
-          Ok(s)
-        } else {
-          Err(NovaError::InternalError)
-        }
-      })
-      .collect::<Result<Vec<_>, _>>()?;
+    let S = S.iter().map(|s| s.pad()).collect::<Vec<_>>();
 
     // Pad (W,E) for each instance
     let W = zip_with!(iter, (W, S), |w, s| w.pad(s)).collect::<Vec<RelaxedR1CSWitness<E>>>();
@@ -156,9 +146,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       let num_instances_field = E::Scalar::from(num_instances as u64);
       transcript.absorb(b"n", &num_instances_field);
     }
-    for u in U.iter() {
-      transcript.absorb(b"U", u);
-    }
+    transcript.absorb(b"U", &U);
 
     let (polys_W, polys_E): (Vec<_>, Vec<_>) = W.into_iter().map(|w| (w.W, w.E)).unzip();
 
@@ -371,12 +359,9 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       &batched_u.e,
     )?;
 
-    let (evals_Az, evals_Bz, evals_Cz): (Vec<_>, Vec<_>, Vec<_>) =
-      evals_Az_Bz_Cz.into_iter().multiunzip();
-
     Ok(Self {
       sc_proof_outer,
-      claims_outer: (evals_Az, evals_Bz, evals_Cz),
+      claims_outer: evals_Az_Bz_Cz,
       evals_E,
       sc_proof_inner,
       evals_W,
@@ -395,9 +380,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       let num_instances_field = E::Scalar::from(num_instances as u64);
       transcript.absorb(b"n", &num_instances_field);
     }
-    for u in U.iter() {
-      transcript.absorb(b"U", u);
-    }
+    transcript.absorb(b"U", &U);
 
     let num_instances = U.len();
 
@@ -410,14 +393,14 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     let num_rounds_y_max = *num_rounds_y.iter().max().unwrap();
 
     // Define τ polynomials of the appropriate size for each instance
-    let polys_tau = {
-      let tau = transcript.squeeze(b"t")?;
+    let tau = transcript.squeeze(b"t")?;
+    let all_taus = PowPolynomial::squares(&tau, num_rounds_x_max);
 
-      num_rounds_x
-        .iter()
-        .map(|&num_rounds| PowPolynomial::new(&tau, num_rounds))
-        .collect::<Vec<_>>()
-    };
+    let polys_tau = num_rounds_x
+      .iter()
+      .map(|&num_rounds_x| PowPolynomial::evals_with_powers(&all_taus, num_rounds_x))
+      .map(MultilinearPolynomial::new)
+      .collect::<Vec<_>>();
 
     // Sample challenge for random linear-combination of outer claims
     let outer_r = transcript.squeeze(b"out_r")?;
@@ -441,20 +424,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // Extract evaluations into a vector [(Azᵢ, Bzᵢ, Czᵢ, Eᵢ)]
     // TODO: This is a multizip, simplify
-    let ABCE_evals = zip_with!(
-      iter,
-      (
-        self.evals_E,
-        self.claims_outer.0,
-        self.claims_outer.1,
-        self.claims_outer.2
-      ),
-      |eval_E, claim_Az, claim_Bz, claim_Cz| (*claim_Az, *claim_Bz, *claim_Cz, *eval_E)
-    )
-    .collect::<Vec<_>>();
+    let ABCE_evals = || self.claims_outer.iter().zip_eq(self.evals_E.iter());
 
     // Add evaluations of Az, Bz, Cz, E to transcript
-    for (claim_Az, claim_Bz, claim_Cz, eval_E) in ABCE_evals.iter() {
+    for ((claim_Az, claim_Bz, claim_Cz), eval_E) in ABCE_evals() {
       transcript.absorb(
         b"claims_outer",
         &[*claim_Az, *claim_Bz, *claim_Cz, *eval_E].as_slice(),
@@ -467,15 +440,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // Compute expected claim for all instances ∑ᵢ rⁱ⋅τ(rₓ)⋅(Azᵢ⋅Bzᵢ − uᵢ⋅Czᵢ − Eᵢ)
     let claim_outer_final_expected = zip_with!(
-      (
-        ABCE_evals.iter().copied(),
-        U.iter(),
-        evals_tau,
-        outer_r_powers.iter()
-      ),
+      (ABCE_evals(), U.iter(), evals_tau, outer_r_powers.iter()),
       |ABCE_eval, u, eval_tau, r| {
-        let (claim_Az, claim_Bz, claim_Cz, eval_E) = ABCE_eval;
-        *r * eval_tau * (claim_Az * claim_Bz - u.u * claim_Cz - eval_E)
+        let ((claim_Az, claim_Bz, claim_Cz), eval_E) = ABCE_eval;
+        *r * eval_tau * (*claim_Az * claim_Bz - u.u * claim_Cz - eval_E)
       }
     )
     .sum::<E::Scalar>();
@@ -491,10 +459,11 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // Compute inner claims Mzᵢ = (Azᵢ + r⋅Bzᵢ + r²⋅Czᵢ),
     // which are batched by Sumcheck into one claim:  ∑ᵢ r³ⁱ⋅Mzᵢ
-    let claims_inner = ABCE_evals
-      .into_iter()
-      .map(|(claim_Az, claim_Bz, claim_Cz, _)| {
-        claim_Az + inner_r * claim_Bz + inner_r_square * claim_Cz
+    let claims_inner = self
+      .claims_outer
+      .iter()
+      .map(|(claim_Az, claim_Bz, claim_Cz)| {
+        *claim_Az + inner_r * claim_Bz + inner_r_square * claim_Cz
       })
       .collect::<Vec<_>>();
 
@@ -515,15 +484,17 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     let evals_Z = zip_with!(iter, (self.evals_W, U, r_y), |eval_W, U, r_y| {
       let eval_X = {
         // constant term
-        let mut poly_X = vec![(0, U.u)];
-        //remaining inputs
-        poly_X.extend(
-          U.X
+        let poly_X = iter::once((0, U.u))
+          .chain(
+            //remaining inputs
+            U.X
             .iter()
             .enumerate()
-            .map(|(i, x_i)| (i + 1, *x_i))
-            .collect::<Vec<(usize, E::Scalar)>>(),
-        );
+            // filter_map uses the sparsity of the polynomial, if irrelevant
+            // we should replace by UniPoly
+            .filter_map(|(i, x_i)| (!x_i.is_zero_vartime()).then_some((i + 1, *x_i))),
+          )
+          .collect();
         SparsePolynomial::new(r_y.len() - 1, poly_X).evaluate(&r_y[1..])
       };
       (E::Scalar::ONE - r_y[0]) * eval_W + r_y[0] * eval_X

src/spartan/batched_ppsnark.rs

@@ -177,29 +177,20 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     W: &[RelaxedR1CSWitness<E>],
   ) -> Result<Self, NovaError> {
     // Pad shapes so that num_vars = num_cons = Nᵢ and check the sizes are correct
-    let S = S
-      .par_iter()
-      .map(|s| {
-        let s = s.pad();
-        if s.is_regular_shape() {
-          Ok(s)
-        } else {
-          Err(NovaError::InternalError)
-        }
-      })
-      .collect::<Result<Vec<_>, _>>()?;
+    let S = S.par_iter().map(|s| s.pad()).collect::<Vec<_>>();
 
     // N[i] = max(|Aᵢ|+|Bᵢ|+|Cᵢ|, 2*num_varsᵢ, num_consᵢ)
-    let N = pk.S_repr.iter().map(|s| s.N).collect::<Vec<_>>();
-    assert!(N.iter().all(|&Ni| Ni.is_power_of_two()));
+    let Nis = pk.S_repr.iter().map(|s| s.N).collect::<Vec<_>>();
+    assert!(Nis.iter().all(|&Ni| Ni.is_power_of_two()));
+    let N_max = *Nis.iter().max().unwrap();
 
     let num_instances = U.len();
 
     // Pad [(Wᵢ,Eᵢ)] to the next power of 2 (not to Ni)
     let W = zip_with!(par_iter, (W, S), |w, s| w.pad(s)).collect::<Vec<RelaxedR1CSWitness<E>>>();
 
     // number of rounds of sum-check
-    let num_rounds_sc = N.iter().max().unwrap().log_2();
+    let num_rounds_sc = N_max.log_2();
 
     // Initialize transcript with vk || [Uᵢ]
     let mut transcript = E::TE::new(b"BatchedRelaxedR1CSSNARK");
@@ -208,12 +199,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       let num_instances_field = E::Scalar::from(num_instances as u64);
       transcript.absorb(b"n", &num_instances_field);
     }
-    for u in U.iter() {
-      transcript.absorb(b"U", u);
-    }
+    transcript.absorb(b"U", &U);
 
     // Append public inputs to Wᵢ: Zᵢ = [Wᵢ, uᵢ, Xᵢ]
-    let polys_Z = zip_with!(par_iter, (W, U, N), |W, U, Ni| {
+    let polys_Z = zip_with!(par_iter, (W, U, Nis), |W, U, Ni| {
       // poly_Z will be resized later, so we preallocate the correct capacity
       let mut poly_Z = Vec::with_capacity(*Ni);
       poly_Z.extend(W.W.iter().chain([&U.u]).chain(U.X.iter()));
@@ -249,21 +238,23 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // Compute eq(tau) for each instance in log2(Ni) variables
     let tau = transcript.squeeze(b"t")?;
-    let (polys_tau, coords_tau): (Vec<_>, Vec<_>) = N
-      .iter()
+    let all_taus = PowPolynomial::squares(&tau, N_max.log_2());
+
+    let (polys_tau, coords_tau): (Vec<_>, Vec<_>) = Nis
+      .par_iter()
       .map(|&N_i| {
         let log_Ni = N_i.log_2();
-        let poly = PowPolynomial::new(&tau, log_Ni);
-        let evals = poly.evals();
-        let coords = poly.coordinates();
+        let eqp: EqPolynomial<_> = all_taus[..log_Ni].iter().cloned().collect();
+        let evals = eqp.evals();
+        let coords = eqp.r;
         (evals, coords)
       })
       .unzip();
 
     // Pad [Az, Bz, Cz] to Ni
     polys_Az_Bz_Cz
       .par_iter_mut()
-      .zip_eq(N.par_iter())
+      .zip_eq(Nis.par_iter())
       .for_each(|(az_bz_cz, &Ni)| {
         az_bz_cz
           .par_iter_mut()
@@ -298,7 +289,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     // Pad Zᵢ, E to Nᵢ
     let polys_Z = polys_Z
       .into_par_iter()
-      .zip_eq(N.par_iter())
+      .zip_eq(Nis.par_iter())
       .map(|(mut poly_Z, &Ni)| {
         poly_Z.resize(Ni, E::Scalar::ZERO);
         poly_Z
@@ -309,7 +300,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     // but it makes it easier to handle the batching at the end.
     let polys_E = polys_E
       .into_par_iter()
-      .zip_eq(N.par_iter())
+      .zip_eq(Nis.par_iter())
       .map(|(mut poly_E, &Ni)| {
         poly_E.resize(Ni, E::Scalar::ZERO);
         poly_E
@@ -318,7 +309,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     let polys_W = polys_W
       .into_par_iter()
-      .zip_eq(N.par_iter())
+      .zip_eq(Nis.par_iter())
       .map(|(mut poly_W, &Ni)| {
         poly_W.resize(Ni, E::Scalar::ZERO);
         poly_W
@@ -330,7 +321,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     // L_col(i) = z(col(i)) for all i in [0..Nᵢ]
     let polys_L_row_col = zip_with!(
       par_iter,
-      (S, N, polys_Z, polys_tau),
+      (S, Nis, polys_Z, polys_tau),
       |S, Ni, poly_Z, poly_tau| {
         let mut L_row = vec![poly_tau[0]; *Ni]; // we place mem_row[0] since resized row is appended with 0s
         let mut L_col = vec![poly_Z[Ni - 1]; *Ni]; // we place mem_col[Ni-1] since resized col is appended with Ni-1
@@ -495,13 +486,12 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
       // Sample new random variable for eq polynomial
       let rho = transcript.squeeze(b"r")?;
-      let N_max = N.iter().max().unwrap();
       let all_rhos = PowPolynomial::squares(&rho, N_max.log_2());
 
       let instances = zip_with!(
         (
           pk.S_repr.par_iter(),
-          N.par_iter(),
+          Nis.par_iter(),
           polys_mem_oracles.par_iter(),
           mem_aux.into_par_iter()
         ),
@@ -765,9 +755,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       let num_instances_field = E::Scalar::from(num_instances as u64);
       transcript.absorb(b"n", &num_instances_field);
     }
-    for u in U.iter() {
-      transcript.absorb(b"U", u);
-    }
+    transcript.absorb(b"U", &U);
 
     // Decompress commitments
     let comms_Az_Bz_Cz = self
@@ -907,14 +895,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
         let num_rounds_i = rand_sc.len();
         let num_vars_log = num_vars.log_2();
 
-        let eq_rho = {
-          let rho_coords = PowPolynomial::new(&rho, num_rounds_i).coordinates();
-          EqPolynomial::new(rho_coords).evaluate(rand_sc)
-        };
+        let eq_rho = PowPolynomial::new(&rho, num_rounds_i).evaluate(rand_sc);
 
         let (eq_tau, eq_masked_tau) = {
-          let tau_coords = PowPolynomial::new(&tau, num_rounds_i).coordinates();
-          let eq_tau = EqPolynomial::new(tau_coords);
+          let eq_tau: EqPolynomial<_> = PowPolynomial::new(&tau, num_rounds_i).into();
 
           let eq_tau_at_rand = eq_tau.evaluate(rand_sc);
           let eq_masked_tau = MaskedEqPolynomial::new(&eq_tau, num_vars_log).evaluate(rand_sc);
@@ -941,13 +925,15 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
           let X = {
             // constant term
-            let mut poly_X = vec![(0, U.u)];
-            //remaining inputs
-            poly_X.extend(
-              (0..U.X.len())
-                .map(|i| (i + 1, U.X[i]))
-                .collect::<Vec<(usize, E::Scalar)>>(),
-            );
+            let poly_X = std::iter::once((0, U.u))
+              .chain(
+                //remaining inputs
+                (0..U.X.len())
+                // filter_map uses the sparsity of the polynomial, if irrelevant
+                // we should replace by UniPoly
+                .filter_map(|i| (!U.X[i].is_zero_vartime()).then_some((i + 1, U.X[i]))),
+              )
+              .collect();
             SparsePolynomial::new(num_vars_log, poly_X).evaluate(&rand_sc_unpad[1..])
           };