Merge pull request #165 from mir-protocol/partial_witness_vec

Replace `HashMap` with `Vec`s in `PartialWitness`

src/bin/bench_recursion.rs

@@ -34,6 +34,7 @@ fn bench_prove<F: Field + Extendable<D>, const D: usize>() -> Result<()> {
         },
     };
 
+    let inputs = PartialWitness::new(config.num_wires);
     let mut builder = CircuitBuilder::<F, D>::new(config);
 
     let zero = builder.zero();
@@ -49,7 +50,6 @@ fn bench_prove<F: Field + Extendable<D>, const D: usize>() -> Result<()> {
     builder.add_extension(zero_ext, zero_ext);
 
     let circuit = builder.build();
-    let inputs = PartialWitness::new();
     let proof_with_pis = circuit.prove(inputs)?;
     let proof_bytes = serde_cbor::to_vec(&proof_with_pis).unwrap();
     info!("Proof length: {} bytes", proof_bytes.len());

src/gadgets/arithmetic_extension.rs

@@ -603,8 +603,8 @@ mod tests {
 
         let config = CircuitConfig::large_config();
 
+        let mut pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, D>::new(config);
-        let mut pw = PartialWitness::new();
 
         let vs = FF::rand_vec(3);
         let ts = builder.add_virtual_extension_targets(3);
@@ -640,6 +640,7 @@ mod tests {
 
         let config = CircuitConfig::large_config();
 
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
         let x = FF::rand();
@@ -654,7 +655,7 @@ mod tests {
         builder.assert_equal_extension(zt, comp_zt_unsafe);
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/insert.rs

@@ -64,6 +64,7 @@ mod tests {
         type FF = QuarticCrandallField;
         let len = 1 << len_log;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
         let v = (0..len - 1)
             .map(|_| builder.constant_extension(FF::rand()))
@@ -83,7 +84,7 @@ mod tests {
         }
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/interpolation.rs

@@ -70,6 +70,7 @@ mod tests {
         type F = CrandallField;
         type FF = QuarticCrandallField;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
 
         let len = 4;
@@ -99,7 +100,7 @@ mod tests {
         builder.assert_equal_extension(eval, true_eval_target);
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -109,6 +110,7 @@ mod tests {
         type F = CrandallField;
         type FF = QuarticCrandallField;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
 
         let len = 2;
@@ -133,7 +135,7 @@ mod tests {
         builder.assert_equal_extension(eval, true_eval_target);
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/random_access.rs

@@ -50,6 +50,7 @@ mod tests {
         type FF = QuarticCrandallField;
         let len = 1 << len_log;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
         let vec = FF::rand_vec(len);
         let v: Vec<_> = vec.iter().map(|x| builder.constant_extension(*x)).collect();
@@ -61,7 +62,7 @@ mod tests {
         }
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/select.rs

@@ -48,8 +48,8 @@ mod tests {
         type F = CrandallField;
         type FF = QuarticCrandallField;
         let config = CircuitConfig::large_config();
+        let mut pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
-        let mut pw = PartialWitness::new();
 
         let (x, y) = (FF::rand(), FF::rand());
         let xt = builder.add_virtual_extension_target();

src/gadgets/split_base.rs

@@ -99,6 +99,7 @@ mod tests {
     fn test_split_base() -> Result<()> {
         type F = CrandallField;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
         let x = F::from_canonical_usize(0b110100000); // 416 = 1532 in base 6.
         let xt = builder.constant(x);
@@ -115,7 +116,7 @@ mod tests {
         builder.assert_leading_zeros(xt, 64 - 9);
         let data = builder.build();
 
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -124,6 +125,7 @@ mod tests {
     fn test_base_sum() -> Result<()> {
         type F = CrandallField;
         let config = CircuitConfig::large_config();
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
 
         let n = thread_rng().gen_range(0, 1 << 10);
@@ -147,7 +149,7 @@ mod tests {
 
         let data = builder.build();
 
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gates/gate_testing.rs

@@ -124,8 +124,8 @@ pub(crate) fn test_eval_fns<F: Extendable<D>, G: Gate<F, D>, const D: usize>(
     let constants = F::Extension::rand_vec(gate.num_constants());
 
     let config = CircuitConfig::large_config();
+    let mut pw = PartialWitness::new(config.num_wires);
     let mut builder = CircuitBuilder::<F, D>::new(config);
-    let mut pw = PartialWitness::new();
 
     let wires_t = builder.add_virtual_extension_targets(wires.len());
     let constants_t = builder.add_virtual_extension_targets(constants.len());

src/gates/gmimc.rs

@@ -345,7 +345,7 @@ mod tests {
 
         let permutation_inputs = (0..W).map(F::from_canonical_usize).collect::<Vec<_>>();
 
-        let mut witness = PartialWitness::new();
+        let mut witness = PartialWitness::new(gate.num_wires());
         witness.set_wire(
             Wire {
                 gate: 0,

src/hash/merkle_proofs.rs

@@ -154,8 +154,8 @@ mod tests {
     fn test_recursive_merkle_proof() -> Result<()> {
         type F = CrandallField;
         let config = CircuitConfig::large_config();
+        let mut pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config);
-        let mut pw = PartialWitness::new();
 
         let log_n = 8;
         let n = 1 << log_n;

src/iop/challenger.rs

@@ -399,6 +399,7 @@ mod tests {
             num_routed_wires: 27,
             ..CircuitConfig::default()
         };
+        let mut witness = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, 4>::new(config.clone());
         let mut recursive_challenger = RecursiveChallenger::new(&mut builder);
         let mut recursive_outputs_per_round: Vec<Vec<Target>> = Vec::new();
@@ -409,7 +410,6 @@ mod tests {
             );
         }
         let circuit = builder.build();
-        let mut witness = PartialWitness::new();
         generate_partial_witness(
             &mut witness,
             &circuit.prover_only.generators,

src/iop/witness.rs

@@ -26,22 +26,23 @@ impl<F: Field> Witness<F> {
 
 #[derive(Clone, Debug)]
 pub struct PartialWitness<F: Field> {
-    pub(crate) target_values: HashMap<Target, F>,
+    pub(crate) wire_values: Vec<Vec<Option<F>>>,
+    pub(crate) virtual_target_values: Vec<Option<F>>,
 }
 
 impl<F: Field> PartialWitness<F> {
-    pub fn new() -> Self {
+    pub fn new(num_wires: usize) -> Self {
         PartialWitness {
-            target_values: HashMap::new(),
+            wire_values: vec![vec![None; num_wires]],
+            virtual_target_values: vec![],
         }
     }
 
-    pub fn is_empty(&self) -> bool {
-        self.target_values.is_empty()
-    }
-
     pub fn get_target(&self, target: Target) -> F {
-        self.target_values[&target]
+        match target {
+            Target::Wire(Wire { gate, input }) => self.wire_values[gate][input].unwrap(),
+            Target::VirtualTarget { index } => self.virtual_target_values[index].unwrap(),
+        }
     }
 
     pub fn get_targets(&self, targets: &[Target]) -> Vec<F> {
@@ -76,7 +77,10 @@ impl<F: Field> PartialWitness<F> {
     }
 
     pub fn try_get_target(&self, target: Target) -> Option<F> {
-        self.target_values.get(&target).cloned()
+        match target {
+            Target::Wire(Wire { gate, input }) => self.wire_values[gate][input],
+            Target::VirtualTarget { index } => self.virtual_target_values[index],
+        }
     }
 
     pub fn get_wire(&self, wire: Wire) -> F {
@@ -88,7 +92,15 @@ impl<F: Field> PartialWitness<F> {
     }
 
     pub fn contains(&self, target: Target) -> bool {
-        self.target_values.contains_key(&target)
+        match target {
+            Target::Wire(Wire { gate, input }) => {
+                self.wire_values.len() > gate && self.wire_values[gate][input].is_some()
+            }
+            Target::VirtualTarget { index } => {
+                self.virtual_target_values.len() > index
+                    && self.virtual_target_values[index].is_some()
+            }
+        }
     }
 
     pub fn contains_all(&self, targets: &[Target]) -> bool {
@@ -96,13 +108,36 @@ impl<F: Field> PartialWitness<F> {
     }
 
     pub fn set_target(&mut self, target: Target, value: F) {
-        let opt_old_value = self.target_values.insert(target, value);
-        if let Some(old_value) = opt_old_value {
-            assert_eq!(
-                old_value, value,
-                "Target was set twice with different values: {:?}",
-                target
-            );
+        match target {
+            Target::Wire(Wire { gate, input }) => {
+                if gate >= self.wire_values.len() {
+                    self.wire_values
+                        .resize(gate + 1, vec![None; self.wire_values[0].len()]);
+                }
+                if let Some(old_value) = self.wire_values[gate][input] {
+                    assert_eq!(
+                        old_value, value,
+                        "Target was set twice with different values: {:?}",
+                        target
+                    );
+                } else {
+                    self.wire_values[gate][input] = Some(value);
+                }
+            }
+            Target::VirtualTarget { index } => {
+                if index >= self.virtual_target_values.len() {
+                    self.virtual_target_values.resize(index + 1, None);
+                }
+                if let Some(old_value) = self.virtual_target_values[index] {
+                    assert_eq!(
+                        old_value, value,
+                        "Target was set twice with different values: {:?}",
+                        target
+                    );
+                } else {
+                    self.virtual_target_values[index] = Some(value);
+                }
+            }
         }
     }
 
@@ -162,16 +197,19 @@ impl<F: Field> PartialWitness<F> {
     }
 
     pub fn extend<I: Iterator<Item = (Target, F)>>(&mut self, pairs: I) {
-        self.target_values.extend(pairs);
+        for (t, v) in pairs {
+            self.set_target(t, v);
+        }
     }
 
     pub fn full_witness(self, degree: usize, num_wires: usize) -> Witness<F> {
         let mut wire_values = vec![vec![F::ZERO; degree]; num_wires];
-        self.target_values.into_iter().for_each(|(t, v)| {
-            if let Target::Wire(Wire { gate, input }) = t {
-                wire_values[input][gate] = v;
+        assert!(self.wire_values.len() <= degree);
+        for i in 0..self.wire_values.len() {
+            for j in 0..num_wires {
+                wire_values[j][i] = self.wire_values[i][j].unwrap_or(F::ZERO);
             }
-        });
+        }
         Witness { wire_values }
     }
 
@@ -212,9 +250,3 @@ impl<F: Field> PartialWitness<F> {
         Ok(())
     }
 }
-
-impl<F: Field> Default for PartialWitness<F> {
-    fn default() -> Self {
-        Self::new()
-    }
-}

src/plonk/recursive_verifier.rs

@@ -383,7 +383,7 @@ mod tests {
             }
             let data = builder.build();
             (
-                data.prove(PartialWitness::new())?,
+                data.prove(PartialWitness::new(config.num_wires))?,
                 data.verifier_only,
                 data.common,
             )
@@ -391,7 +391,7 @@ mod tests {
         verify(proof_with_pis.clone(), &vd, &cd)?;
 
         let mut builder = CircuitBuilder::<F, D>::new(config.clone());
-        let mut pw = PartialWitness::new();
+        let mut pw = PartialWitness::new(config.num_wires);
         let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
         set_proof_target(&proof_with_pis, &pt, &mut pw);
 
@@ -438,7 +438,7 @@ mod tests {
                 }
                 let data = builder.build();
                 (
-                    data.prove(PartialWitness::new())?,
+                    data.prove(PartialWitness::new(config.num_wires))?,
                     data.verifier_only,
                     data.common,
                 )
@@ -446,7 +446,7 @@ mod tests {
             verify(proof_with_pis.clone(), &vd, &cd)?;
 
             let mut builder = CircuitBuilder::<F, D>::new(config.clone());
-            let mut pw = PartialWitness::new();
+            let mut pw = PartialWitness::new(config.num_wires);
             let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
             set_proof_target(&proof_with_pis, &pt, &mut pw);
 
@@ -464,7 +464,7 @@ mod tests {
 
         verify(proof_with_pis.clone(), &vd, &cd)?;
         let mut builder = CircuitBuilder::<F, D>::new(config.clone());
-        let mut pw = PartialWitness::new();
+        let mut pw = PartialWitness::new(config.num_wires);
         let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
         set_proof_target(&proof_with_pis, &pt, &mut pw);
 

src/util/reducing.rs

@@ -249,6 +249,7 @@ mod tests {
 
         let config = CircuitConfig::large_config();
 
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
         let alpha = FF::rand();
@@ -264,7 +265,7 @@ mod tests {
         builder.assert_equal_extension(manual_reduce, circuit_reduce);
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -276,6 +277,7 @@ mod tests {
 
         let config = CircuitConfig::large_config();
 
+        let pw = PartialWitness::new(config.num_wires);
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
         let alpha = FF::rand();
@@ -294,7 +296,7 @@ mod tests {
         builder.assert_equal_extension(manual_reduce, circuit_reduce);
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(pw)?;
 
         verify(proof, &data.verifier_only, &data.common)
     }