feat(ecc): FpChip::range_check now works with max_bits < n * (k-1)

axiom-crypto · Jul 7, 2023 · 54f8bd2 · 54f8bd2
1 parent 97ab27d
commit 54f8bd2
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 26 deletions.
diff --git a/halo2-ecc/src/fields/fp.rs b/halo2-ecc/src/fields/fp.rs
@@ -15,6 +15,7 @@ use halo2_base::{
 };
 use num_bigint::{BigInt, BigUint};
 use num_traits::One;
+use std::cmp;
 use std::{cmp::max, marker::PhantomData};
 
 pub type BaseFieldChip<'range, C> =
@@ -298,7 +299,8 @@ impl<'range, F: PrimeField, Fp: PrimeField> FieldChip<F> for FpChip<'range, F, F
     }
 
     /// # Assumptions
-    /// * `max_bits` in `(n * (k - 1), n * k]`
+    /// * `max_bits <= n * k` where `n = self.limb_bits` and `k = self.num_limbs`
+    /// * `a.truncation.limbs.len() = self.num_limbs`
     fn range_check(
         &self,
         ctx: &mut Context<F>,
@@ -307,15 +309,14 @@ impl<'range, F: PrimeField, Fp: PrimeField> FieldChip<F> for FpChip<'range, F, F
     ) {
         let n = self.limb_bits;
         let a = a.into();
-        let k = a.truncation.limbs.len();
-        debug_assert!(max_bits > n * (k - 1) && max_bits <= n * k);
-        let last_limb_bits = max_bits - n * (k - 1);
+        let mut remaining_bits = max_bits;
 
         debug_assert!(a.value.bits() as usize <= max_bits);
 
         // range check limbs of `a` are in [0, 2^n) except last limb should be in [0, 2^last_limb_bits)
-        for (i, cell) in a.truncation.limbs.into_iter().enumerate() {
-            let limb_bits = if i == k - 1 { last_limb_bits } else { n };
+        for cell in a.truncation.limbs {
+            let limb_bits = cmp::min(n, remaining_bits);
+            remaining_bits -= limb_bits;
             self.range.range_check(ctx, cell, limb_bits);
         }
     }

diff --git a/halo2-ecc/src/fields/tests/fp/mod.rs b/halo2-ecc/src/fields/tests/fp/mod.rs
@@ -1,9 +1,12 @@
+use std::env::set_var;
+
 use crate::fields::fp::FpChip;
 use crate::fields::{FieldChip, PrimeField};
 use crate::halo2_proofs::{
     dev::MockProver,
     halo2curves::bn256::{Fq, Fr},
 };
+
 use halo2_base::gates::builder::{GateThreadBuilder, RangeCircuitBuilder};
 use halo2_base::gates::RangeChip;
 use halo2_base::utils::biguint_to_fe;
@@ -15,39 +18,56 @@ pub mod assert_eq;
 
 const K: usize = 10;
 
-fn fp_mul_test<F: PrimeField>(
-    ctx: &mut Context<F>,
+fn fp_chip_test(
+    k: usize,
     lookup_bits: usize,
     limb_bits: usize,
     num_limbs: usize,
-    _a: Fq,
-    _b: Fq,
+    f: impl Fn(&mut Context<Fr>, &FpChip<Fr, Fq>),
 ) {
-    std::env::set_var("LOOKUP_BITS", lookup_bits.to_string());
-    let range = RangeChip::<F>::default(lookup_bits);
-    let chip = FpChip::<F, Fq>::new(&range, limb_bits, num_limbs);
+    set_var("LOOKUP_BITS", lookup_bits.to_string());
+    let range = RangeChip::<Fr>::default(lookup_bits);
+    let chip = FpChip::<Fr, Fq>::new(&range, limb_bits, num_limbs);
 
-    let [a, b] = [_a, _b].map(|x| chip.load_private(ctx, x));
-    let c = chip.mul(ctx, a, b);
+    let mut builder = GateThreadBuilder::mock();
+    f(builder.main(0), &chip);
 
-    assert_eq!(c.0.truncation.to_bigint(limb_bits), c.0.value);
-    assert_eq!(c.native().value(), &biguint_to_fe(&(c.value() % modulus::<F>())));
-    assert_eq!(c.0.value, fe_to_biguint(&(_a * _b)).into())
+    builder.config(k, Some(10));
+    let circuit = RangeCircuitBuilder::mock(builder);
+    MockProver::run(k as u32, &circuit, vec![]).unwrap().assert_satisfied();
 }
 
 #[test]
 fn test_fp() {
-    let k = K;
-    let a = Fq::random(OsRng);
-    let b = Fq::random(OsRng);
+    let limb_bits = 88;
+    let num_limbs = 3;
+    fp_chip_test(K, K - 1, limb_bits, num_limbs, |ctx, chip| {
+        let _a = Fq::random(OsRng);
+        let _b = Fq::random(OsRng);
 
-    let mut builder = GateThreadBuilder::<Fr>::mock();
-    fp_mul_test(builder.main(0), k - 1, 88, 3, a, b);
+        let [a, b] = [_a, _b].map(|x| chip.load_private(ctx, x));
+        let c = chip.mul(ctx, a, b);
 
-    builder.config(k, Some(10));
-    let circuit = RangeCircuitBuilder::mock(builder);
+        assert_eq!(c.0.truncation.to_bigint(limb_bits), c.0.value);
+        assert_eq!(c.native().value(), &biguint_to_fe(&(c.value() % modulus::<Fr>())));
+        assert_eq!(c.0.value, fe_to_biguint(&(_a * _b)).into());
+    });
+}
 
-    MockProver::run(k as u32, &circuit, vec![]).unwrap().assert_satisfied();
+#[test]
+fn test_range_check() {
+    fp_chip_test(K, K - 1, 88, 3, |ctx, chip| {
+        let mut range_test = |x, bits| {
+            let x = chip.load_private(ctx, x);
+            chip.range_check(ctx, x, bits);
+        };
+        let a = -Fq::one();
+        range_test(a, Fq::NUM_BITS as usize);
+        range_test(Fq::one(), 1);
+        range_test(Fq::from(u64::MAX), 64);
+        range_test(Fq::zero(), 1);
+        range_test(Fq::zero(), 0);
+    });
 }
 
 #[cfg(feature = "dev-graph")]