From b813727df2111fd6b982fab174873a5fa40e4b36 Mon Sep 17 00:00:00 2001
From: Greg Zaverucha <gregz@microsoft.com>
Date: Fri, 4 Aug 2023 13:40:00 -0700
Subject: [PATCH] Proverkey (#2)

* Do not duplicate R1CS instance in ProverKey

The R1CSShape object was being stored in both the implementation of the
RelaxedR1CSSNARKTrait's ProverKey and the Spartan SNARK's ProverKey. Now we
store it once in the top level object (RelaxedR1CSSNARKTrait) and pass
it to the Spartan implementation.

This saves memory and makes serialization of the ProverKey about twice
as fast. Both are significant when there are a large number of
constraints.

Signed-off-by: Greg Zaverucha <gregz@microsoft.com>

* Clippy fix & whitespace

Signed-off-by: Greg Zaverucha <gregz@microsoft.com>

* Whitespace

Signed-off-by: Greg Zaverucha <gregz@microsoft.com>

* cargo fmt

---------

Signed-off-by: Greg Zaverucha <gregz@microsoft.com>
Co-authored-by: Srinath Setty <srinath@microsoft.com>
---
 src/lib.rs             |  2 +-
 src/spartan/ppsnark.rs | 17 ++++++++---------
 src/spartan/snark.rs   | 25 ++++++++++++-------------
 src/traits/snark.rs    |  1 +
 4 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index ad15ff9..39f2ba4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -110,7 +110,7 @@ impl<G: Group, S: RelaxedR1CSSNARKTrait<G>, C: Circuit<G::Scalar>> SNARK<G, S, C
     );
 
     // prove the instance using Spartan
-    let snark = S::prove(&pk.ck, &pk.pk, &u_relaxed, &w_relaxed)?;
+    let snark = S::prove(&pk.ck, &pk.pk, &pk.S, &u_relaxed, &w_relaxed)?;
 
     Ok(SNARK {
       comm_W: u.comm_W,
diff --git a/src/spartan/ppsnark.rs b/src/spartan/ppsnark.rs
index ac47fab..8faead2 100644
--- a/src/spartan/ppsnark.rs
+++ b/src/spartan/ppsnark.rs
@@ -729,7 +729,6 @@ impl<G: Group> SumcheckEngine<G> for InnerSumcheckInstance<G> {
 #[serde(bound = "")]
 pub struct ProverKey<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> {
   pk_ee: EE::ProverKey,
-  S: R1CSShape<G>,
   S_repr: R1CSShapeSparkRepr<G>,
   S_comm: R1CSShapeSparkCommitment<G>,
   vk_digest: G::Scalar, // digest of verifier's key
@@ -944,7 +943,6 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
 
     let pk = ProverKey {
       pk_ee,
-      S,
       S_repr,
       S_comm,
       vk_digest: vk.digest,
@@ -957,19 +955,20 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
   fn prove(
     ck: &CommitmentKey<G>,
     pk: &Self::ProverKey,
+    S: &R1CSShape<G>,
     U: &RelaxedR1CSInstance<G>,
     W: &RelaxedR1CSWitness<G>,
   ) -> Result<Self, SpartanError> {
-    let W = W.pad(&pk.S); // pad the witness
+    let W = W.pad(S); // pad the witness
     let mut transcript = G::TE::new(b"RelaxedR1CSSNARK");
 
     // a list of polynomial evaluation claims that will be batched
     let mut w_u_vec = Vec::new();
 
     // sanity check that R1CSShape has certain size characteristics
-    assert_eq!(pk.S.num_cons.next_power_of_two(), pk.S.num_cons);
-    assert_eq!(pk.S.num_vars.next_power_of_two(), pk.S.num_vars);
-    assert!(pk.S.num_io < pk.S.num_vars);
+    assert_eq!(S.num_cons.next_power_of_two(), S.num_cons);
+    assert_eq!(S.num_vars.next_power_of_two(), S.num_vars);
+    assert!(S.num_io < S.num_vars);
 
     // append the verifier key (which includes commitment to R1CS matrices) and the RelaxedR1CSInstance to the transcript
     transcript.absorb(b"vk", &pk.vk_digest);
@@ -979,7 +978,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
     let z = concat(vec![W.W.clone(), vec![U.u], U.X.clone()]);
 
     // compute Az, Bz, Cz
-    let (mut Az, mut Bz, mut Cz) = pk.S.multiply_vec(&z)?;
+    let (mut Az, mut Bz, mut Cz) = S.multiply_vec(&z)?;
 
     // commit to Az, Bz, Cz
     let (comm_Az, (comm_Bz, comm_Cz)) = rayon::join(
@@ -1017,7 +1016,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
     // (2) send commitments to the following two oracles
     // E_row(i) = eq(tau, row(i)) for all i
     // E_col(i) = z(col(i)) for all i
-    let (mem_row, mem_col, E_row, E_col) = pk.S_repr.evaluation_oracles(&pk.S, &tau, &z);
+    let (mem_row, mem_col, E_row, E_col) = pk.S_repr.evaluation_oracles(S, &tau, &z);
     let (comm_E_row, comm_E_col) =
       rayon::join(|| G::CE::commit(ck, &E_row), || G::CE::commit(ck, &E_col));
 
@@ -1331,7 +1330,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
     // we need to prove that eval_z = z(r_prod) = (1-r_prod[0]) * W.w(r_prod[1..]) + r_prod[0] * U.x(r_prod[1..]).
     // r_prod was padded, so we now remove the padding
     let r_prod_unpad = {
-      let l = pk.S_repr.N.log_2() - (2 * pk.S.num_vars).log_2();
+      let l = pk.S_repr.N.log_2() - (2 * S.num_vars).log_2();
       r_prod[l..].to_vec()
     };
 
diff --git a/src/spartan/snark.rs b/src/spartan/snark.rs
index 4356976..1461d03 100644
--- a/src/spartan/snark.rs
+++ b/src/spartan/snark.rs
@@ -29,7 +29,6 @@ use serde::{Deserialize, Serialize};
 #[serde(bound = "")]
 pub struct ProverKey<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> {
   pk_ee: EE::ProverKey,
-  S: R1CSShape<G>,
   vk_digest: G::Scalar, // digest of the verifier's key
 }
 
@@ -70,7 +69,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
   ) -> Result<(Self::ProverKey, Self::VerifierKey), SpartanError> {
     let (pk_ee, vk_ee) = EE::setup(ck);
 
-    let S = S.pad();
+    let S = &S.pad();
 
     let vk = {
       let mut vk = VerifierKey {
@@ -84,7 +83,6 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
 
     let pk = ProverKey {
       pk_ee,
-      S,
       vk_digest: vk.digest,
     };
 
@@ -95,16 +93,17 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
   fn prove(
     ck: &CommitmentKey<G>,
     pk: &Self::ProverKey,
+    S: &R1CSShape<G>,
     U: &RelaxedR1CSInstance<G>,
     W: &RelaxedR1CSWitness<G>,
   ) -> Result<Self, SpartanError> {
-    let W = W.pad(&pk.S); // pad the witness
+    let W = W.pad(S); // pad the witness
     let mut transcript = G::TE::new(b"RelaxedR1CSSNARK");
 
     // sanity check that R1CSShape has certain size characteristics
-    assert_eq!(pk.S.num_cons.next_power_of_two(), pk.S.num_cons);
-    assert_eq!(pk.S.num_vars.next_power_of_two(), pk.S.num_vars);
-    assert!(pk.S.num_io < pk.S.num_vars);
+    assert_eq!(S.num_cons.next_power_of_two(), S.num_cons);
+    assert_eq!(S.num_vars.next_power_of_two(), S.num_vars);
+    assert!(S.num_io < S.num_vars);
 
     // append the digest of vk (which includes R1CS matrices) and the RelaxedR1CSInstance to the transcript
     transcript.absorb(b"vk", &pk.vk_digest);
@@ -114,8 +113,8 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
     let mut z = concat(vec![W.W.clone(), vec![U.u], U.X.clone()]);
 
     let (num_rounds_x, num_rounds_y) = (
-      (pk.S.num_cons as f64).log2() as usize,
-      ((pk.S.num_vars as f64).log2() as usize + 1),
+      (S.num_cons as f64).log2() as usize,
+      ((S.num_vars as f64).log2() as usize + 1),
     );
 
     // outer sum-check
@@ -125,8 +124,8 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
 
     let mut poly_tau = MultilinearPolynomial::new(EqPolynomial::new(tau).evals());
     let (mut poly_Az, mut poly_Bz, poly_Cz, mut poly_uCz_E) = {
-      let (poly_Az, poly_Bz, poly_Cz) = pk.S.multiply_vec(&z)?;
-      let poly_uCz_E = (0..pk.S.num_cons)
+      let (poly_Az, poly_Bz, poly_Cz) = S.multiply_vec(&z)?;
+      let poly_uCz_E = (0..S.num_cons)
         .map(|i| U.u * poly_Cz[i] + W.E[i])
         .collect::<Vec<G::Scalar>>();
       (
@@ -207,7 +206,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
           (A_evals, B_evals, C_evals)
         };
 
-      let (evals_A, evals_B, evals_C) = compute_eval_table_sparse(&pk.S, &evals_rx);
+      let (evals_A, evals_B, evals_C) = compute_eval_table_sparse(S, &evals_rx);
 
       assert_eq!(evals_A.len(), evals_B.len());
       assert_eq!(evals_A.len(), evals_C.len());
@@ -218,7 +217,7 @@ impl<G: Group, EE: EvaluationEngineTrait<G, CE = G::CE>> RelaxedR1CSSNARKTrait<G
     };
 
     let poly_z = {
-      z.resize(pk.S.num_vars * 2, G::Scalar::ZERO);
+      z.resize(S.num_vars * 2, G::Scalar::ZERO);
       z
     };
 
diff --git a/src/traits/snark.rs b/src/traits/snark.rs
index 8b136b5..5ed56c8 100644
--- a/src/traits/snark.rs
+++ b/src/traits/snark.rs
@@ -28,6 +28,7 @@ pub trait RelaxedR1CSSNARKTrait<G: Group>:
   fn prove(
     ck: &CommitmentKey<G>,
     pk: &Self::ProverKey,
+    S: &R1CSShape<G>,
     U: &RelaxedR1CSInstance<G>,
     W: &RelaxedR1CSWitness<G>,
   ) -> Result<Self, SpartanError>;