Fix witness generation performance (#269)

See #266. This avoids the quadratic costs (w.r.t. partition size), as we will now only enumerate watchers the first time a representative is assigned.

src/iop/generator.rs

@@ -49,11 +49,13 @@ pub(crate) fn generate_partial_witness<F: RichField + Extendable<D>, const D: us
                 remaining_generators -= 1;
             }
 
+            // Merge any generated values into our witness, and get a list of newly-populated
+            // targets' representatives.
+            let new_target_reps = witness.extend_returning_parents(buffer.target_values.drain(..));
+
             // Enqueue unfinished generators that were watching one of the newly populated targets.
-            for &(watch, _) in &buffer.target_values {
-                let watch_index = witness.target_index(watch);
-                let watch_rep_index = witness.forest[watch_index].parent;
-                let opt_watchers = generator_indices_by_watches.get(&watch_rep_index);
+            for watch in new_target_reps {
+                let opt_watchers = generator_indices_by_watches.get(&watch);
                 if let Some(watchers) = opt_watchers {
                     for &watching_generator_idx in watchers {
                         if !generator_is_expired[watching_generator_idx] {
@@ -62,8 +64,6 @@ pub(crate) fn generate_partial_witness<F: RichField + Extendable<D>, const D: us
                     }
                 }
             }
-
-            witness.extend(buffer.target_values.drain(..));
         }
 
         pending_generator_indices = next_pending_generator_indices;

src/iop/witness.rs

@@ -203,6 +203,14 @@ impl<F: Field> Witness<F> for PartitionWitness<F> {
     }
 
     fn set_target(&mut self, target: Target, value: F) {
+        self.set_target_returning_parent(target, value);
+    }
+}
+
+impl<F: Field> PartitionWitness<F> {
+    /// Set a `Target`. On success, returns the representative index of the newly-set target. If the
+    /// target was already set, returns `None`.
+    fn set_target_returning_parent(&mut self, target: Target, value: F) -> Option<usize> {
         let parent_index = self.forest[self.target_index(target)].parent;
         let parent_value = &mut self.forest[parent_index].value;
         if let Some(old_value) = *parent_value {
@@ -211,13 +219,21 @@ impl<F: Field> Witness<F> for PartitionWitness<F> {
                 "Partition containing {:?} was set twice with different values",
                 target
             );
+            None
         } else {
             *parent_value = Some(value);
+            Some(parent_index)
         }
     }
-}
 
-impl<F: Field> PartitionWitness<F> {
+    /// Returns the representative indices of any newly-set targets.
+    pub(crate) fn extend_returning_parents<'a, I: 'a + Iterator<Item = (Target, F)>>(
+        &'a mut self,
+        pairs: I,
+    ) -> impl Iterator<Item = usize> + 'a {
+        pairs.flat_map(move |(t, v)| self.set_target_returning_parent(t, v))
+    }
+
     pub fn target_index(&self, target: Target) -> usize {
         match target {
             Target::Wire(Wire { gate, input }) => gate * self.num_wires + input,