Rework batches in ppsnark (#305)

* refactor: Streamline matrix multiply & some zip_with instances

- Reworked vector and matrix multiplication operations for improved memory use in `sparse.rs`.
- Modified `batch_eval_prove` function in `snark.rs` for more cleaner assertion checks related to polynomial sizes.
- Enhanced `InnerSumcheckInstance` within `setup` function for parallel computation in `ppsnark.rs`.

* refactor: Streamline transcript absorption in Spartan batched code

- Optimized the transcript insertion process in `batched.rs` by replacing loop-absorbing RelaxedR1CSInstances with a single absorb statement

* refactor: Refactor 'claims_outer' and related functions in BatchedRelaxedR1CSSNARK

- Refactored the structure of `claims_outer` in `BatchedRelaxedR1CSSNARK` from a tuple of vectors to a vector of tuples, streamlining its usage and unpacking.
- Adjusted the `prove` and `verify` functions' code in accordance with the new `claims_outer` implementation.
- Enhanced the extraction and manipulation of `ABCE_evals` in the `verify` function resulting in simpler and more efficient usage.

* fix: simplify inner final claim

* refactor: Optimize Spartan batch evaluation and polynomial computation

- Refined tau polynomial definitions using `PowPolynomial::squares` function to improve efficiency.
- uses in verify the same efficient method as exists in proving.

* fix: streamline usage of PowPolynomial

* refactor: Enhance efficiency of generating power polynomials

- Introduced `#[must_use]` annotation to the `squares` function in `power.rs` to ensure the returned result is handled appropriately.
- avoid re-computing powers of tau over and over in batched_ppsnark

* fix: one forgotten instance of the poly_X refactor

* refactor: slightly simplify padding

* fix: remove misleading comment

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..])
  };