Replace HashMap by Vec in generate_partial_witness

src/gates/gmimc.rs

@@ -364,7 +364,14 @@ mod tests {
         }
 
         let generators = gate.generators(0, &[]);
-        generate_partial_witness(&mut witness, &generators, &mut TimingTree::default());
+        generate_partial_witness(
+            &mut witness,
+            &generators,
+            33,
+            100,
+            100,
+            &mut TimingTree::default(),
+        );
 
         let expected_outputs: [F; W] =
             gmimc_permute_naive(permutation_inputs.try_into().unwrap(), constants);

src/iop/challenger.rs

@@ -413,6 +413,9 @@ mod tests {
         generate_partial_witness(
             &mut witness,
             &circuit.prover_only.generators,
+            33,
+            100,
+            100,
             &mut TimingTree::default(),
         );
         let recursive_output_values_per_round: Vec<Vec<F>> = recursive_outputs_per_round

src/iop/generator.rs

@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::convert::identity;
 use std::fmt::Debug;
 
@@ -17,17 +16,26 @@ use crate::util::timing::TimingTree;
 pub(crate) fn generate_partial_witness<F: Field>(
     witness: &mut PartialWitness<F>,
     generators: &[Box<dyn WitnessGenerator<F>>],
+    num_routed_wires: usize,
+    degree: usize,
+    max_virtual_target: usize,
     timing: &mut TimingTree,
 ) {
+    let ind = |t: Target| -> usize {
+        match t {
+            Target::Wire(Wire { gate, input }) => gate * num_routed_wires + input,
+            Target::VirtualTarget { index } => degree * num_routed_wires + index,
+        }
+    };
+    let max_ind = ind(Target::VirtualTarget {
+        index: max_virtual_target,
+    });
     // Index generator indices by their watched targets.
-    let mut generator_indices_by_watches = HashMap::new();
+    let mut generator_indices_by_watches = vec![Vec::new(); max_ind];
     timed!(timing, "index generators by their watched targets", {
         for (i, generator) in generators.iter().enumerate() {
             for watch in generator.watch_list() {
-                generator_indices_by_watches
-                    .entry(watch)
-                    .or_insert_with(Vec::new)
-                    .push(i);
+                generator_indices_by_watches[ind(watch)].push(i);
             }
         }
     });
@@ -56,11 +64,9 @@ pub(crate) fn generate_partial_witness<F: Field>(
             }
 
             // Enqueue unfinished generators that were watching one of the newly populated targets.
-            for (watch, _) in &buffer.target_values {
-                if let Some(watching_generator_indices) = generator_indices_by_watches.get(watch) {
-                    for &watching_generator_idx in watching_generator_indices {
-                        next_pending_generator_indices.push(watching_generator_idx);
-                    }
+            for &(watch, _) in &buffer.target_values {
+                for &watching_generator_idx in &generator_indices_by_watches[ind(watch)] {
+                    next_pending_generator_indices.push(watching_generator_idx);
                 }
             }
 

src/plonk/circuit_builder.rs

@@ -576,6 +576,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             gate_instances: self.gate_instances,
             public_inputs: self.public_inputs,
             marked_targets: self.marked_targets,
+            max_virtual_target_index: self.virtual_target_index,
         };
 
         // The HashSet of gates will have a non-deterministic order. When converting to a Vec, we

src/plonk/circuit_data.rs

@@ -136,6 +136,7 @@ pub(crate) struct ProverOnlyCircuitData<F: Extendable<D>, const D: usize> {
     pub public_inputs: Vec<Target>,
     /// A vector of marked targets. The values assigned to these targets will be displayed by the prover.
     pub marked_targets: Vec<MarkedTargets<D>>,
+    pub max_virtual_target_index: usize,
 }
 
 /// Circuit data required by the verifier, but not the prover.

src/plonk/prover.rs

@@ -37,7 +37,14 @@ pub(crate) fn prove<F: Extendable<D>, const D: usize>(
     timed!(
         timing,
         &format!("run {} generators", prover_data.generators.len()),
-        generate_partial_witness(&mut partial_witness, &prover_data.generators, &mut timing)
+        generate_partial_witness(
+            &mut partial_witness,
+            &prover_data.generators,
+            config.num_routed_wires,
+            degree,
+            prover_data.max_virtual_target_index,
+            &mut timing
+        )
     );
 
     let public_inputs = partial_witness.get_targets(&prover_data.public_inputs);