From 233cf2f482fbe1888ba7ac9d0932f0115bb83ca1 Mon Sep 17 00:00:00 2001
From: Hanting Zhang <hantingz@usc.edu>
Date: Thu, 18 Apr 2024 01:00:13 -0700
Subject: [PATCH] write results too

---
 src/data.rs     | 16 ++++++++++++++++
 src/lib.rs      | 37 ++++++++++++++++++-------------------
 src/nifs.rs     |  2 +-
 src/r1cs/mod.rs | 13 +++++++++++++
 4 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/src/data.rs b/src/data.rs
index f402e3c1e..7bb9e8825 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -31,6 +31,7 @@ pub struct DataConfig {
   #[cfg(not(target_arch = "wasm32"))]
   section_counters: HashMap<String, usize>,
   write_data: bool,
+  witness_size: usize,
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -47,6 +48,7 @@ pub fn init_config() -> Mutex<DataConfig> {
     root_dir,
     section_counters: HashMap::new(),
     write_data: WRITE,
+    witness_size: 0,
   };
 
   Mutex::new(config)
@@ -150,3 +152,17 @@ pub fn set_write_data(write_data: bool) {
   let mut config = mutex.lock().unwrap();
   config.write_data = write_data;
 }
+
+/// Are we configured to write data?
+pub fn witness_size() -> usize {
+  let mutex = ARECIBO_CONFIG.get_or_init(init_config);
+  let config = mutex.lock().unwrap();
+  config.witness_size
+}
+
+/// Set the configuration for writing data.
+pub fn set_witness_size(witness_size: usize) {
+  let mutex = ARECIBO_CONFIG.get_or_init(init_config);
+  let mut config = mutex.lock().unwrap();
+  config.witness_size = witness_size;
+}
diff --git a/src/lib.rs b/src/lib.rs
index f5762a5c8..8616cd232 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -28,7 +28,7 @@ pub mod spartan;
 pub mod supernova;
 pub mod traits;
 
-use crate::data::{write_arecibo_data, write_data};
+use crate::data::{set_witness_size, write_arecibo_data, write_data};
 use crate::digest::{DigestComputer, SimpleDigestible};
 use crate::{
   bellpepper::{
@@ -491,9 +491,23 @@ where
     };
 
     if write_data() {
-      write_arecibo_data(format!("sparse_matrices_{:?}", pp.digest()), "A", &r1cs_primary.A);
-      write_arecibo_data(format!("sparse_matrices_{:?}", pp.digest()), "B", &r1cs_primary.B);
-      write_arecibo_data(format!("sparse_matrices_{:?}", pp.digest()), "C", &r1cs_primary.C);
+      write_arecibo_data(
+        format!("sparse_matrices_{:?}", pp.digest()),
+        "A",
+        &r1cs_primary.A,
+      );
+      write_arecibo_data(
+        format!("sparse_matrices_{:?}", pp.digest()),
+        "B",
+        &r1cs_primary.B,
+      );
+      write_arecibo_data(
+        format!("sparse_matrices_{:?}", pp.digest()),
+        "C",
+        &r1cs_primary.C,
+      );
+
+      set_witness_size(r1cs_primary.A.num_cols());
     }
 
     Ok(Self {
@@ -601,21 +615,6 @@ where
       &mut self.buffer_primary.ABC_Z_2,
     )?;
 
-    if write_data() {
-      let W = l_w_primary.W;
-      write_arecibo_data(
-        format!("witness_{:?}", pp.digest()),
-        format!("len_{}", W.len()),
-        &W,
-      );
-      let T = &self.buffer_primary.T;
-      write_arecibo_data(
-        format!("cross_term_{:?}", pp.digest()),
-        format!("len_{}", T.len()),
-        &T,
-      );
-    }
-
     let mut cs_secondary = SatisfyingAssignment::<Dual<E1>>::with_capacity(
       pp.circuit_shape_secondary.r1cs_shape.num_io + 1,
       pp.circuit_shape_secondary.r1cs_shape.num_vars,
diff --git a/src/nifs.rs b/src/nifs.rs
index f0630e533..bb5d5f338 100644
--- a/src/nifs.rs
+++ b/src/nifs.rs
@@ -126,7 +126,7 @@ impl<E: Engine> NIFS<E> {
     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)?;
+    let comm_T = S.commit_T_into(ck, pp_digest, 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);
diff --git a/src/r1cs/mod.rs b/src/r1cs/mod.rs
index 8edc796ce..c95eadb9f 100644
--- a/src/r1cs/mod.rs
+++ b/src/r1cs/mod.rs
@@ -4,6 +4,7 @@ pub(crate) mod util;
 
 use crate::{
   constants::{BN_LIMB_WIDTH, BN_N_LIMBS},
+  data::{witness_size, write_arecibo_data, write_data},
   digest::{DigestComputer, SimpleDigestible},
   errors::NovaError,
   gadgets::{f_to_nat, nat_to_limbs, scalar_as_base},
@@ -465,6 +466,7 @@ impl<E: Engine> R1CSShape<E> {
   pub fn commit_T_into(
     &self,
     ck: &CommitmentKey<E>,
+    pp_digest: &E::Scalar,
     U1: &RelaxedR1CSInstance<E>,
     W1: &RelaxedR1CSWitness<E>,
     U2: &R1CSInstance<E>,
@@ -473,6 +475,11 @@ impl<E: Engine> R1CSShape<E> {
     ABC_Z_1: &mut R1CSResult<E>,
     ABC_Z_2: &mut R1CSResult<E>,
   ) -> Result<Commitment<E>, NovaError> {
+    if write_data() && self.A.num_cols() == witness_size() {
+      let witness = [&W1.W[..], &[U1.u], &U1.X[..]].concat();
+      write_arecibo_data(format!("witness_{:?}", pp_digest), "", &witness);
+    }
+
     tracing::info_span!("AZ_1, BZ_1, CZ_1")
       .in_scope(|| self.multiply_witness_into(&W1.W, &U1.u, &U1.X, ABC_Z_1))?;
 
@@ -482,6 +489,12 @@ impl<E: Engine> R1CSShape<E> {
       CZ: CZ_1,
     } = ABC_Z_1;
 
+    if write_data() && self.A.num_cols() == witness_size() {
+      write_arecibo_data(format!("result_{:?}", pp_digest), "AZ", &AZ_1);
+      write_arecibo_data(format!("result_{:?}", pp_digest), "BZ", &BZ_1);
+      write_arecibo_data(format!("result_{:?}", pp_digest), "CZ", &CZ_1);
+    }
+
     tracing::info_span!("AZ_2, BZ_2, CZ_2")
       .in_scope(|| self.multiply_witness_into(&W2.W, &E::Scalar::ONE, &U2.X, ABC_Z_2))?;