Merge brance 'v2' into v2-bench

kzg_prover/src/chips/range/range_check.rs

@@ -1,113 +1,112 @@
+use crate::chips::range::utils::decompose_fp_to_byte_pairs;
 use halo2_proofs::arithmetic::Field;
 use halo2_proofs::circuit::{AssignedCell, Region, Value};
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
 use halo2_proofs::plonk::{Advice, Column, ConstraintSystem, Error, Expression, Fixed};
 use halo2_proofs::poly::Rotation;
-
 use std::fmt::Debug;
 
-use crate::chips::range::utils::decompose_fp_to_bytes;
-
-/// Configuration for the Range Check Chip
+/// Configuration for the Range Check u64 Chip
+/// Used to verify that an element lies in the u64 range.
 ///
-/// # Type Parameters
+/// # Fields
 ///
-/// * `N_BYTES`: Number of bytes in which the element to be checked should lie
+/// * `zs`: Four advice columns - contain the truncated right-shifted values of the element to be checked
 ///
-/// # Fields
+/// # Assumptions
 ///
-/// * `zs`: Advice columns - contain the truncated right-shifted values of the element to be checked
+/// * The lookup table `range_u16` is by default loaded with values from 0 to 2^16 - 1.
 ///
 /// Patterned after [halo2_gadgets](https://github.com/privacy-scaling-explorations/halo2/blob/main/halo2_gadgets/src/utilities/decompose_running_sum.rs)
 #[derive(Debug, Copy, Clone)]
-pub struct RangeCheckConfig<const N_BYTES: usize> {
-    zs: [Column<Advice>; N_BYTES],
+pub struct RangeCheckU64Config {
+    zs: [Column<Advice>; 4],
 }
 
-/// Helper chip that verfiies that the element witnessed in a given cell lies within a given range defined by N_BYTES.
-/// For example, Let's say we want to constraint 0x1f2f3f4f to be within the range N_BYTES=4.
+/// Helper chip that verfiies that the element witnessed in a given cell lies within the u64 range.
+/// For example, Let's say we want to constraint 0x1f2f3f4f5f6f7f8f to be a u64.
+/// Note that the lookup table `range` is by default loaded with values from 0 to 2^16 - 1.
 /// `z` is the advice column that contains the element to be checked.
 ///
-/// `z = 0x1f2f3f4f`
-/// `zs[0] = (0x1f2f3f4f - 0x4f) / 2^8 = 0x1f2f3f`
-/// `zs[1] = (0x1f2f3f - 0x3f) / 2^8 = 0x1f2f`
-/// `zs[2] = (0x1f2f - 0x2f) / 2^8 = 0x1f`
-/// `zs[3] = (0x1f - 0x1f) / 2^8 = 0x00`
+/// `z = 0x1f2f3f4f5f6f7f8f`
+/// `zs[0] = (0x1f2f3f4f5f6f7f8f - 0x7f8f) / 2^16 = 0x1f2f3f4f5f6f`
+/// `zs[1] = (0x1f2f3f4f5f6f - 0xf5f6f) / 2^16 = 0x1f2f3f4f`
+/// `zs[2] = (0x1f2f3f4f - 0x3f4f) / 2^16 = 0x1f2f`
+/// `zs[3] = (0x1f2f - 0x1f2f) / 2^16 = 0x00`
 ///
-///  z          | zs[0]      | zs[1]      | zs[2]      | zs[3]      |
-///  ---------  | ---------- | ---------- | ---------- | ---------- |
-///  0x1f2f3f4f | 0x1f2f3f   | 0x1f2f     | 0x1f       | 0x00       |
+///  z                  | zs[0]            | zs[1]         | zs[2]        | zs[3]      |
+///  ---------          | ----------       | ----------    | ----------   | ---------- |
+///  0x1f2f3f4f5f6f7f8f | 0x1f2f3f4f5f6f   | 0x1f2f3f4f     | 0x1f2f       | 0x00       |
 ///
-/// Column zs[0], at offset 0, contains the truncated right-shifted value z - ks[0] / 2^8 (shift right by 8 bits) where ks[0] is the 0-th decomposition big-endian of the element to be checked
-/// Column zs[1], at offset 0, contains the truncated right-shifted value zs[0] - ks[1] / 2^8 (shift right by 8 bits) where ks[1] is the 1-th decomposition big-endian of the element to be checked
-/// Column zs[2], at offset 0, contains the truncated right-shifted value zs[1] - ks[2] / 2^8 (shift right by 8 bits) where ks[2] is the 2-th decomposition big-endian of the element to be checked
-/// Column zs[3], at offset 0, contains the truncated right-shifted value zs[2] - ks[3] / 2^8 (shift right by 8 bits) where ks[3] is the 3-th decomposition big-endian of the element to be checked
+/// Column zs[0], at offset 0, contains the truncated right-shifted value z - ks[0] / 2^16 (shift right by 16 bits) where ks[0] is the 0-th decomposition big-endian of the element to be checked
+/// Column zs[1], at offset 0, contains the truncated right-shifted value zs[0] - ks[1] / 2^16 (shift right by 16 bits) where ks[1] is the 1-th decomposition big-endian of the element to be checked
+/// Column zs[2], at offset 0, contains the truncated right-shifted value zs[1] - ks[2] / 2^16 (shift right by 16 bits) where ks[2] is the 2-th decomposition big-endian of the element to be checked
+/// Column zs[3], at offset 0, contains the truncated right-shifted value zs[2] - ks[3] / 2^16 (shift right by 16 bits) where ks[3] is the 3-th decomposition big-endian of the element to be checked
 ///
 /// The contraints that are enforced are:
 /// 1.
-/// z - 2^8⋅zs[0] = ks[0] ∈ lookup_u8
+/// z - 2^16⋅zs[0] = ks[0] ∈ range_u16
 ///
 /// 2.
-/// for i = 0..=N_BYTES - 2:
-///     zs[i] - 2^8⋅zs[i+1] = ks[i]  ∈ lookup_u8
+/// for i = 0..=2:
+///     zs[i] - 2^16⋅zs[i+1] = ks[i]  ∈ range_u16
 ///
 /// 3.
-/// zs[N_BYTES - 1] == 0
+/// zs[3] == 0
 #[derive(Debug, Clone)]
-pub struct RangeCheckChip<const N_BYTES: usize> {
-    config: RangeCheckConfig<N_BYTES>,
+pub struct RangeCheckU64Chip {
+    config: RangeCheckU64Config,
 }
 
-impl<const N_BYTES: usize> RangeCheckChip<N_BYTES> {
-    pub fn construct(config: RangeCheckConfig<N_BYTES>) -> Self {
+impl RangeCheckU64Chip {
+    pub fn construct(config: RangeCheckU64Config) -> Self {
         Self { config }
     }
 
     /// Configures the Range Chip
+    /// Note: the lookup table should be loaded with values from `0` to `2^16 - 1` otherwise the range check will fail.
     pub fn configure(
         meta: &mut ConstraintSystem<Fp>,
         z: Column<Advice>,
-        zs: [Column<Advice>; N_BYTES],
-        range: Column<Fixed>,
-    ) -> RangeCheckConfig<N_BYTES> {
-        meta.annotate_lookup_any_column(range, || "LOOKUP_MAXBITS_RANGE");
-
+        zs: [Column<Advice>; 4],
+        range_u16: Column<Fixed>,
+    ) -> RangeCheckU64Config {
         // Constraint that the difference between the element to be checked and the 0-th truncated right-shifted value of the element to be within the range.
-        // z - 2^8⋅zs[0] = ks[0] ∈ lookup_u8
+        // z - 2^16⋅zs[0] = ks[0] ∈ range_u16
         meta.lookup_any(
-            "range u8 check for difference between the element to be checked and the 0-th truncated right-shifted value of the element",
+            "range check in u16 for difference between the element to be checked and the 0-th truncated right-shifted value of the element",
             |meta| {
                 let element = meta.query_advice(z, Rotation::cur());
 
                 let zero_truncation = meta.query_advice(zs[0], Rotation::cur());
 
-                let u8_range = meta.query_fixed(range, Rotation::cur());
+                let range_u16 = meta.query_fixed(range_u16, Rotation::cur());
 
-                let diff = element - zero_truncation * Expression::Constant(Fp::from(1 << 8));
+                let diff = element - zero_truncation * Expression::Constant(Fp::from(1 << 16));
 
-                vec![(diff, u8_range)]
+                vec![(diff, range_u16)]
             },
         );
 
-        // For i = 0..=N_BYTES - 2: Constraint that the difference between the i-th truncated right-shifted value and the (i+1)-th truncated right-shifted value to be within the range.
-        // zs[i] - 2^8⋅zs[i+1] = ks[i]  ∈ lookup_u8
-        for i in 0..=N_BYTES - 2 {
+        // For i = 0..=2: Constraint that the difference between the i-th truncated right-shifted value and the (i+1)-th truncated right-shifted value to be within the range.
+        // zs[i] - 2^16⋅zs[i+1] = ks[i]  ∈ range_u16
+        for i in 0..=2 {
             meta.lookup_any(
-                format!("range u8 check for difference between the {}-th truncated right-shifted value and the {}-th truncated right-shifted value", i, i+1).as_str(),
+                format!("range check in u16 for difference between the {}-th truncated right-shifted value and the {}-th truncated right-shifted value", i, i+1).as_str(),
                 |meta| {
                     let i_truncation = meta.query_advice(zs[i], Rotation::cur());
                     let i_plus_one_truncation = meta.query_advice(zs[i + 1], Rotation::cur());
 
-                    let u8_range = meta.query_fixed(range, Rotation::cur());
+                    let range_u16 = meta.query_fixed(range_u16, Rotation::cur());
 
-                    let diff = i_truncation - i_plus_one_truncation * Expression::Constant(Fp::from(1 << 8));
+                    let diff = i_truncation - i_plus_one_truncation * Expression::Constant(Fp::from(1 << 16));
 
-                    vec![(diff, u8_range)]
+                    vec![(diff, range_u16)]
                 },
             );
         }
 
-        RangeCheckConfig { zs }
+        RangeCheckU64Config { zs }
     }
 
     /// Assign the truncated right-shifted values of the element to be checked to the corresponding columns zs at offset 0 starting from the element to be checked.
@@ -116,25 +115,25 @@ impl<const N_BYTES: usize> RangeCheckChip<N_BYTES> {
         region: &mut Region<'_, Fp>,
         element: &AssignedCell<Fp, Fp>,
     ) -> Result<(), Error> {
-        // Decompose the element in #N_BYTES bytes
+        // Decompose the element in 4 byte pairs.
         let ks = element
             .value()
             .copied()
-            .map(|x| decompose_fp_to_bytes(x, N_BYTES))
-            .transpose_vec(N_BYTES);
+            .map(|x| decompose_fp_to_byte_pairs(x, 4))
+            .transpose_vec(4);
 
         // Initalize an empty vector of cells for the truncated right-shifted values of the element to be checked.
-        let mut zs = Vec::with_capacity(N_BYTES);
+        let mut zs = Vec::with_capacity(4);
         let mut z = element.clone();
 
-        // Calculate 1 / 2^8
-        let two_pow_eight_inv = Value::known(Fp::from(1 << 8).invert().unwrap());
+        // Calculate 1 / 2^16
+        let two_pow_sixteen_inv = Value::known(Fp::from(1 << 16).invert().unwrap());
 
         // Perform the assignment of the truncated right-shifted values to zs columns.
         for (i, k) in ks.iter().enumerate() {
             let zs_next = {
                 let k = k.map(|byte| Fp::from(byte as u64));
-                let zs_next_val = (z.value().copied() - k) * two_pow_eight_inv;
+                let zs_next_val = (z.value().copied() - k) * two_pow_sixteen_inv;
                 region.assign_advice(
                     || format!("zs_{:?}", i),
                     self.config.zs[i],
@@ -148,7 +147,7 @@ impl<const N_BYTES: usize> RangeCheckChip<N_BYTES> {
         }
 
         // Constrain the final running sum output to be zero.
-        region.constrain_constant(zs[N_BYTES - 1].cell(), Fp::from(0))?;
+        region.constrain_constant(zs[3].cell(), Fp::from(0))?;
 
         Ok(())
     }