Change PartialWitness to use Vecs

src/bin/bench_recursion.rs

@@ -49,7 +49,7 @@ 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 inputs = PartialWitness::new(0, 0, 0);
     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

@@ -604,7 +604,7 @@ mod tests {
         let config = CircuitConfig::large_config();
 
         let mut builder = CircuitBuilder::<F, D>::new(config);
-        let mut pw = PartialWitness::new();
+        let mut pw = PartialWitness::new(0, 0, 0);
 
         let vs = FF::rand_vec(3);
         let ts = builder.add_virtual_extension_targets(3);
@@ -654,7 +654,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(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/insert.rs

@@ -83,7 +83,7 @@ mod tests {
         }
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/interpolation.rs

@@ -99,7 +99,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(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -133,7 +133,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(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/random_access.rs

@@ -61,7 +61,7 @@ mod tests {
         }
 
         let data = builder.build();
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gadgets/select.rs

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

src/gadgets/split_base.rs

@@ -115,7 +115,7 @@ mod tests {
         builder.assert_leading_zeros(xt, 64 - 9);
         let data = builder.build();
 
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -147,7 +147,7 @@ mod tests {
 
         let data = builder.build();
 
-        let proof = data.prove(PartialWitness::new())?;
+        let proof = data.prove(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }

src/gates/gate_testing.rs

@@ -125,7 +125,7 @@ pub(crate) fn test_eval_fns<F: Extendable<D>, G: Gate<F, D>, const D: usize>(
 
     let config = CircuitConfig::large_config();
     let mut builder = CircuitBuilder::<F, D>::new(config);
-    let mut pw = PartialWitness::new();
+    let mut pw = PartialWitness::new(0, 0, 0);
 
     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(0, 0, 0);
         witness.set_wire(
             Wire {
                 gate: 0,

src/hash/merkle_proofs.rs

@@ -155,7 +155,7 @@ mod tests {
         type F = CrandallField;
         let config = CircuitConfig::large_config();
         let mut builder = CircuitBuilder::<F, 4>::new(config);
-        let mut pw = PartialWitness::new();
+        let mut pw = PartialWitness::new(0, 0, 0);
 
         let log_n = 8;
         let n = 1 << log_n;

src/iop/challenger.rs

@@ -409,7 +409,7 @@ mod tests {
             );
         }
         let circuit = builder.build();
-        let mut witness = PartialWitness::new();
+        let mut witness = PartialWitness::new(0, 0, 0);
         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(degree: usize, num_wires: usize, max_virtual_target: usize) -> Self {
         PartialWitness {
-            target_values: HashMap::new(),
+            wire_values: vec![vec![None; num_wires]; degree],
+            virtual_target_values: vec![None; max_virtual_target],
         }
     }
 
-    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,10 @@ 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[gate][input].is_some(),
+            Target::VirtualTarget { index } => self.virtual_target_values[index].is_some(),
+        }
     }
 
     pub fn contains_all(&self, targets: &[Target]) -> bool {
@@ -96,13 +103,29 @@ 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 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 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 +185,18 @@ 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;
+        for i in 0..degree {
+            for j in 0..num_wires {
+                wire_values[j][i] = self.wire_values[i][j].unwrap_or(F::ZERO);
             }
-        });
+        }
         Witness { wire_values }
     }
 
@@ -215,6 +240,6 @@ impl<F: Field> PartialWitness<F> {
 
 impl<F: Field> Default for PartialWitness<F> {
     fn default() -> Self {
-        Self::new()
+        Self::new(0, 0, 0)
     }
 }

src/plonk/recursive_verifier.rs

@@ -387,7 +387,7 @@ mod tests {
             }
             let data = builder.build();
             (
-                data.prove(PartialWitness::new())?,
+                data.prove(PartialWitness::new(0, 0, 0))?,
                 data.verifier_only,
                 data.common,
             )
@@ -395,7 +395,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(0, 0, 0);
         let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
         set_proof_target(&proof_with_pis, &pt, &mut pw);
 
@@ -442,7 +442,7 @@ mod tests {
                 }
                 let data = builder.build();
                 (
-                    data.prove(PartialWitness::new())?,
+                    data.prove(PartialWitness::new(1 << 14, config.num_wires, 1000))?,
                     data.verifier_only,
                     data.common,
                 )
@@ -450,7 +450,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(1 << 14, config.num_wires, 100000);
             let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
             set_proof_target(&proof_with_pis, &pt, &mut pw);
 
@@ -468,7 +468,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(1 << 14, config.num_wires, 100000);
         let pt = proof_to_proof_target(&proof_with_pis, &mut builder);
         set_proof_target(&proof_with_pis, &pt, &mut pw);
 

src/util/reducing.rs

@@ -264,7 +264,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(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }
@@ -294,7 +294,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(PartialWitness::new(0, 0, 0))?;
 
         verify(proof, &data.verifier_only, &data.common)
     }