More wires for ConstantGate (#332)

* More wires for ConstantGate

* fix

* fix

src/bin/bench_recursion.rs

@@ -25,6 +25,7 @@ fn bench_prove<F: RichField + Extendable<D>, const D: usize>() -> Result<()> {
     let config = CircuitConfig {
         num_wires: 126,
         num_routed_wires: 33,
+        constant_gate_size: 6,
         security_bits: 128,
         rate_bits: 3,
         num_challenges: 3,

src/gates/constant.rs

@@ -1,3 +1,5 @@
+use std::ops::Range;
+
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
 use crate::field::field_types::{Field, RichField};
@@ -10,29 +12,38 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 
 /// A gate which takes a single constant parameter and outputs that value.
-pub struct ConstantGate;
+#[derive(Copy, Clone, Debug)]
+pub struct ConstantGate {
+    pub(crate) num_consts: usize,
+}
 
 impl ConstantGate {
-    pub const CONST_INPUT: usize = 0;
+    pub fn consts_inputs(&self) -> Range<usize> {
+        0..self.num_consts
+    }
 
-    pub const WIRE_OUTPUT: usize = 0;
+    pub fn wires_outputs(&self) -> Range<usize> {
+        0..self.num_consts
+    }
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
     fn id(&self) -> String {
-        "ConstantGate".into()
+        format!("{:?}", self)
     }
 
     fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
-        let input = vars.local_constants[Self::CONST_INPUT];
-        let output = vars.local_wires[Self::WIRE_OUTPUT];
-        vec![output - input]
+        self.consts_inputs()
+            .zip(self.wires_outputs())
+            .map(|(con, out)| vars.local_constants[con] - vars.local_wires[out])
+            .collect()
     }
 
     fn eval_unfiltered_base(&self, vars: EvaluationVarsBase<F>) -> Vec<F> {
-        let input = vars.local_constants[Self::CONST_INPUT];
-        let output = vars.local_wires[Self::WIRE_OUTPUT];
-        vec![output - input]
+        self.consts_inputs()
+            .zip(self.wires_outputs())
+            .map(|(con, out)| vars.local_constants[con] - vars.local_wires[out])
+            .collect()
     }
 
     fn eval_unfiltered_recursively(
@@ -40,9 +51,12 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
         builder: &mut CircuitBuilder<F, D>,
         vars: EvaluationTargets<D>,
     ) -> Vec<ExtensionTarget<D>> {
-        let input = vars.local_constants[Self::CONST_INPUT];
-        let output = vars.local_wires[Self::WIRE_OUTPUT];
-        vec![builder.sub_extension(output, input)]
+        self.consts_inputs()
+            .zip(self.wires_outputs())
+            .map(|(con, out)| {
+                builder.sub_extension(vars.local_constants[con], vars.local_wires[out])
+            })
+            .collect()
     }
 
     fn generators(
@@ -52,17 +66,18 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
     ) -> Vec<Box<dyn WitnessGenerator<F>>> {
         let gen = ConstantGenerator {
             gate_index,
-            constant: local_constants[0],
+            gate: *self,
+            constants: local_constants[self.consts_inputs()].to_vec(),
         };
         vec![Box::new(gen.adapter())]
     }
 
     fn num_wires(&self) -> usize {
-        1
+        self.num_consts
     }
 
     fn num_constants(&self) -> usize {
-        1
+        self.num_consts
     }
 
     fn degree(&self) -> usize {
@@ -70,14 +85,15 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ConstantGate {
     }
 
     fn num_constraints(&self) -> usize {
-        1
+        self.num_consts
     }
 }
 
 #[derive(Debug)]
 struct ConstantGenerator<F: Field> {
     gate_index: usize,
-    constant: F,
+    gate: ConstantGate,
+    constants: Vec<F>,
 }
 
 impl<F: Field> SimpleGenerator<F> for ConstantGenerator<F> {
@@ -86,11 +102,13 @@ impl<F: Field> SimpleGenerator<F> for ConstantGenerator<F> {
     }
 
     fn run_once(&self, _witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
-        let wire = Wire {
-            gate: self.gate_index,
-            input: ConstantGate::WIRE_OUTPUT,
-        };
-        out_buffer.set_target(Target::Wire(wire), self.constant);
+        for (con, out) in self.gate.consts_inputs().zip(self.gate.wires_outputs()) {
+            let wire = Wire {
+                gate: self.gate_index,
+                input: out,
+            };
+            out_buffer.set_wire(wire, self.constants[con]);
+        }
     }
 }
 
@@ -101,14 +119,19 @@ mod tests {
     use crate::field::goldilocks_field::GoldilocksField;
     use crate::gates::constant::ConstantGate;
     use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
+    use crate::plonk::circuit_data::CircuitConfig;
 
     #[test]
     fn low_degree() {
-        test_low_degree::<GoldilocksField, _, 4>(ConstantGate)
+        let num_consts = CircuitConfig::large_config().constant_gate_size;
+        let gate = ConstantGate { num_consts };
+        test_low_degree::<GoldilocksField, _, 2>(gate)
     }
 
     #[test]
     fn eval_fns() -> Result<()> {
-        test_eval_fns::<GoldilocksField, _, 4>(ConstantGate)
+        let num_consts = CircuitConfig::large_config().constant_gate_size;
+        let gate = ConstantGate { num_consts };
+        test_eval_fns::<GoldilocksField, _, 2>(gate)
     }
 }

src/gates/gate_tree.rs

@@ -238,7 +238,7 @@ mod tests {
 
         let gates = vec![
             GateRef::new(NoopGate),
-            GateRef::new(ConstantGate),
+            GateRef::new(ConstantGate { num_consts: 4 }),
             GateRef::new(ArithmeticExtensionGate { num_ops: 4 }),
             GateRef::new(BaseSumGate::<4>::new(4)),
             GateRef::new(GMiMCGate::<F, D, 12>::new()),

src/plonk/circuit_builder.rs

@@ -82,6 +82,9 @@ pub struct CircuitBuilder<F: RichField + Extendable<D>, const D: usize> {
     // chunk_size, and contains `(g, i, c)`, if the gate `g`, at index `i`, already contains `c` copies
     // of switches
     pub(crate) current_switch_gates: Vec<Option<(SwitchGate<F, D>, usize, usize)>>,
+
+    /// An available `ConstantGate` instance, if any.
+    free_constant: Option<(usize, usize)>,
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
@@ -101,6 +104,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             free_arithmetic: HashMap::new(),
             free_random_access: HashMap::new(),
             current_switch_gates: Vec::new(),
+            free_constant: None,
         };
         builder.check_config();
         builder
@@ -194,7 +198,10 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         );
 
         let index = self.gate_instances.len();
-        self.add_generators(gate_type.generators(index, &constants));
+
+        // 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
+        // `build` instead.
 
         // Register this gate type if we haven't seen it before.
         let gate_ref = GateRef::new(gate_type);
@@ -291,16 +298,36 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             return target;
         }
 
-        let gate = self.add_gate(ConstantGate, vec![c]);
-        let target = Target::Wire(Wire {
-            gate,
-            input: ConstantGate::WIRE_OUTPUT,
-        });
+        let (gate, instance) = self.constant_gate_instance();
+        let target = Target::wire(gate, instance);
+        self.gate_instances[gate].constants[instance] = c;
+
         self.constants_to_targets.insert(c, target);
         self.targets_to_constants.insert(target, c);
+
         target
     }
 
+    /// Returns the gate index and copy index of a free `ConstantGate` slot, potentially adding a
+    /// new `ConstantGate` if needed.
+    fn constant_gate_instance(&mut self) -> (usize, usize) {
+        if self.free_constant.is_none() {
+            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 = self.add_gate(ConstantGate { num_consts }, vec![F::ZERO; num_consts]);
+            self.free_constant = Some((gate, 0));
+        }
+
+        let (gate, instance) = self.free_constant.unwrap();
+        if instance + 1 < self.config.constant_gate_size {
+            self.free_constant = Some((gate, instance + 1));
+        } else {
+            self.free_constant = None;
+        }
+        (gate, instance)
+    }
+
     pub fn constants(&mut self, constants: &[F]) -> Vec<Target> {
         constants.iter().map(|&c| self.constant(c)).collect()
     }
@@ -695,6 +722,15 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             constants_sigmas_cap: constants_sigmas_cap.clone(),
         };
 
+        // Add gate generators.
+        self.add_generators(
+            self.gate_instances
+                .iter()
+                .enumerate()
+                .flat_map(|(index, gate)| gate.gate_ref.0.generators(index, &gate.constants))
+                .collect(),
+        );
+
         // Index generator indices by their watched targets.
         let mut generator_indices_by_watches = BTreeMap::new();
         for (i, generator) in self.generators.iter().enumerate() {

src/plonk/circuit_data.rs

@@ -25,6 +25,7 @@ use crate::util::timing::TimingTree;
 pub struct CircuitConfig {
     pub num_wires: usize,
     pub num_routed_wires: usize,
+    pub constant_gate_size: usize,
     pub security_bits: usize,
     pub rate_bits: usize,
     /// The number of challenge points to generate, for IOPs that have soundness errors of (roughly)
@@ -53,6 +54,7 @@ impl CircuitConfig {
         Self {
             num_wires: 143,
             num_routed_wires: 28,
+            constant_gate_size: 6,
             security_bits: 100,
             rate_bits: 3,
             num_challenges: 2,
@@ -71,6 +73,7 @@ impl CircuitConfig {
         Self {
             num_wires: 143,
             num_routed_wires: 64,
+            constant_gate_size: 6,
             security_bits: 4,
             rate_bits: 3,
             num_challenges: 3,

src/plonk/proof.rs

@@ -321,7 +321,7 @@ mod tests {
         const D: usize = 4;
 
         let mut config = CircuitConfig::large_config();
-        config.fri_config.reduction_strategy = FriReductionStrategy::Fixed(vec![2, 1]);
+        config.fri_config.reduction_strategy = FriReductionStrategy::Fixed(vec![1, 1]);
         config.fri_config.num_query_rounds = 50;
 
         let pw = PartialWitness::new();