Fix type errors and move copy constraints check to Witness

src/prover.rs

@@ -23,6 +23,11 @@ pub(crate) fn prove<F: Extendable<D>, const D: usize>(
     inputs: PartialWitness<F>,
 ) -> Proof<F, D> {
     let fri_config = &common_data.config.fri_config;
+    let config = &common_data.config;
+    let num_wires = config.num_wires;
+    let num_challenges = config.num_challenges;
+    let quotient_degree = common_data.quotient_degree();
+    let degree = common_data.degree();
 
     let start_proof_gen = Instant::now();
 
@@ -33,6 +38,11 @@ pub(crate) fn prove<F: Extendable<D>, const D: usize>(
         "to generate witness"
     );
 
+    let witness = timed!(
+        partial_witness.full_witness(degree, num_wires),
+        "to compute full witness"
+    );
+
     timed!(
         witness
             .check_copy_constraints(&prover_data.copy_constraints, &prover_data.gate_instances)
@@ -40,14 +50,6 @@ pub(crate) fn prove<F: Extendable<D>, const D: usize>(
         "to check copy constraints"
     );
 
-    let config = &common_data.config;
-    let num_wires = config.num_wires;
-    let num_challenges = config.num_challenges;
-    let quotient_degree = common_data.quotient_degree();
-    let degree = common_data.degree();
-
-    let witness = partial_witness.full_witness(degree, num_wires);
-
     let wires_values: Vec<PolynomialValues<F>> = timed!(
         witness
             .wire_values
@@ -169,7 +171,7 @@ fn compute_zs<F: Extendable<D>, const D: usize>(
     witness: &Witness<F>,
     betas: &[F],
     gammas: &[F],
-    prover_data: &ProverOnlyCircuitData<F>,
+    prover_data: &ProverOnlyCircuitData<F, D>,
     common_data: &CommonCircuitData<F, D>,
 ) -> Vec<PolynomialValues<F>> {
     (0..common_data.config.num_challenges)
@@ -181,7 +183,7 @@ fn compute_z<F: Extendable<D>, const D: usize>(
     witness: &Witness<F>,
     beta: F,
     gamma: F,
-    prover_data: &ProverOnlyCircuitData<F>,
+    prover_data: &ProverOnlyCircuitData<F, D>,
     common_data: &CommonCircuitData<F, D>,
 ) -> PolynomialValues<F> {
     let subgroup = &prover_data.subgroup;

src/witness.rs

@@ -19,6 +19,41 @@ impl<F: Field> Witness<F> {
     pub fn get_wire(&self, gate: usize, input: usize) -> F {
         self.wire_values[input][gate]
     }
+
+    /// Checks that the copy constraints are satisfied in the witness.
+    pub fn check_copy_constraints<const D: usize>(
+        &self,
+        copy_constraints: &[(Target, Target)],
+        gate_instances: &[GateInstance<F, D>],
+    ) -> Result<()>
+    where
+        F: Extendable<D>,
+    {
+        for &(a, b) in copy_constraints {
+            // TODO: Take care of public inputs once they land.
+            if let (
+                Target::Wire(Wire {
+                    gate: a_gate,
+                    input: a_input,
+                }),
+                Target::Wire(Wire {
+                    gate: b_gate,
+                    input: b_input,
+                }),
+            ) = (a, b)
+            {
+                let va = self.get_wire(a_gate, a_input);
+                let vb = self.get_wire(b_gate, b_input);
+                ensure!(
+                    va == vb,
+                    "Copy constraint between wire {} of gate #{} (`{}`) and wire {} of gate #{} (`{}`) is not satisfied. \
+                    Got values of {} and {} respectively.",
+                    a_input, a_gate, gate_instances[a_gate].gate_type.0.id(), b_input, b_gate,
+                    gate_instances[b_gate].gate_type.0.id(), va, vb);
+            }
+        }
+        Ok(())
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -145,31 +180,6 @@ impl<F: Field> PartialWitness<F> {
         });
         Witness { wire_values }
     }
-
-    /// Checks that the copy constraints are satisfied in the witness.
-    pub fn check_copy_constraints<const D: usize>(
-        &self,
-        copy_constraints: &[(Target, Target)],
-        gate_instances: &[GateInstance<F, D>],
-    ) -> Result<()>
-    where
-        F: Extendable<D>,
-    {
-        for &(a, b) in copy_constraints {
-            // TODO: Take care of public inputs once they land.
-            if let (Target::Wire(wa), Target::Wire(wb)) = (a, b) {
-                let va = self.target_values.get(&a).copied().unwrap_or(F::ZERO);
-                let vb = self.target_values.get(&b).copied().unwrap_or(F::ZERO);
-                ensure!(
-                    va == vb,
-                    "Copy constraint between wire {} of gate #{} (`{}`) and wire {} of gate #{} (`{}`) is not satisfied. \
-                    Got values of {} and {} respectively.",
-                    wa.input, wa.gate, gate_instances[wa.gate].gate_type.0.id(), wb.input, wb.gate,
-                    gate_instances[wb.gate].gate_type.0.id(), va, vb);
-            }
-        }
-        Ok(())
-    }
 }
 
 impl<F: Field> Default for PartialWitness<F> {