implement intoiterator for lookup

src/gadgets/lookup.rs

@@ -1,6 +1,8 @@
 //! This module implements lookup gadget for applications built with Nova.
 use std::cmp::max;
+use std::collections::btree_map::Iter;
 use std::collections::BTreeMap;
+use std::collections::btree_map::Values;
 
 use bellpepper_core::{num::AllocatedNum, ConstraintSystem, LinearCombination, SynthesisError};
 use std::cmp::Ord;
@@ -508,7 +510,7 @@ impl<'a, G: Group> LookupTraceBuilder<'a, G> {
 /// Lookup in R1CS
 #[derive(Clone, Debug)]
 pub struct Lookup<F: PrimeField> {
-  pub(crate) map_aux: BTreeMap<F, (F, F)>, // (value, counter)
+  map_aux: BTreeMap<F, (F, F)>, // (value, counter)
   rw_counter: F,
   pub(crate) table_type: TableType, // read only or read-write
   pub(crate) max_cap_rwcounter_log2: usize, // max cap for rw_counter operation in bits
@@ -536,21 +538,15 @@ impl<F: PrimeField> Lookup<F> {
     }
   }
 
-  /// get table vector
-  /// very costly operation
-  pub fn get_table(&self) -> Vec<(F, F, F)> {
-    self
-      .map_aux
-      .iter()
-      .map(|(addr, (value, counter))| (*addr, *value, *counter))
-      .collect()
-  }
-
   /// table size
   pub fn table_size(&self) -> usize {
     self.map_aux.len()
   }
 
+  pub fn values(&self) -> Values<'_, F, (F, F)> {
+    self.map_aux.values()
+  }
+
   fn rw_operation(&mut self, addr: F, external_value: Option<F>) -> (F, F)
   where
     F: Ord,
@@ -577,8 +573,15 @@ impl<F: PrimeField> Lookup<F> {
     self.rw_counter = write_counter;
     (read_value, read_counter)
   }
+}
 
-  // fn write(&mut self, addr: AllocatedNum<F>, value: F) {}
+impl<'a, F: PrimeField> IntoIterator for &'a Lookup<F> {
+  type Item = (&'a F, &'a (F, F));
+  type IntoIter = Iter<'a, F, (F, F)>;
+
+  fn into_iter(self) -> Self::IntoIter {
+    self.map_aux.iter()
+  }
 }
 
 /// c = a + b where a, b is AllocatedNum

src/lib.rs

@@ -1710,11 +1710,10 @@ mod tests {
       G2: Group<Base = <G1 as Group>::Scalar>,
     {
       fn new(
-        initial_table: &Lookup<G1::Scalar>,
+        initial_table: Lookup<G1::Scalar>,
         ro_consts_circuit: ROConstantsCircuit<G2>,
       ) -> (Vec<Self>, Lookup<G1::Scalar>, G1::Scalar) {
         let n = initial_table.table_size();
-        let initial_table = initial_table.clone();
 
         let initial_index = (n - 4) / 2;
         let max_value_bits = (n - 1).log_2() + 1; // + 1 as a buffer
@@ -1976,7 +1975,7 @@ mod tests {
     };
 
     let (circuit_primaries, final_table, expected_intermediate_gamma) =
-      HeapifyCircuit::new(&initial_table, ro_consts);
+      HeapifyCircuit::new(initial_table.clone(), ro_consts);
 
     let circuit_secondary = TrivialCircuit::default();
 
@@ -2067,16 +2066,15 @@ mod tests {
 
     // lookup snark prove/verify
     type EE = crate::provider::ipa_pc::EvaluationEngine<G1>;
-    let (pk, vk) =
-      LookupSNARK::<G1, EE>::setup(&pp.ck_primary, &initial_table.get_table()).unwrap();
+    let (pk, vk) = LookupSNARK::<G1, EE>::setup(&pp.ck_primary, &initial_table).unwrap();
     let snark_proof = LookupSNARK::<G1, EE>::prove(
       &pp.ck_primary,
       &pk,
       (alpha, gamma),
       read_row,
       write_row,
-      initial_table.get_table(),
-      final_table.get_table(),
+      &initial_table,
+      &final_table,
     )
     .unwrap();
 

src/spartan/lookupsnark.rs

@@ -3,7 +3,7 @@ use crate::{
   constants::NUM_CHALLENGE_BITS,
   digest::{DigestComputer, SimpleDigestible},
   errors::NovaError,
-  gadgets::utils::scalar_as_base,
+  gadgets::{lookup::Lookup, utils::scalar_as_base},
   spartan::{
     math::Math,
     polys::{
@@ -120,17 +120,19 @@ where
   /// setup
   pub fn setup(
     ck: &CommitmentKey<G>,
-    initial_table: &Vec<(G::Scalar, G::Scalar, G::Scalar)>,
+    initial_table: &Lookup<G::Scalar>,
   ) -> Result<(ProverKey<G, EE>, VerifierKey<G, EE>), NovaError> {
     // check the provided commitment key meets minimal requirements
     // assert!(ck.length() >= Self::commitment_key_floor()(S));
-    let init_values: Vec<<G as Group>::Scalar> =
-      initial_table.iter().map(|(_, value, _)| *value).collect();
+    let init_values: Vec<<G as Group>::Scalar> = initial_table
+      .into_iter()
+      .map(|(_, (value, _))| *value)
+      .collect();
 
     let comm_init_value = G::CE::commit(ck, &init_values);
 
     let (pk_ee, vk_ee) = EE::setup(ck);
-    let table_size = initial_table.len();
+    let table_size = initial_table.table_size();
 
     let vk = VerifierKey::new(vk_ee, table_size, comm_init_value);
 
@@ -150,8 +152,8 @@ where
     challenges: (G::Scalar, G::Scalar),
     read_row: G::Scalar,
     write_row: G::Scalar,
-    initial_table: Vec<(G::Scalar, G::Scalar, G::Scalar)>,
-    final_table: Vec<(G::Scalar, G::Scalar, G::Scalar)>,
+    initial_table: &Lookup<G::Scalar>,
+    final_table: &Lookup<G::Scalar>,
   ) -> Result<Self, NovaError> {
     // a list of polynomial evaluation claims that will be batched
     let mut w_u_vec = Vec::new();
@@ -163,13 +165,13 @@ where
     };
     // init_row
     let initial_row: Vec<G::Scalar> = initial_table
-      .iter()
-      .map(|(addr, value, counter)| hash_func(addr, value, counter))
+      .into_iter()
+      .map(|(addr, (value, counter))| hash_func(addr, value, counter))
       .collect();
     // audit_row
     let audit_row: Vec<G::Scalar> = final_table
-      .iter()
-      .map(|(addr, value, counter)| hash_func(addr, value, counter))
+      .into_iter()
+      .map(|(addr, (value, counter))| hash_func(addr, value, counter))
       .collect();
     let mut transcript = G::TE::new(b"LookupSNARK");
     // append the verifier key (which includes commitment to R1CS matrices) and the read_row/write_row to the transcript
@@ -179,12 +181,12 @@ where
     transcript.absorb(b"alpha", &fingerprint_alpha);
     transcript.absorb(b"gamma", &fingerprint_gamma);
 
-    let init_values: Vec<<G as Group>::Scalar> =
-      initial_table.iter().map(|(_, value, _)| *value).collect();
-    let final_values: Vec<<G as Group>::Scalar> =
-      final_table.iter().map(|(_, value, _)| *value).collect();
-    let final_counters: Vec<<G as Group>::Scalar> =
-      final_table.iter().map(|(_, _, counter)| *counter).collect();
+    let init_values: Vec<<G as Group>::Scalar> = initial_table
+      .into_iter()
+      .map(|(_, (value, _))| *value)
+      .collect();
+    let (final_values, final_counters): (Vec<_>, Vec<_>) =
+      final_table.values().copied().unzip();
     let comm_init_value = pk.comm_init_value;
     let (comm_final_value, comm_final_counter) = rayon::join(
       || G::CE::commit(ck, &final_values),