Merge pull request #525 from mir-protocol/remove_const_gate

Generate constants in `RandomAccessGate`

plonky2/src/gates/constant.rs

@@ -1,5 +1,4 @@
 use plonky2_field::extension_field::Extendable;
-use plonky2_field::field_types::Field;
 use plonky2_field::packed_field::PackedField;
 
 use crate::gates::gate::Gate;
@@ -7,10 +6,7 @@ use crate::gates::packed_util::PackedEvaluableBase;
 use crate::gates::util::StridedConstraintConsumer;
 use crate::hash::hash_types::RichField;
 use crate::iop::ext_target::ExtensionTarget;
-use crate::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
-use crate::iop::target::Target;
-use crate::iop::wire::Wire;
-use crate::iop::witness::PartitionWitness;
+use crate::iop::generator::WitnessGenerator;
 use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::plonk::vars::{
     EvaluationTargets, EvaluationVars, EvaluationVarsBase, EvaluationVarsBaseBatch,
@@ -77,23 +73,10 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
 
     fn generators(
         &self,
-        gate_index: usize,
-        local_constants: &[F],
+        _gate_index: usize,
+        _local_constants: &[F],
     ) -> Vec<Box<dyn WitnessGenerator<F>>> {
-        (0..self.num_consts)
-            .map(|i| {
-                let g: Box<dyn WitnessGenerator<F>> = Box::new(
-                    ConstantGenerator {
-                        gate_index,
-                        gate: *self,
-                        i,
-                        constant: local_constants[self.const_input(i)],
-                    }
-                    .adapter(),
-                );
-                g
-            })
-            .collect()
+        vec![]
     }
 
     fn num_wires(&self) -> usize {
@@ -111,6 +94,12 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
     fn num_constraints(&self) -> usize {
         self.num_consts
     }
+
+    fn extra_constant_wires(&self) -> Vec<(usize, usize)> {
+        (0..self.num_consts)
+            .map(|i| (self.const_input(i), self.wire_output(i)))
+            .collect()
+    }
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> PackedEvaluableBase<F, D> for ConstantGate {
@@ -125,28 +114,6 @@ impl<F: RichField + Extendable<D>, const D: usize> PackedEvaluableBase<F, D> for
     }
 }
 
-#[derive(Debug)]
-struct ConstantGenerator<F: Field> {
-    gate_index: usize,
-    gate: ConstantGate,
-    i: usize,
-    constant: F,
-}
-
-impl<F: Field> SimpleGenerator<F> for ConstantGenerator<F> {
-    fn dependencies(&self) -> Vec<Target> {
-        Vec::new()
-    }
-
-    fn run_once(&self, _witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
-        let wire = Wire {
-            gate: self.gate_index,
-            input: self.gate.wire_output(self.i),
-        };
-        out_buffer.set_wire(wire, self.constant);
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
@@ -159,7 +126,7 @@ mod tests {
 
     #[test]
     fn low_degree() {
-        let num_consts = CircuitConfig::standard_recursion_config().constant_gate_size;
+        let num_consts = CircuitConfig::standard_recursion_config().num_constants;
         let gate = ConstantGate { num_consts };
         test_low_degree::<GoldilocksField, _, 2>(gate)
     }
@@ -169,7 +136,7 @@ mod tests {
         const D: usize = 2;
         type C = PoseidonGoldilocksConfig;
         type F = <C as GenericConfig<D>>::F;
-        let num_consts = CircuitConfig::standard_recursion_config().constant_gate_size;
+        let num_consts = CircuitConfig::standard_recursion_config().num_constants;
         let gate = ConstantGate { num_consts };
         test_eval_fns::<F, C, _, D>(gate)
     }

plonky2/src/gates/gate.rs

@@ -180,6 +180,15 @@ pub trait Gate<F: RichField + Extendable<D>, const D: usize>: 'static + Send + S
         self.generators(0, &vec![F::ZERO; self.num_constants()])
             .len()
     }
+
+    /// Enables gates to store some "routed constants", if they have both unused constants and
+    /// unused routed wires.
+    ///
+    /// Each entry in the returned `Vec` has the form `(constant_index, wire_index)`. `wire_index`
+    /// must correspond to a *routed* wire.
+    fn extra_constant_wires(&self) -> Vec<(usize, usize)> {
+        vec![]
+    }
 }
 
 /// A wrapper around an `Rc<Gate>` which implements `PartialEq`, `Eq` and `Hash` based on gate IDs.

plonky2/src/gates/random_access.rs

@@ -26,14 +26,16 @@ use crate::plonk::vars::{
 pub(crate) struct RandomAccessGate<F: RichField + Extendable<D>, const D: usize> {
     pub bits: usize,
     pub num_copies: usize,
+    pub num_extra_constants: usize,
     _phantom: PhantomData<F>,
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> RandomAccessGate<F, D> {
-    fn new(num_copies: usize, bits: usize) -> Self {
+    fn new(num_copies: usize, bits: usize, num_extra_constants: usize) -> Self {
         Self {
             bits,
             num_copies,
+            num_extra_constants,
             _phantom: PhantomData,
         }
     }
@@ -45,7 +47,12 @@ impl<F: RichField + Extendable<D>, const D: usize> RandomAccessGate<F, D> {
             // Need `(2 + vec_size + bits) * num_copies` wires
             config.num_wires / (2 + vec_size + bits),
         );
-        Self::new(max_copies, bits)
+        let max_extra_constants = config.num_routed_wires - (2 + vec_size) * max_copies;
+        Self::new(
+            max_copies,
+            bits,
+            max_extra_constants.min(config.num_constants),
+        )
     }
 
     fn vec_size(&self) -> usize {
@@ -68,12 +75,17 @@ impl<F: RichField + Extendable<D>, const D: usize> RandomAccessGate<F, D> {
         (2 + self.vec_size()) * copy + 2 + i
     }
 
-    fn start_of_intermediate_wires(&self) -> usize {
+    fn start_extra_constants(&self) -> usize {
         (2 + self.vec_size()) * self.num_copies
     }
 
-    pub(crate) fn num_routed_wires(&self) -> usize {
-        self.start_of_intermediate_wires()
+    fn wire_extra_constant(&self, i: usize) -> usize {
+        debug_assert!(i < self.num_extra_constants);
+        self.start_extra_constants() + i
+    }
+
+    pub fn num_routed_wires(&self) -> usize {
+        self.start_extra_constants() + self.num_extra_constants
     }
 
     /// An intermediate wire where the prover gives the (purported) binary decomposition of the
@@ -81,7 +93,7 @@ impl<F: RichField + Extendable<D>, const D: usize> RandomAccessGate<F, D> {
     pub fn wire_bit(&self, i: usize, copy: usize) -> usize {
         debug_assert!(i < self.bits);
         debug_assert!(copy < self.num_copies);
-        self.start_of_intermediate_wires() + copy * self.bits + i
+        self.num_routed_wires() + copy * self.bits + i
     }
 }
 
@@ -129,6 +141,11 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for RandomAccessGa
             constraints.push(list_items[0] - claimed_element);
         }
 
+        constraints.extend(
+            (0..self.num_extra_constants)
+                .map(|i| vars.local_constants[i] - vars.local_wires[self.wire_extra_constant(i)]),
+        );
+
         constraints
     }
 
@@ -189,6 +206,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for RandomAccessGa
             constraints.push(builder.sub_extension(list_items[0], claimed_element));
         }
 
+        constraints.extend((0..self.num_extra_constants).map(|i| {
+            builder.sub_extension(
+                vars.local_constants[i],
+                vars.local_wires[self.wire_extra_constant(i)],
+            )
+        }));
+
         constraints
     }
 
@@ -217,7 +241,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for RandomAccessGa
     }
 
     fn num_constants(&self) -> usize {
-        0
+        self.num_extra_constants
     }
 
     fn degree(&self) -> usize {
@@ -226,7 +250,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for RandomAccessGa
 
     fn num_constraints(&self) -> usize {
         let constraints_per_copy = self.bits + 2;
-        self.num_copies * constraints_per_copy
+        self.num_copies * constraints_per_copy + self.num_extra_constants
+    }
+
+    fn extra_constant_wires(&self) -> Vec<(usize, usize)> {
+        (0..self.num_extra_constants)
+            .map(|i| (i, self.wire_extra_constant(i)))
+            .collect()
     }
 }
 
@@ -270,6 +300,10 @@ impl<F: RichField + Extendable<D>, const D: usize> PackedEvaluableBase<F, D>
             debug_assert_eq!(list_items.len(), 1);
             yield_constr.one(list_items[0] - claimed_element);
         }
+        yield_constr.many(
+            (0..self.num_extra_constants)
+                .map(|i| vars.local_constants[i] - vars.local_wires[self.wire_extra_constant(i)]),
+        );
     }
 }
 
@@ -344,7 +378,7 @@ mod tests {
 
     #[test]
     fn low_degree() {
-        test_low_degree::<GoldilocksField, _, 4>(RandomAccessGate::new(4, 4));
+        test_low_degree::<GoldilocksField, _, 4>(RandomAccessGate::new(4, 4, 1));
     }
 
     #[test]
@@ -352,7 +386,7 @@ mod tests {
         const D: usize = 2;
         type C = PoseidonGoldilocksConfig;
         type F = <C as GenericConfig<D>>::F;
-        test_eval_fns::<F, C, _, D>(RandomAccessGate::new(4, 4))
+        test_eval_fns::<F, C, _, D>(RandomAccessGate::new(4, 4, 1))
     }
 
     #[test]
@@ -369,6 +403,7 @@ mod tests {
             lists: Vec<Vec<F>>,
             access_indices: Vec<usize>,
             claimed_elements: Vec<F>,
+            constants: &[F],
         ) -> Vec<FF> {
             let num_copies = lists.len();
             let vec_size = lists[0].len();
@@ -387,6 +422,7 @@ mod tests {
                     bit_vals.push(F::from_bool(((access_index >> i) & 1) != 0));
                 }
             }
+            v.extend(constants);
             v.extend(bit_vals);
 
             v.iter().map(|&x| x.into()).collect()
@@ -404,8 +440,10 @@ mod tests {
         let gate = RandomAccessGate::<F, D> {
             bits,
             num_copies,
+            num_extra_constants: 1,
             _phantom: PhantomData,
         };
+        let constants = F::rand_vec(gate.num_constants());
 
         let good_claimed_elements = lists
             .iter()
@@ -413,19 +451,26 @@ mod tests {
             .map(|(l, &i)| l[i])
             .collect();
         let good_vars = EvaluationVars {
-            local_constants: &[],
+            local_constants: &constants.iter().map(|&x| x.into()).collect::<Vec<_>>(),
             local_wires: &get_wires(
                 bits,
                 lists.clone(),
                 access_indices.clone(),
                 good_claimed_elements,
+                &constants,
             ),
             public_inputs_hash: &HashOut::rand(),
         };
         let bad_claimed_elements = F::rand_vec(4);
         let bad_vars = EvaluationVars {
-            local_constants: &[],
-            local_wires: &get_wires(bits, lists, access_indices, bad_claimed_elements),
+            local_constants: &constants.iter().map(|&x| x.into()).collect::<Vec<_>>(),
+            local_wires: &get_wires(
+                bits,
+                lists,
+                access_indices,
+                bad_claimed_elements,
+                &constants,
+            ),
             public_inputs_hash: &HashOut::rand(),
         };
 

plonky2/src/iop/generator.rs

@@ -307,3 +307,31 @@ impl<F: Field> SimpleGenerator<F> for NonzeroTestGenerator {
         out_buffer.set_target(self.dummy, dummy_value);
     }
 }
+
+/// Generator used to fill an extra constant.
+#[derive(Debug, Clone)]
+pub(crate) struct ConstantGenerator<F: Field> {
+    pub gate_index: usize,
+    pub constant_index: usize,
+    pub wire_index: usize,
+    pub constant: F,
+}
+
+impl<F: Field> ConstantGenerator<F> {
+    pub fn set_constant(&mut self, c: F) {
+        self.constant = c;
+    }
+}
+
+impl<F: Field> SimpleGenerator<F> for ConstantGenerator<F> {
+    fn dependencies(&self) -> Vec<Target> {
+        vec![]
+    }
+
+    fn run_once(&self, _witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
+        out_buffer.set_target(
+            Target::wire(self.gate_index, self.wire_index),
+            self.constant,
+        );
+    }
+}

plonky2/src/plonk/circuit_builder.rs

@@ -27,7 +27,7 @@ use crate::hash::hash_types::{HashOutTarget, MerkleCapTarget, RichField};
 use crate::hash::merkle_proofs::MerkleProofTarget;
 use crate::iop::ext_target::ExtensionTarget;
 use crate::iop::generator::{
-    CopyGenerator, RandomValueGenerator, SimpleGenerator, WitnessGenerator,
+    ConstantGenerator, CopyGenerator, RandomValueGenerator, SimpleGenerator, WitnessGenerator,
 };
 use crate::iop::target::{BoolTarget, Target};
 use crate::iop::wire::Wire;
@@ -83,6 +83,9 @@ pub struct CircuitBuilder<F: RichField + Extendable<D>, const D: usize> {
 
     /// Map between gate type and the current gate of this type with available slots.
     current_slots: HashMap<GateRef<F, D>, CurrentSlot<F, D>>,
+
+    /// List of constant generators used to fill the constant wires.
+    constant_generators: Vec<ConstantGenerator<F>>,
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
@@ -102,6 +105,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             arithmetic_results: HashMap::new(),
             targets_to_constants: HashMap::new(),
             current_slots: HashMap::new(),
+            constant_generators: Vec::new(),
         };
         builder.check_config();
         builder
@@ -206,15 +210,26 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     }
 
     /// Adds a gate to the circuit, and returns its index.
-    pub fn add_gate<G: Gate<F, D>>(&mut self, gate_type: G, constants: Vec<F>) -> usize {
+    pub fn add_gate<G: Gate<F, D>>(&mut self, gate_type: G, mut constants: Vec<F>) -> usize {
         self.check_gate_compatibility(&gate_type);
-        assert_eq!(
-            gate_type.num_constants(),
-            constants.len(),
-            "Number of constants doesn't match."
-        );
 
-        let index = self.gate_instances.len();
+        assert!(
+            constants.len() <= gate_type.num_constants(),
+            "Too many constants."
+        );
+        constants.resize(gate_type.num_constants(), F::ZERO);
+
+        let gate_index = self.gate_instances.len();
+
+        self.constant_generators
+            .extend(gate_type.extra_constant_wires().into_iter().map(
+                |(constant_index, wire_index)| ConstantGenerator {
+                    gate_index,
+                    constant_index,
+                    wire_index,
+                    constant: F::ZERO, // Placeholder; will be replaced later.
+                },
+            ));
 
         // Note that we can't immediately add this gate's generators, because the list of constants
         // could be modified later, i.e. in the case of `ConstantGate`. We will add them later in
@@ -229,7 +244,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             constants,
         });
 
-        index
+        gate_index
     }
 
     fn check_gate_compatibility<G: Gate<F, D>>(&self, gate: &G) {
@@ -240,6 +255,13 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             gate.num_wires(),
             self.config.num_wires
         );
+        assert!(
+            gate.num_constants() <= self.config.num_constants,
+            "{:?} requires {} constants, but our CircuitConfig has only {}",
+            gate.id(),
+            gate.num_constants(),
+            self.config.num_constants
+        );
     }
 
     pub fn connect_extension(&mut self, src: ExtensionTarget<D>, dst: ExtensionTarget<D>) {
@@ -321,14 +343,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             return target;
         }
 
-        let num_consts = self.config.constant_gate_size;
-        // We will fill this `ConstantGate` with zero constants initially.
-        // These will be overwritten by `constant` as the gate instances are filled.
-        let gate = ConstantGate { num_consts };
-        let (gate, instance) = self.find_slot(gate, &[], &vec![F::ZERO; num_consts]);
-        let target = Target::wire(gate, instance);
-        self.gate_instances[gate].constants[instance] = c;
-
+        let target = self.add_virtual_target();
         self.constants_to_targets.insert(c, target);
         self.targets_to_constants.insert(target, c);
 
@@ -650,6 +665,32 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             self.connect(hash_part, Target::wire(pi_gate, wire))
         }
 
+        // Make sure we have enough constant generators. If not, add a `ConstantGate`.
+        while self.constants_to_targets.len() > self.constant_generators.len() {
+            self.add_gate(
+                ConstantGate {
+                    num_consts: self.config.num_constants,
+                },
+                vec![],
+            );
+        }
+
+        // For each constant-target pair used in the circuit, use a constant generator to fill this target.
+        for ((c, t), mut const_gen) in self
+            .constants_to_targets
+            .clone()
+            .into_iter()
+            .zip(self.constant_generators.clone())
+        {
+            // Set the constant in the constant polynomial.
+            self.gate_instances[const_gen.gate_index].constants[const_gen.constant_index] = c;
+            // Generate a copy between the target and the routable wire.
+            self.connect(Target::wire(const_gen.gate_index, const_gen.wire_index), t);
+            // Set the constant in the generator (it's initially set with a dummy value).
+            const_gen.set_constant(c);
+            self.add_simple_generator(const_gen);
+        }
+
         info!(
             "Degree before blinding & padding: {}",
             self.gate_instances.len()

plonky2/src/plonk/circuit_data.rs

@@ -33,7 +33,7 @@ use crate::util::timing::TimingTree;
 pub struct CircuitConfig {
     pub num_wires: usize,
     pub num_routed_wires: usize,
-    pub constant_gate_size: usize,
+    pub num_constants: usize,
     /// Whether to use a dedicated gate for base field arithmetic, rather than using a single gate
     /// for both base field and extension field arithmetic.
     pub use_base_arithmetic_gate: bool,
@@ -64,7 +64,7 @@ impl CircuitConfig {
         Self {
             num_wires: 135,
             num_routed_wires: 80,
-            constant_gate_size: 5,
+            num_constants: 2,
             use_base_arithmetic_gate: true,
             security_bits: 100,
             num_challenges: 2,