Enforce copy constraints in partial witness generation.

src/circuit_builder.rs

@@ -236,7 +236,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             .collect()
     }
 
-    fn sigma_vecs(&self, k_is: &[F]) -> Vec<PolynomialValues<F>> {
+    fn sigma_vecs(&self, k_is: &[F]) -> (Vec<PolynomialValues<F>>, TargetPartitions) {
         let degree = self.gate_instances.len();
         let degree_log = log2_strict(degree);
         let mut target_partitions = TargetPartitions::new();
@@ -256,7 +256,11 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
 
         let wire_partitions = target_partitions.to_wire_partitions();
-        wire_partitions.get_sigma_polys(degree_log, k_is)
+
+        (
+            wire_partitions.get_sigma_polys(degree_log, k_is),
+            target_partitions,
+        )
     }
 
     /// Builds a "full circuit", with both prover and verifier data.
@@ -278,7 +282,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         );
 
         let k_is = get_unique_coset_shifts(degree, self.config.num_routed_wires);
-        let sigma_vecs = self.sigma_vecs(&k_is);
+        let (sigma_vecs, targets_partition) = self.sigma_vecs(&k_is);
         let sigmas_commitment = ListPolynomialCommitment::new(
             sigma_vecs.into_iter().map(|v| v.ifft()).collect(),
             self.config.fri_config.rate_bits,
@@ -297,6 +301,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             generators,
             constants_commitment,
             sigmas_commitment,
+            targets_partition,
         };
 
         // The HashSet of gates will have a non-deterministic order. When converting to a Vec, we

src/circuit_data.rs

@@ -5,6 +5,7 @@ use crate::field::field::Field;
 use crate::fri::FriConfig;
 use crate::gates::gate::GateRef;
 use crate::generator::WitnessGenerator;
+use crate::permutation_argument::TargetPartitions;
 use crate::polynomial::commitment::ListPolynomialCommitment;
 use crate::proof::{Hash, HashTarget, Proof};
 use crate::prover::prove;
@@ -104,6 +105,8 @@ pub(crate) struct ProverOnlyCircuitData<F: Field> {
     pub constants_commitment: ListPolynomialCommitment<F>,
     /// Commitments to the sigma polynomial.
     pub sigmas_commitment: ListPolynomialCommitment<F>,
+    /// Partition of the targets into copy-constrained sets.
+    pub targets_partition: TargetPartitions,
 }
 
 /// Circuit data required by the verifier, but not the prover.

src/gates/gmimc.rs

@@ -323,6 +323,8 @@ mod tests {
     use crate::gates::gmimc::{GMiMCGate, W};
     use crate::generator::generate_partial_witness;
     use crate::gmimc::gmimc_permute_naive;
+    use crate::permutation_argument::TargetPartitions;
+    use crate::target::Target;
     use crate::wire::Wire;
     use crate::witness::PartialWitness;
 
@@ -368,7 +370,13 @@ mod tests {
         }
 
         let generators = gate.0.generators(0, &[]);
-        generate_partial_witness(&mut witness, &generators);
+        let mut tp = TargetPartitions::new();
+        for g in 0..10 {
+            for i in 0..config.num_routed_wires {
+                tp.add_partition(Target::wire(g, i));
+            }
+        }
+        generate_partial_witness(&mut witness, &generators, &tp);
 
         let expected_outputs: [F; W] =
             gmimc_permute_naive(permutation_inputs.try_into().unwrap(), constants);

src/generator.rs

@@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 
 use crate::field::field::Field;
+use crate::permutation_argument::TargetPartitions;
 use crate::target::Target;
 use crate::witness::PartialWitness;
 
@@ -10,6 +11,7 @@ use crate::witness::PartialWitness;
 pub(crate) fn generate_partial_witness<F: Field>(
     witness: &mut PartialWitness<F>,
     generators: &[Box<dyn WitnessGenerator<F>>],
+    target_partition: &TargetPartitions,
 ) {
     // Index generator indices by their watched targets.
     let mut generator_indices_by_watches = HashMap::new();
@@ -21,6 +23,8 @@ pub(crate) fn generate_partial_witness<F: Field>(
         }
     }
 
+    target_partition.generate_copies(witness, &witness.all_populated_targets());
+
     // Build a list of "pending" generators which are queued to be run. Initially, all generators
     // are queued.
     let mut pending_generator_indices: HashSet<_> = (0..generators.len()).collect();
@@ -33,11 +37,14 @@ pub(crate) fn generate_partial_witness<F: Field>(
         let mut next_pending_generator_indices = HashSet::new();
 
         for &generator_idx in &pending_generator_indices {
-            let (result, finished) = generators[generator_idx].run(&witness);
+            let (mut result, finished) = generators[generator_idx].run(&witness);
             if finished {
                 expired_generator_indices.insert(generator_idx);
             }
 
+            let new_targets = result.all_populated_targets();
+            target_partition.generate_copies(&mut result, &new_targets);
+
             // Enqueue unfinished generators that were watching one of the newly populated targets.
             for watch in result.target_values.keys() {
                 if let Some(watching_generator_indices) = generator_indices_by_watches.get(watch) {

src/permutation_argument.rs

@@ -6,6 +6,7 @@ use crate::field::field::Field;
 use crate::polynomial::polynomial::PolynomialValues;
 use crate::target::Target;
 use crate::wire::Wire;
+use crate::witness::PartialWitness;
 
 #[derive(Debug, Clone)]
 pub struct TargetPartitions {
@@ -82,6 +83,26 @@ impl TargetPartitions {
             indices,
         }
     }
+    /// For the given set of targets, find any copy constraints involving those targets and populate
+    /// the witness with copies as needed.
+    pub fn generate_copies<F: Field>(&self, witness: &mut PartialWitness<F>, targets: &[Target]) {
+        let mut result = PartialWitness::new();
+
+        for &target in targets {
+            let value = witness.get_target(target);
+            let partition = self.get_partition(target);
+
+            for &sibling in partition {
+                if witness.contains(sibling) {
+                    // This sibling's value was already set; make sure it has the same value.
+                    assert_eq!(witness.get_target(sibling), value);
+                } else {
+                    result.set_target(sibling, value);
+                }
+            }
+        }
+        witness.extend(result);
+    }
 }
 
 pub struct WirePartitions {

src/plonk_challenger.rs

@@ -289,6 +289,7 @@ mod tests {
     use crate::field::crandall_field::CrandallField;
     use crate::field::field::Field;
     use crate::generator::generate_partial_witness;
+    use crate::permutation_argument::TargetPartitions;
     use crate::plonk_challenger::{Challenger, RecursiveChallenger};
     use crate::target::Target;
     use crate::witness::PartialWitness;
@@ -337,7 +338,7 @@ mod tests {
 
         let config = CircuitConfig {
             num_wires: 12 + 12 + 3 + 101,
-            num_routed_wires: 27,
+            num_routed_wires: 200,
             ..CircuitConfig::default()
         };
         let mut builder = CircuitBuilder::<F, 4>::new(config);
@@ -351,7 +352,11 @@ mod tests {
         }
         let circuit = builder.build();
         let mut witness = PartialWitness::new();
-        generate_partial_witness(&mut witness, &circuit.prover_only.generators);
+        generate_partial_witness(
+            &mut witness,
+            &circuit.prover_only.generators,
+            &circuit.prover_only.targets_partition,
+        );
         let recursive_output_values_per_round: Vec<Vec<F>> = recursive_outputs_per_round
             .iter()
             .map(|outputs| witness.get_targets(outputs))

src/prover.rs

@@ -34,7 +34,11 @@ pub(crate) fn prove<F: Extendable<D>, const D: usize>(
     let mut witness = inputs;
     info!("Running {} generators", prover_data.generators.len());
     timed!(
-        generate_partial_witness(&mut witness, &prover_data.generators),
+        generate_partial_witness(
+            &mut witness,
+            &prover_data.generators,
+            &prover_data.targets_partition
+        ),
         "to generate witness"
     );
 

src/witness.rs

@@ -59,6 +59,10 @@ impl<F: Field> PartialWitness<F> {
         targets.iter().all(|&t| self.contains(t))
     }
 
+    pub fn all_populated_targets(&self) -> Vec<Target> {
+        self.target_values.keys().cloned().collect()
+    }
+
     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 {