Have ArithmeticExtensionGate adapt based on available wires (#287)

src/gadgets/arithmetic_extension.rs

@@ -4,7 +4,7 @@ use crate::field::extension_field::target::{ExtensionAlgebraTarget, ExtensionTar
 use crate::field::extension_field::FieldExtension;
 use crate::field::extension_field::{Extendable, OEF};
 use crate::field::field_types::{Field, RichField};
-use crate::gates::arithmetic::{ArithmeticExtensionGate, NUM_ARITHMETIC_OPS};
+use crate::gates::arithmetic::ArithmeticExtensionGate;
 use crate::iop::generator::{GeneratedValues, SimpleGenerator};
 use crate::iop::target::Target;
 use crate::iop::witness::{PartitionWitness, Witness};
@@ -21,12 +21,15 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             .get(&(const_0, const_1))
             .copied()
             .unwrap_or_else(|| {
-                let gate = self.add_gate(ArithmeticExtensionGate, vec![const_0, const_1]);
+                let gate = self.add_gate(
+                    ArithmeticExtensionGate::new_from_config(&self.config),
+                    vec![const_0, const_1],
+                );
                 (gate, 0)
             });
 
         // Update `free_arithmetic` with new values.
-        if i < NUM_ARITHMETIC_OPS - 1 {
+        if i < ArithmeticExtensionGate::<D>::num_ops(&self.config) - 1 {
             self.free_arithmetic
                 .insert((const_0, const_1), (gate, i + 1));
         } else {

src/gates/arithmetic.rs

@@ -9,16 +9,29 @@ use crate::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator};
 use crate::iop::target::Target;
 use crate::iop::witness::{PartitionWitness, Witness};
 use crate::plonk::circuit_builder::CircuitBuilder;
+use crate::plonk::circuit_data::CircuitConfig;
 use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 
-/// Number of arithmetic operations performed by an arithmetic gate.
-pub const NUM_ARITHMETIC_OPS: usize = 4;
-
 /// A gate which can a linear combination `c0*x*y+c1*z` twice with the same `x`.
 #[derive(Debug)]
-pub struct ArithmeticExtensionGate<const D: usize>;
+pub struct ArithmeticExtensionGate<const D: usize> {
+    /// Number of arithmetic operations performed by an arithmetic gate.
+    pub num_ops: usize,
+}
 
 impl<const D: usize> ArithmeticExtensionGate<D> {
+    pub fn new_from_config(config: &CircuitConfig) -> Self {
+        Self {
+            num_ops: Self::num_ops(config),
+        }
+    }
+
+    /// Determine the maximum number of operations that can fit in one gate for the given config.
+    pub(crate) fn num_ops(config: &CircuitConfig) -> usize {
+        let wires_per_op = 4 * D;
+        config.num_routed_wires / wires_per_op
+    }
+
     pub fn wires_ith_multiplicand_0(i: usize) -> Range<usize> {
         4 * D * i..4 * D * i + D
     }
@@ -43,7 +56,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
         let const_1 = vars.local_constants[1];
 
         let mut constraints = Vec::new();
-        for i in 0..NUM_ARITHMETIC_OPS {
+        for i in 0..self.num_ops {
             let multiplicand_0 = vars.get_local_ext_algebra(Self::wires_ith_multiplicand_0(i));
             let multiplicand_1 = vars.get_local_ext_algebra(Self::wires_ith_multiplicand_1(i));
             let addend = vars.get_local_ext_algebra(Self::wires_ith_addend(i));
@@ -62,7 +75,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
         let const_1 = vars.local_constants[1];
 
         let mut constraints = Vec::new();
-        for i in 0..NUM_ARITHMETIC_OPS {
+        for i in 0..self.num_ops {
             let multiplicand_0 = vars.get_local_ext(Self::wires_ith_multiplicand_0(i));
             let multiplicand_1 = vars.get_local_ext(Self::wires_ith_multiplicand_1(i));
             let addend = vars.get_local_ext(Self::wires_ith_addend(i));
@@ -85,7 +98,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
         let const_1 = vars.local_constants[1];
 
         let mut constraints = Vec::new();
-        for i in 0..NUM_ARITHMETIC_OPS {
+        for i in 0..self.num_ops {
             let multiplicand_0 = vars.get_local_ext_algebra(Self::wires_ith_multiplicand_0(i));
             let multiplicand_1 = vars.get_local_ext_algebra(Self::wires_ith_multiplicand_1(i));
             let addend = vars.get_local_ext_algebra(Self::wires_ith_addend(i));
@@ -109,7 +122,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
         gate_index: usize,
         local_constants: &[F],
     ) -> Vec<Box<dyn WitnessGenerator<F>>> {
-        (0..NUM_ARITHMETIC_OPS)
+        (0..self.num_ops)
             .map(|i| {
                 let g: Box<dyn WitnessGenerator<F>> = Box::new(
                     ArithmeticExtensionGenerator {
@@ -126,7 +139,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
     }
 
     fn num_wires(&self) -> usize {
-        NUM_ARITHMETIC_OPS * 4 * D
+        self.num_ops * 4 * D
     }
 
     fn num_constants(&self) -> usize {
@@ -138,7 +151,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for ArithmeticExte
     }
 
     fn num_constraints(&self) -> usize {
-        NUM_ARITHMETIC_OPS * D
+        self.num_ops * D
     }
 }
 
@@ -196,13 +209,17 @@ mod tests {
     use crate::field::crandall_field::CrandallField;
     use crate::gates::arithmetic::ArithmeticExtensionGate;
     use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
+    use crate::plonk::circuit_data::CircuitConfig;
 
     #[test]
     fn low_degree() {
-        test_low_degree::<CrandallField, _, 4>(ArithmeticExtensionGate)
+        let gate = ArithmeticExtensionGate::new_from_config(&CircuitConfig::large_config());
+        test_low_degree::<CrandallField, _, 4>(gate);
     }
+
     #[test]
     fn eval_fns() -> Result<()> {
-        test_eval_fns::<CrandallField, _, 4>(ArithmeticExtensionGate)
+        let gate = ArithmeticExtensionGate::new_from_config(&CircuitConfig::large_config());
+        test_eval_fns::<CrandallField, _, 4>(gate)
     }
 }

src/gates/gate_tree.rs

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

src/plonk/circuit_builder.rs

@@ -12,7 +12,7 @@ use crate::field::fft::fft_root_table;
 use crate::field::field_types::RichField;
 use crate::fri::commitment::PolynomialBatchCommitment;
 use crate::fri::FriParams;
-use crate::gates::arithmetic::{ArithmeticExtensionGate, NUM_ARITHMETIC_OPS};
+use crate::gates::arithmetic::ArithmeticExtensionGate;
 use crate::gates::constant::ConstantGate;
 use crate::gates::gate::{Gate, GateInstance, GateRef, PrefixedGate};
 use crate::gates::gate_tree::Tree;
@@ -509,7 +509,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let zero = self.zero_extension();
         let remaining_arithmetic_gates = self.free_arithmetic.values().copied().collect::<Vec<_>>();
         for (gate, i) in remaining_arithmetic_gates {
-            for j in i..NUM_ARITHMETIC_OPS {
+            for j in i..ArithmeticExtensionGate::<D>::num_ops(&self.config) {
                 let wires_multiplicand_0 = ExtensionTarget::from_range(
                     gate,
                     ArithmeticExtensionGate::<D>::wires_ith_multiplicand_0(j),