diff --git a/src/gadgets/arithmetic_extension.rs b/src/gadgets/arithmetic_extension.rs index 24499760..e2654dcc 100644 --- a/src/gadgets/arithmetic_extension.rs +++ b/src/gadgets/arithmetic_extension.rs @@ -12,33 +12,6 @@ use crate::plonk::circuit_builder::CircuitBuilder; use crate::util::bits_u64; impl, const D: usize> CircuitBuilder { - /// 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::::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, diff --git a/src/gadgets/arithmetic_u32.rs b/src/gadgets/arithmetic_u32.rs index 2b83b03d..4f60dde7 100644 --- a/src/gadgets/arithmetic_u32.rs +++ b/src/gadgets/arithmetic_u32.rs @@ -36,14 +36,7 @@ impl, const D: usize> CircuitBuilder { 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, const D: usize> CircuitBuilder { U32ArithmeticGate::::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, const D: usize> CircuitBuilder { 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::::wire_ith_multiplicand_0(copy), + ), + x.0, + ); + self.connect( + Target::wire( + gate_index, + U32ArithmeticGate::::wire_ith_multiplicand_1(copy), + ), + y.0, + ); + self.connect( + Target::wire(gate_index, U32ArithmeticGate::::wire_ith_addend(copy)), + z.0, + ); + + let output_low = U32Target(Target::wire( + gate_index, + U32ArithmeticGate::::wire_ith_output_low_half(copy), + )); + let output_high = U32Target(Target::wire( + gate_index, + U32ArithmeticGate::::wire_ith_output_high_half(copy), + )); + + (output_low, output_high) + } } diff --git a/src/gadgets/permutation.rs b/src/gadgets/permutation.rs index fd4a897f..a0c9b087 100644 --- a/src/gadgets/permutation.rs +++ b/src/gadgets/permutation.rs @@ -73,20 +73,8 @@ impl, const D: usize> CircuitBuilder { 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::::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, const D: usize> CircuitBuilder { 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) } diff --git a/src/plonk/circuit_builder.rs b/src/plonk/circuit_builder.rs index 019bc71e..f915493a 100644 --- a/src/plonk/circuit_builder.rs +++ b/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, const D: usize> { /// Memoized results of `arithmetic_extension` calls. pub(crate) arithmetic_results: HashMap, ExtensionTarget>, - /// 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, - - // `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, 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 } impl, const D: usize> CircuitBuilder { @@ -110,11 +95,7 @@ impl, const D: usize> CircuitBuilder { 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, const D: usize> CircuitBuilder { 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, const D: usize> CircuitBuilder { 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 { constants.iter().map(|&c| self.constant(c)).collect() } @@ -846,3 +807,152 @@ impl, const D: usize> CircuitBuilder { } } } + +/// +pub struct BatchedGates, 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, 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, const D: usize> BatchedGates { + 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, const D: usize> CircuitBuilder { + /// 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::::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, 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::::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) + } +}