perf: improve range gate implementation

CHANGELOG.md

@@ -38,6 +38,7 @@ and follow [semantic versioning](https://semver.org/) for our releases.
 - [#297](https://github.com/EspressoSystems/jellyfish/pull/297) Updated `tagged-base64` dependency to the `crates.io` package
 - [#299](https://github.com/EspressoSystems/jellyfish/pull/299) For Merkle tree, `DigestAlgorithm` now returns a `Result` type.
 - [#302](https://github.com/EspressoSystems/jellyfish/pull/302) Followup APIs for non-native ECC circuit support.
+- [#323](https://github.com/EspressoSystems/jellyfish/pull/323) Improve performance of range gate in ultra plonk.
 
 ### Removed
 

relation/src/gadgets/emulated.rs

@@ -150,6 +150,8 @@ impl<F: PrimeField> PlonkCircuit<F> {
 
     /// Constrain that a*b=c in the emulated field.
     /// Checking that a * b - k * E::MODULUS = c.
+    /// This function doesn't perform emulated variable validaty check on the
+    /// input a, b and c. We assume that they are already performed elsewhere.
     pub fn emulated_mul_gate<E: EmulationConfig<F>>(
         &mut self,
         a: &EmulatedVariable<E>,
@@ -292,6 +294,8 @@ impl<F: PrimeField> PlonkCircuit<F> {
     }
 
     /// Constrain that a*b=c in the emulated field for a constant b.
+    /// This function doesn't perform emulated variable validaty check on the
+    /// input a and c. We assume that they are already performed elsewhere.
     pub fn emulated_mul_constant_gate<E: EmulationConfig<F>>(
         &mut self,
         a: &EmulatedVariable<E>,
@@ -318,11 +322,6 @@ impl<F: PrimeField> PlonkCircuit<F> {
             E::NUM_LIMBS,
         );
 
-        // range checking for output c
-        c.0.iter()
-            .map(|v| self.enforce_in_range(*v, E::B))
-            .collect::<Result<Vec<_>, CircuitError>>()?;
-
         // enforcing a * b - k * E::MODULUS = c mod 2^t
 
         // first compare the first limb
@@ -413,6 +412,8 @@ impl<F: PrimeField> PlonkCircuit<F> {
 
     /// Constrain that a+b=c in the emulated field.
     /// Checking whether a + b = k * E::MODULUS + c
+    /// This function doesn't perform emulated variable validaty check on the
+    /// input a, b and c. We assume that they are already performed elsewhere.
     pub fn emulated_add_gate<E: EmulationConfig<F>>(
         &mut self,
         a: &EmulatedVariable<E>,
@@ -482,6 +483,8 @@ impl<F: PrimeField> PlonkCircuit<F> {
     }
 
     /// Constrain that a+b=c in the emulated field.
+    /// This function doesn't perform emulated variable validaty check on the
+    /// input a and c. We assume that they are already performed elsewhere.
     pub fn emulated_add_constant_gate<E: EmulationConfig<F>>(
         &mut self,
         a: &EmulatedVariable<E>,

relation/src/gadgets/range.rs

@@ -17,7 +17,7 @@ impl<F: PrimeField> PlonkCircuit<F> {
     /// Constrain a variable to be within the [0, 2^`bit_len`) range
     /// Return error if the variable is invalid.
     pub fn enforce_in_range(&mut self, a: Variable, bit_len: usize) -> Result<(), CircuitError> {
-        if self.support_lookup() && bit_len % self.range_bit_len()? == 0 {
+        if self.support_lookup() {
             self.range_gate_with_lookup(a, bit_len)?;
         } else {
             self.range_gate_internal(a, bit_len)?;
@@ -95,7 +95,7 @@ impl<F: PrimeField> PlonkCircuit<F> {
 /// Private helper function for range gate
 impl<F: PrimeField> PlonkCircuit<F> {
     // internal of a range check gate
-    fn range_gate_internal(
+    pub(crate) fn range_gate_internal(
         &mut self,
         a: Variable,
         bit_len: usize,

relation/src/gadgets/ultraplonk/range.rs

@@ -16,7 +16,7 @@ impl<F: PrimeField> PlonkCircuit<F> {
     /// Constrain a variable to be within the [0, 2^{bit_len}) range
     /// Return error if one of the following holds:
     /// 1. the variable is invalid;
-    /// 2. `RANGE_BIT_LEN` equals zero or does not divide `bit_len`;
+    /// 2. `RANGE_BIT_LEN` equals zero;
     /// 3. the circuit does not support lookup.
     pub(crate) fn range_gate_with_lookup(
         &mut self,
@@ -28,24 +28,25 @@ impl<F: PrimeField> PlonkCircuit<F> {
         if bit_len == 0 {
             return Err(ParameterError("bit_len cannot be zero".to_string()));
         }
-        if bit_len % range_bit_len != 0 {
-            return Err(ParameterError(
-                "circuit.range_bit_len does not divide bit_len".to_string(),
-            ));
-        }
         self.check_var_bound(a)?;
-        let len = bit_len / range_bit_len;
+        let leftover = bit_len % range_bit_len;
+        let lookup_len = bit_len / range_bit_len;
+        let len = lookup_len + if leftover > 0 { 1 } else { 0 };
         let reprs_le = decompose_le(self.witness(a)?, len, range_bit_len);
         let reprs_le_vars: Vec<Variable> = reprs_le
             .iter()
             .map(|&val| self.create_variable(val))
             .collect::<Result<Vec<_>, CircuitError>>()?;
 
         // add range gates for decomposed variables
-        for &var in reprs_le_vars.iter() {
+        for &var in reprs_le_vars[..lookup_len].iter() {
             self.add_range_check_variable(var)?;
         }
 
+        if leftover > 0 {
+            self.range_gate_internal(reprs_le_vars[lookup_len], leftover)?;
+        }
+
         // add linear combination gates
         self.decomposition_gate(reprs_le_vars, a, F::from(range_size as u64))?;
 
@@ -124,24 +125,30 @@ mod test {
 
     #[test]
     fn test_range_gate_with_lookup() -> Result<(), CircuitError> {
-        test_range_gate_with_lookup_helper::<FqEd254>()?;
-        test_range_gate_with_lookup_helper::<FqEd377>()?;
-        test_range_gate_with_lookup_helper::<FqEd381>()?;
-        test_range_gate_with_lookup_helper::<Fq377>()
+        (2 * RANGE_BIT_LEN_FOR_TEST..3 * RANGE_BIT_LEN_FOR_TEST)
+            .map(|bitlen| -> Result<(), CircuitError> {
+                test_range_gate_with_lookup_helper::<FqEd254>(bitlen)?;
+                test_range_gate_with_lookup_helper::<FqEd377>(bitlen)?;
+                test_range_gate_with_lookup_helper::<FqEd381>(bitlen)?;
+                test_range_gate_with_lookup_helper::<Fq377>(bitlen)
+            })
+            .collect::<Result<Vec<_>, CircuitError>>()?;
+        Ok(())
     }
-    fn test_range_gate_with_lookup_helper<F: PrimeField>() -> Result<(), CircuitError> {
+    fn test_range_gate_with_lookup_helper<F: PrimeField>(
+        bit_len: usize,
+    ) -> Result<(), CircuitError> {
         let mut circuit: PlonkCircuit<F> = PlonkCircuit::new_ultra_plonk(RANGE_BIT_LEN_FOR_TEST);
         let mut rng = test_rng();
-        let bit_len = RANGE_BIT_LEN_FOR_TEST * 4;
 
         // Good path
         let a = (0..10)
             .map(|_| circuit.create_variable(F::from(rng.gen_range(0..u32::MAX))))
             .collect::<Result<Vec<_>, CircuitError>>()?;
         for &var in a.iter() {
-            circuit.range_gate_with_lookup(var, bit_len)?;
+            circuit.range_gate_with_lookup(var, 32)?;
         }
-        circuit.range_gate_with_lookup(circuit.zero(), RANGE_BIT_LEN_FOR_TEST)?;
+        circuit.range_gate_with_lookup(circuit.zero(), bit_len)?;
         assert!(circuit.check_circuit_satisfiability(&[]).is_ok());
 
         // Error paths
@@ -153,23 +160,19 @@ mod test {
         *circuit.witness_mut(a[0]) = tmp;
 
         let mut circuit: PlonkCircuit<F> = PlonkCircuit::new_ultra_plonk(RANGE_BIT_LEN_FOR_TEST);
-        // Should fail when the value = 2^RANGE_BIT_LEN_FOR_TEST
-        let a_var = circuit.create_variable(F::from(1u32 << RANGE_BIT_LEN_FOR_TEST))?;
-        circuit.range_gate_with_lookup(a_var, RANGE_BIT_LEN_FOR_TEST)?;
+        // Should fail when the value = 2^bit_len
+        let a_var = circuit.create_variable(F::from(1u64 << bit_len))?;
+        circuit.range_gate_with_lookup(a_var, bit_len)?;
         assert!(circuit.check_circuit_satisfiability(&[]).is_err());
 
-        // Should fail when the value = 2^{2*RANGE_BIT_LEN_FOR_TEST}
-        let a_var = circuit.create_variable(F::from(1u32 << (2 * RANGE_BIT_LEN_FOR_TEST)))?;
-        circuit.range_gate_with_lookup(a_var, 2 * RANGE_BIT_LEN_FOR_TEST)?;
+        // Should fail when the value = 2^{2*bit_len}
+        let a_var = circuit.create_variable(F::from(1u64 << (2 * bit_len)))?;
+        circuit.range_gate_with_lookup(a_var, 2 * bit_len)?;
         assert!(circuit.check_circuit_satisfiability(&[]).is_err());
 
         let zero_var = circuit.zero();
         // bit_len = 0
         assert!(circuit.range_gate_with_lookup(zero_var, 0).is_err());
-        // bit_len % RANGE_BIT_LEN_FOR_TEST != 0
-        assert!(circuit
-            .range_gate_with_lookup(zero_var, bit_len + 1)
-            .is_err());
         // Check variable out of bound error.
         assert!(circuit
             .range_gate_with_lookup(circuit.num_vars(), bit_len)