Remove large vector allocations (microsoft#144)

* remove large vector allocations

* add suggestions

src/lib.rs

@@ -27,12 +27,15 @@ pub mod traits;
 
 use once_cell::sync::OnceCell;
 
-use crate::bellpepper::{
-  r1cs::{NovaShape, NovaWitness},
-  shape_cs::ShapeCS,
-  solver::SatisfyingAssignment,
-};
 use crate::digest::{DigestComputer, SimpleDigestible};
+use crate::{
+  bellpepper::{
+    r1cs::{NovaShape, NovaWitness},
+    shape_cs::ShapeCS,
+    solver::SatisfyingAssignment,
+  },
+  r1cs::R1CSResult,
+};
 use bellpepper_core::{ConstraintSystem, SynthesisError};
 use circuit::{NovaAugmentedCircuit, NovaAugmentedCircuitInputs, NovaAugmentedCircuitParams};
 use constants::{BN_LIMB_WIDTH, BN_N_LIMBS, NUM_FE_WITHOUT_IO_FOR_CRHF, NUM_HASH_BITS};
@@ -229,6 +232,21 @@ where
   }
 }
 
+/// A resource buffer for [`RecursiveSNARK`] for storing scratch values that are computed by `prove_step`,
+/// which allows the reuse of memory allocations and avoids unnecessary new allocations in the critical section.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(bound = "")]
+pub struct ResourceBuffer<E: Engine> {
+  l_w: Option<R1CSWitness<E>>,
+  l_u: Option<R1CSInstance<E>>,
+
+  ABC_Z_1: R1CSResult<E>,
+  ABC_Z_2: R1CSResult<E>,
+
+  /// buffer for `commit_T`
+  T: Vec<E::Scalar>,
+}
+
 /// A SNARK that proves the correct execution of an incremental computation
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(bound = "")]
@@ -247,6 +265,12 @@ where
   r_U_secondary: RelaxedR1CSInstance<E2>,
   l_w_secondary: R1CSWitness<E2>,
   l_u_secondary: R1CSInstance<E2>,
+
+  /// Buffer for memory needed by the primary fold-step
+  buffer_primary: ResourceBuffer<E1>,
+  /// Buffer for memory needed by the secondary fold-step
+  buffer_secondary: ResourceBuffer<E2>,
+
   i: usize,
   zi_primary: Vec<E1::Scalar>,
   zi_secondary: Vec<E2::Scalar>,
@@ -272,6 +296,9 @@ where
       return Err(NovaError::InvalidInitialInputLength);
     }
 
+    let r1cs_primary = &pp.r1cs_shape_primary;
+    let r1cs_secondary = &pp.r1cs_shape_secondary;
+
     // base case for the primary
     let mut cs_primary = SatisfyingAssignment::<E1>::new();
     let inputs_primary: NovaAugmentedCircuitInputs<E2> = NovaAugmentedCircuitInputs::new(
@@ -325,9 +352,8 @@ where
     // IVC proof for the secondary circuit
     let l_w_secondary = w_secondary;
     let l_u_secondary = u_secondary;
-    let r_W_secondary = RelaxedR1CSWitness::<E2>::default(&pp.r1cs_shape_secondary);
-    let r_U_secondary =
-      RelaxedR1CSInstance::<E2>::default(&pp.ck_secondary, &pp.r1cs_shape_secondary);
+    let r_W_secondary = RelaxedR1CSWitness::<E2>::default(r1cs_secondary);
+    let r_U_secondary = RelaxedR1CSInstance::<E2>::default(&pp.ck_secondary, r1cs_secondary);
 
     assert!(
       !(zi_primary.len() != pp.F_arity_primary || zi_secondary.len() != pp.F_arity_secondary),
@@ -344,6 +370,22 @@ where
       .map(|v| v.get_value().ok_or(SynthesisError::AssignmentMissing))
       .collect::<Result<Vec<<E2 as Engine>::Scalar>, _>>()?;
 
+    let buffer_primary = ResourceBuffer {
+      l_w: None,
+      l_u: None,
+      ABC_Z_1: R1CSResult::default(r1cs_primary),
+      ABC_Z_2: R1CSResult::default(r1cs_primary),
+      T: r1cs::default_T(r1cs_primary),
+    };
+
+    let buffer_secondary = ResourceBuffer {
+      l_w: None,
+      l_u: None,
+      ABC_Z_1: R1CSResult::default(r1cs_secondary),
+      ABC_Z_2: R1CSResult::default(r1cs_secondary),
+      T: r1cs::default_T(r1cs_secondary),
+    };
+
     Ok(Self {
       z0_primary: z0_primary.to_vec(),
       z0_secondary: z0_secondary.to_vec(),
@@ -353,6 +395,9 @@ where
       r_U_secondary,
       l_w_secondary,
       l_u_secondary,
+
+      buffer_primary,
+      buffer_secondary,
       i: 0,
       zi_primary,
       zi_secondary,
@@ -374,16 +419,24 @@ where
       return Ok(());
     }
 
+    // save the inputs before proceeding to the `i+1`th step
+    let r_U_primary_i = self.r_U_primary.clone();
+    let r_U_secondary_i = self.r_U_secondary.clone();
+    let l_u_secondary_i = self.l_u_secondary.clone();
+
     // fold the secondary circuit's instance
-    let (nifs_secondary, (r_U_secondary, r_W_secondary)) = NIFS::prove(
+    let nifs_secondary = NIFS::prove_mut(
       &pp.ck_secondary,
       &pp.ro_consts_secondary,
       &scalar_as_base::<E1>(pp.digest()),
       &pp.r1cs_shape_secondary,
-      &self.r_U_secondary,
-      &self.r_W_secondary,
+      &mut self.r_U_secondary,
+      &mut self.r_W_secondary,
       &self.l_u_secondary,
       &self.l_w_secondary,
+      &mut self.buffer_secondary.T,
+      &mut self.buffer_secondary.ABC_Z_1,
+      &mut self.buffer_secondary.ABC_Z_2,
     )?;
 
     let mut cs_primary = SatisfyingAssignment::<E1>::with_capacity(
@@ -395,8 +448,8 @@ where
       E1::Scalar::from(self.i as u64),
       self.z0_primary.to_vec(),
       Some(self.zi_primary.clone()),
-      Some(self.r_U_secondary.clone()),
-      Some(self.l_u_secondary.clone()),
+      Some(r_U_secondary_i),
+      Some(l_u_secondary_i),
       Some(Commitment::<E2>::decompress(&nifs_secondary.comm_T)?),
     );
 
@@ -412,15 +465,18 @@ where
       cs_primary.r1cs_instance_and_witness(&pp.r1cs_shape_primary, &pp.ck_primary)?;
 
     // fold the primary circuit's instance
-    let (nifs_primary, (r_U_primary, r_W_primary)) = NIFS::prove(
+    let nifs_primary = NIFS::prove_mut(
       &pp.ck_primary,
       &pp.ro_consts_primary,
       &pp.digest(),
       &pp.r1cs_shape_primary,
-      &self.r_U_primary,
-      &self.r_W_primary,
+      &mut self.r_U_primary,
+      &mut self.r_W_primary,
       &l_u_primary,
       &l_w_primary,
+      &mut self.buffer_primary.T,
+      &mut self.buffer_primary.ABC_Z_1,
+      &mut self.buffer_primary.ABC_Z_2,
     )?;
 
     let mut cs_secondary = SatisfyingAssignment::<E2>::with_capacity(
@@ -432,7 +488,7 @@ where
       E2::Scalar::from(self.i as u64),
       self.z0_secondary.to_vec(),
       Some(self.zi_secondary.clone()),
-      Some(self.r_U_primary.clone()),
+      Some(r_U_primary_i),
       Some(l_u_primary),
       Some(Commitment::<E1>::decompress(&nifs_primary.comm_T)?),
     );
@@ -462,14 +518,8 @@ where
     self.l_u_secondary = l_u_secondary;
     self.l_w_secondary = l_w_secondary;
 
-    self.r_U_primary = r_U_primary;
-    self.r_W_primary = r_W_primary;
-
     self.i += 1;
 
-    self.r_U_secondary = r_U_secondary;
-    self.r_W_secondary = r_W_secondary;
-
     Ok(())
   }
 

src/nifs.rs

@@ -4,7 +4,9 @@
 use crate::{
   constants::{NUM_CHALLENGE_BITS, NUM_FE_FOR_RO},
   errors::NovaError,
-  r1cs::{R1CSInstance, R1CSShape, R1CSWitness, RelaxedR1CSInstance, RelaxedR1CSWitness},
+  r1cs::{
+    R1CSInstance, R1CSResult, R1CSShape, R1CSWitness, RelaxedR1CSInstance, RelaxedR1CSWitness,
+  },
   scalar_as_base,
   traits::{commitment::CommitmentTrait, AbsorbInROTrait, Engine, ROTrait},
   Commitment, CommitmentKey, CompressedCommitment,
@@ -78,6 +80,56 @@ impl<E: Engine> NIFS<E> {
     ))
   }
 
+  /// Takes as input a Relaxed R1CS instance-witness tuple `(U1, W1)` and
+  /// an R1CS instance-witness tuple `(U2, W2)` with the same structure `shape`
+  /// and defined with respect to the same `ck`, and updates `(U1, W1)` by folding
+  /// `(U2, W2)` into it with the guarantee that the updated witness `W` satisfies
+  /// the updated instance `U` if and only if `W1` satisfies `U1` and `W2` satisfies `U2`.
+  #[allow(clippy::too_many_arguments)]
+  pub fn prove_mut(
+    ck: &CommitmentKey<E>,
+    ro_consts: &ROConstants<E>,
+    pp_digest: &E::Scalar,
+    S: &R1CSShape<E>,
+    U1: &mut RelaxedR1CSInstance<E>,
+    W1: &mut RelaxedR1CSWitness<E>,
+    U2: &R1CSInstance<E>,
+    W2: &R1CSWitness<E>,
+    T: &mut Vec<E::Scalar>,
+    ABC_Z_1: &mut R1CSResult<E>,
+    ABC_Z_2: &mut R1CSResult<E>,
+  ) -> Result<NIFS<E>, NovaError> {
+    // initialize a new RO
+    let mut ro = E::RO::new(ro_consts.clone(), NUM_FE_FOR_RO);
+
+    // append the digest of pp to the transcript
+    ro.absorb(scalar_as_base::<E>(*pp_digest));
+
+    // append U1 and U2 to transcript
+    U1.absorb_in_ro(&mut ro);
+    U2.absorb_in_ro(&mut ro);
+
+    // compute a commitment to the cross-term
+    let comm_T = S.commit_T_into(ck, U1, W1, U2, W2, T, ABC_Z_1, ABC_Z_2)?;
+
+    // append `comm_T` to the transcript and obtain a challenge
+    comm_T.absorb_in_ro(&mut ro);
+
+    // compute a challenge from the RO
+    let r = ro.squeeze(NUM_CHALLENGE_BITS);
+
+    // fold the instance using `r` and `comm_T`
+    U1.fold_mut(U2, &comm_T, &r);
+
+    // fold the witness using `r` and `T`
+    W1.fold_mut(W2, T, &r)?;
+
+    // return the commitment
+    Ok(Self {
+      comm_T: comm_T.compress(),
+    })
+  }
+
   /// Takes as input a relaxed R1CS instance `U1` and R1CS instance `U2`
   /// with the same shape and defined with respect to the same parameters,
   /// and outputs a folded instance `U` with the same shape,