From 318185d1c010fb721e9da9114c4eaf877799a75e Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Thu, 21 Oct 2021 12:56:02 +0200
Subject: [PATCH] Use only one `BaseSumGate`

---
 src/gadgets/range_check.rs |  4 +---
 src/gadgets/split_base.rs  |  6 +++++-
 src/gadgets/split_join.rs  | 17 +++++++++--------
 src/gates/base_sum.rs      | 14 +++++++++++---
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/src/gadgets/range_check.rs b/src/gadgets/range_check.rs
index f3b58d4a..a4153daf 100644
--- a/src/gadgets/range_check.rs
+++ b/src/gadgets/range_check.rs
@@ -9,9 +9,7 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     /// Checks that `x < 2^n_log` using a `BaseSumGate`.
     pub fn range_check(&mut self, x: Target, n_log: usize) {
-        let gate = self.add_gate(BaseSumGate::<2>::new(n_log), vec![]);
-        let sum = Target::wire(gate, BaseSumGate::<2>::WIRE_SUM);
-        self.connect(x, sum);
+        self.split_le(x, n_log);
     }
 
     /// Returns the first `num_low_bits` little-endian bits of `x`.
diff --git a/src/gadgets/split_base.rs b/src/gadgets/split_base.rs
index 1023e277..43de4375 100644
--- a/src/gadgets/split_base.rs
+++ b/src/gadgets/split_base.rs
@@ -48,13 +48,17 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             BaseSumGate::<2>::START_LIMBS + num_bits <= self.config.num_routed_wires,
             "Not enough routed wires."
         );
-        let gate_index = self.add_gate(BaseSumGate::<2>::new(num_bits), vec![]);
+        let gate_type = BaseSumGate::<2>::new_from_config::<F>(&self.config);
+        let gate_index = self.add_gate(gate_type, vec![]);
         for (limb, wire) in bits
             .clone()
             .zip(BaseSumGate::<2>::START_LIMBS..BaseSumGate::<2>::START_LIMBS + num_bits)
         {
             self.connect(limb.borrow().target, Target::wire(gate_index, wire));
         }
+        for l in gate_type.limbs().skip(num_bits) {
+            self.assert_zero(Target::wire(gate_index, l));
+        }
 
         self.add_simple_generator(BaseSumGenerator::<2> {
             gate_index,
diff --git a/src/gadgets/split_join.rs b/src/gadgets/split_join.rs
index 8875d041..cd920d8a 100644
--- a/src/gadgets/split_join.rs
+++ b/src/gadgets/split_join.rs
@@ -16,35 +16,36 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         if num_bits == 0 {
             return Vec::new();
         }
-        let bits_per_gate = (bits_u64(F::ORDER) - 1)
-            .min(self.config.num_routed_wires - BaseSumGate::<2>::START_LIMBS);
-        let k = ceil_div_usize(num_bits, bits_per_gate);
+        let gate_type = BaseSumGate::<2>::new_from_config::<F>(&self.config);
+        let k = ceil_div_usize(num_bits, gate_type.num_limbs);
         let gates = (0..k)
-            .map(|_| self.add_gate(BaseSumGate::<2>::new(bits_per_gate), vec![]))
+            .map(|_| self.add_gate(gate_type, vec![]))
             .collect::<Vec<_>>();
 
         let mut bits = Vec::with_capacity(num_bits);
         for &gate in &gates {
             let start_limbs = BaseSumGate::<2>::START_LIMBS;
-            for limb_input in start_limbs..start_limbs + bits_per_gate {
+            for limb_input in start_limbs..start_limbs + gate_type.num_limbs {
                 // `new_unsafe` is safe here because BaseSumGate::<2> forces it to be in `{0, 1}`.
                 bits.push(BoolTarget::new_unsafe(Target::wire(gate, limb_input)));
             }
         }
-        bits.drain(num_bits..);
+        for b in bits.drain(num_bits..) {
+            self.assert_zero(b.target);
+        }
 
         let zero = self.zero();
         let mut acc = zero;
         for &gate in gates.iter().rev() {
             let sum = Target::wire(gate, BaseSumGate::<2>::WIRE_SUM);
-            acc = self.mul_const_add(F::from_canonical_usize(1 << bits_per_gate), acc, sum);
+            acc = self.mul_const_add(F::from_canonical_usize(1 << gate_type.num_limbs), acc, sum);
         }
         self.connect(acc, integer);
 
         self.add_simple_generator(WireSplitGenerator {
             integer,
             gates,
-            num_limbs: bits_per_gate,
+            num_limbs: gate_type.num_limbs,
         });
 
         bits
diff --git a/src/gates/base_sum.rs b/src/gates/base_sum.rs
index 896ac37d..298d7090 100644
--- a/src/gates/base_sum.rs
+++ b/src/gates/base_sum.rs
@@ -2,19 +2,21 @@ use std::ops::Range;
 
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
-use crate::field::field_types::{Field, RichField};
+use crate::field::field_types::{Field, PrimeField, RichField};
 use crate::gates::gate::Gate;
 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::plonk_common::{reduce_with_powers, reduce_with_powers_ext_recursive};
 use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
+use crate::util::bits_u64;
 
 /// A gate which can decompose a number into base B little-endian limbs.
-#[derive(Clone, Debug)]
+#[derive(Copy, Clone, Debug)]
 pub struct BaseSumGate<const B: usize> {
-    num_limbs: usize,
+    pub num_limbs: usize,
 }
 
 impl<const B: usize> BaseSumGate<B> {
@@ -22,6 +24,12 @@ impl<const B: usize> BaseSumGate<B> {
         Self { num_limbs }
     }
 
+    pub fn new_from_config<F: PrimeField>(config: &CircuitConfig) -> Self {
+        let num_limbs = ((F::ORDER as f64).log(B as f64).floor() as usize)
+            .min(config.num_routed_wires - Self::START_LIMBS);
+        Self::new(num_limbs)
+    }
+
     pub const WIRE_SUM: usize = 0;
     pub const START_LIMBS: usize = 1;