merge

src/gadgets/arithmetic_extension.rs

@@ -12,33 +12,6 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::util::bits_u64;
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
-    /// Finds the last available arithmetic gate with the given constants or add one if there aren't any.
-    /// Returns `(g,i)` such that there is an arithmetic gate with the given constants at index
-    /// `g` and the gate's `i`-th operation is available.
-    fn find_arithmetic_gate(&mut self, const_0: F, const_1: F) -> (usize, usize) {
-        let (gate, i) = self
-            .free_arithmetic
-            .get(&(const_0, const_1))
-            .copied()
-            .unwrap_or_else(|| {
-                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 < ArithmeticExtensionGate::<D>::num_ops(&self.config) - 1 {
-            self.free_arithmetic
-                .insert((const_0, const_1), (gate, i + 1));
-        } else {
-            self.free_arithmetic.remove(&(const_0, const_1));
-        }
-
-        (gate, i)
-    }
-
     pub fn arithmetic_extension(
         &mut self,
         const_0: F,

src/gadgets/arithmetic_u32.rs

@@ -36,14 +36,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         y: U32Target,
         z: U32Target,
     ) -> (U32Target, U32Target) {
-        let (gate_index, copy) = match self.current_u32_arithmetic_gate {
-            None => {
-                let gate = U32ArithmeticGate::new();
-                let gate_index = self.add_gate(gate, vec![]);
-                (gate_index, 0)
-            }
-            Some((gate_index, copy)) => (gate_index, copy),
-        };
+        let (gate_index, copy) = self.find_u32_arithmetic_gate();
 
         self.connect(
             Target::wire(
@@ -73,12 +66,6 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             U32ArithmeticGate::<F, D>::wire_ith_output_high_half(copy),
         ));
 
-        if copy == NUM_U32_ARITHMETIC_OPS - 1 {
-            self.current_u32_arithmetic_gate = None;
-        } else {
-            self.current_u32_arithmetic_gate = Some((gate_index, copy + 1));
-        }
-
         (output_low, output_high)
     }
 
@@ -103,4 +90,44 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let zero = self.zero_u32();
         self.mul_add_u32(a, b, zero)
     }
+
+    // Returns x * y + z.
+    pub fn sub_u32(
+        &mut self,
+        x: U32Target,
+        y: U32Target,
+        borrow: U32Target,
+    ) -> (U32Target, U32Target) {
+        let (gate_index, copy) = self.find_u32_subtraction_gate();
+
+        self.connect(
+            Target::wire(
+                gate_index,
+                U32ArithmeticGate::<F, D>::wire_ith_multiplicand_0(copy),
+            ),
+            x.0,
+        );
+        self.connect(
+            Target::wire(
+                gate_index,
+                U32ArithmeticGate::<F, D>::wire_ith_multiplicand_1(copy),
+            ),
+            y.0,
+        );
+        self.connect(
+            Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_addend(copy)),
+            z.0,
+        );
+
+        let output_low = U32Target(Target::wire(
+            gate_index,
+            U32ArithmeticGate::<F, D>::wire_ith_output_low_half(copy),
+        ));
+        let output_high = U32Target(Target::wire(
+            gate_index,
+            U32ArithmeticGate::<F, D>::wire_ith_output_high_half(copy),
+        ));
+
+        (output_low, output_high)
+    }
 }

src/gadgets/permutation.rs

@@ -73,20 +73,8 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
         let chunk_size = a1.len();
 
-        if self.current_switch_gates.len() < chunk_size {
-            self.current_switch_gates
-                .extend(vec![None; chunk_size - self.current_switch_gates.len()]);
-        }
-
         let (gate, gate_index, mut next_copy) =
-            match self.current_switch_gates[chunk_size - 1].clone() {
-                None => {
-                    let gate = SwitchGate::<F, D>::new_from_config(&self.config, chunk_size);
-                    let gate_index = self.add_gate(gate.clone(), vec![]);
-                    (gate, gate_index, 0)
-                }
-                Some((gate, idx, next_copy)) => (gate, idx, next_copy),
-            };
+            self.find_switch_gate(chunk_size);
 
         let num_copies = gate.num_copies;
 
@@ -113,13 +101,6 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
         let switch = Target::wire(gate_index, gate.wire_switch_bool(next_copy));
 
-        next_copy += 1;
-        if next_copy == num_copies {
-            self.current_switch_gates[chunk_size - 1] = None;
-        } else {
-            self.current_switch_gates[chunk_size - 1] = Some((gate, gate_index, next_copy));
-        }
-
         (switch, c, d)
     }
 

src/plonk/circuit_builder.rs

@@ -15,6 +15,8 @@ use crate::fri::{FriConfig, FriParams};
 use crate::gadgets::arithmetic_extension::ArithmeticOperation;
 use crate::gadgets::arithmetic_u32::U32Target;
 use crate::gates::arithmetic::ArithmeticExtensionGate;
+use crate::gates::arithmetic_u32::{NUM_U32_ARITHMETIC_OPS, U32ArithmeticGate};
+use crate::gates::subtraction_u32::{NUM_U32_SUBTRACTION_OPS, U32SubtractionGate};
 use crate::gates::constant::ConstantGate;
 use crate::gates::gate::{Gate, GateInstance, GateRef, PrefixedGate};
 use crate::gates::gate_tree::Tree;
@@ -75,24 +77,7 @@ pub struct CircuitBuilder<F: RichField + Extendable<D>, const D: usize> {
     /// Memoized results of `arithmetic_extension` calls.
     pub(crate) arithmetic_results: HashMap<ArithmeticOperation<F, D>, ExtensionTarget<D>>,
 
-    /// A map `(c0, c1) -> (g, i)` from constants `(c0,c1)` to an available arithmetic gate using
-    /// these constants with gate index `g` and already using `i` arithmetic operations.
-    pub(crate) free_arithmetic: HashMap<(F, F), (usize, usize)>,
-
-    /// A map `(c0, c1) -> (g, i)` from constants `vec_size` to an available arithmetic gate using
-    /// these constants with gate index `g` and already using `i` random accesses.
-    pub(crate) free_random_access: HashMap<usize, (usize, usize)>,
-
-    // `current_switch_gates[chunk_size - 1]` contains None if we have no switch gates with the value
-    // 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)>>,
-
-    /// The `U32ArithmeticGate` currently being filled (so new u32 arithmetic operations will be added to this gate before creating a new one)
-    pub(crate) current_u32_arithmetic_gate: Option<(usize, usize)>,
-
-    /// An available `ConstantGate` instance, if any.
-    free_constant: Option<(usize, usize)>,
+    batched_gates: BatchedGates<F, D>
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
@@ -110,11 +95,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             constants_to_targets: HashMap::new(),
             arithmetic_results: HashMap::new(),
             targets_to_constants: HashMap::new(),
-            free_arithmetic: HashMap::new(),
-            free_random_access: HashMap::new(),
-            current_switch_gates: Vec::new(),
-            current_u32_arithmetic_gate: None,
-            free_constant: None,
+            batched_gates: BatchedGates::new(),
         };
         builder.check_config();
         builder
@@ -308,7 +289,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             return target;
         }
 
-        let (gate, instance) = self.constant_gate_instance();
+        let (gate, instance) = self.batched_gates.constant_gate_instance();
         let target = Target::wire(gate, instance);
         self.gate_instances[gate].constants[instance] = c;
 
@@ -318,26 +299,6 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         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()
     }
@@ -846,3 +807,152 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         }
     }
 }
+
+/// 
+pub struct BatchedGates<F: RichField + Extendable<D>, const D: usize> {
+    /// A map `(c0, c1) -> (g, i)` from constants `(c0,c1)` to an available arithmetic gate using
+    /// these constants with gate index `g` and already using `i` arithmetic operations.
+    pub(crate) free_arithmetic: HashMap<(F, F), (usize, usize)>,
+
+    /// `current_switch_gates[chunk_size - 1]` contains None if we have no switch gates with the value
+    /// 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)>>,
+
+    /// The `U32ArithmeticGate` currently being filled (so new u32 arithmetic operations will be added to this gate before creating a new one)
+    pub(crate) current_u32_arithmetic_gate: Option<(usize, usize)>,
+
+    /// The `U32SubtractionGate` currently being filled (so new u32 subtraction operations will be added to this gate before creating a new one)
+    pub(crate) current_u32_subtraction_gate: Option<(usize, usize)>,
+
+    /// An available `ConstantGate` instance, if any.
+    pub(crate) free_constant: Option<(usize, usize)>,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> BatchedGates<F, D> {
+    pub fn new() -> Self {
+        Self {
+            free_arithmetic: HashMap::new(),
+            current_switch_gates: Vec::new(),
+            current_u32_arithmetic_gate: None,
+            current_u32_subtraction_gate: None,
+            free_constant: None,
+        }
+    }
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
+    /// Finds the last available arithmetic gate with the given constants or add one if there aren't any.
+    /// Returns `(g,i)` such that there is an arithmetic gate with the given constants at index
+    /// `g` and the gate's `i`-th operation is available.
+    pub fn find_arithmetic_gate(&mut self, const_0: F, const_1: F) -> (usize, usize) {
+        let (gate, i) = self
+            .batched_gates
+            .free_arithmetic
+            .get(&(const_0, const_1))
+            .copied()
+            .unwrap_or_else(|| {
+                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 < ArithmeticExtensionGate::<D>::num_ops(&self.config) - 1 {
+            self.batched_gates
+                .free_arithmetic
+                .insert((const_0, const_1), (gate, i + 1));
+        } else {
+            self.batched_gates.free_arithmetic.remove(&(const_0, const_1));
+        }
+
+        (gate, i)
+    }
+
+    pub fn find_switch_gate(&mut self, chunk_size: usize) -> (SwitchGate<F, D>, usize, usize) {
+        if self.batched_gates.current_switch_gates.len() < chunk_size {
+            self.batched_gates.current_switch_gates
+                .extend(vec![None; chunk_size - self.current_switch_gates.len()]);
+        }
+
+        let (gate, gate_index, mut next_copy) =
+            match self.current_switch_gates[chunk_size - 1].clone() {
+                None => {
+                    let gate = SwitchGate::<F, D>::new_from_config(self.config.clone(), chunk_size);
+                    let gate_index = self.add_gate(gate.clone(), vec![]);
+                    (gate, gate_index, 0)
+                }
+                Some((gate, idx, next_copy)) => (gate, idx, next_copy),
+            };
+        
+        let num_copies = gate.num_copies;
+        
+        if next_copy == num_copies {
+            self.batched_gates.current_switch_gates[chunk_size - 1] = None;
+        } else {
+            self.batched_gates.current_switch_gates[chunk_size - 1] = Some((gate, gate_index, next_copy + 1));
+        }
+
+        (gate, gate_index, next_copy)
+    }
+
+    pub fn find_u32_arithmetic_gate(&mut self) -> (usize, usize) {
+        let (gate_index, copy) = match self.batched_gates.current_u32_arithmetic_gate {
+            None => {
+                let gate = U32ArithmeticGate::new();
+                let gate_index = self.add_gate(gate, vec![]);
+                (gate_index, 0)
+            }
+            Some((gate_index, copy)) => (gate_index, copy),
+        };
+
+        if copy == NUM_U32_ARITHMETIC_OPS - 1 {
+            self.current_u32_arithmetic_gate = None;
+        } else {
+            self.current_u32_arithmetic_gate = Some((gate_index, copy + 1));
+        }
+
+        (gate_index, copy)
+    }
+
+    pub fn find_u32_subtraction_gate(&mut self) -> (usize, usize) {
+        let (gate_index, copy) = match self.batched_gates.current_u32_subtraction_gate {
+            None => {
+                let gate = U32SubtractionGate::new();
+                let gate_index = self.add_gate(gate, vec![]);
+                (gate_index, 0)
+            }
+            Some((gate_index, copy)) => (gate_index, copy),
+        };
+
+        if copy == NUM_U32_SUBTRACTION_OPS - 1 {
+            self.current_u32_subtraction_gate = None;
+        } else {
+            self.current_u32_subtraction_gate = Some((gate_index, copy + 1));
+        }
+
+        (gate_index, copy)
+    }
+
+    /// 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)
+    }
+}