test-keccak-normalize-table: unit tests for the content of lookup tab…

…les (#1263)

### Description

This PR verifies the properties of the lookup tables used in Keccak.

A port of scroll-tech#326

### Issue Link

This is a good regression test before
#1241

### Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update

### Contents

- A test circuit that populates the fixed lookup tables using the
functions under test.
- Check that the table contain valid values, and nothing else.
- Minor refactoring for testability. Make pure functions that calculate
parameters, and isolate the dependency on environment variables.

### How Has This Been Tested?

    cargo test -p zkevm-circuits --features test --release -- keccak

---------

Co-authored-by: Aurélien Nicolas <info@nau.re>

zkevm-circuits/src/keccak_circuit/table.rs

@@ -7,17 +7,31 @@ use halo2_proofs::{
 };
 use itertools::Itertools;
 
-/// Loads a normalization table with the given parameters
+/// Loads a normalization table with the given parameters and KECCAK_DEGREE.
 pub(crate) fn load_normalize_table<F: Field>(
     layouter: &mut impl Layouter<F>,
     name: &str,
     tables: &[TableColumn; 2],
     range: u64,
 ) -> Result<(), Error> {
-    let part_size = get_num_bits_per_lookup(range as usize);
+    let log_height = get_degree();
+    load_normalize_table_impl(layouter, name, tables, range, log_height)
+}
+
+// Implementation of the above without environment dependency.
+fn load_normalize_table_impl<F: Field>(
+    layouter: &mut impl Layouter<F>,
+    name: &str,
+    tables: &[TableColumn; 2],
+    range: u64,
+    log_height: usize,
+) -> Result<(), Error> {
+    assert!(range <= BIT_SIZE as u64);
+    let part_size = get_num_bits_per_lookup_impl(range as usize, log_height);
     layouter.assign_table(
         || format!("{} table", name),
         |mut table| {
+            // Iterate over all combinations of parts, each taking values in the range.
             for (offset, perm) in (0..part_size)
                 .map(|_| 0u64..range)
                 .multi_cartesian_product()
@@ -113,3 +127,168 @@ pub(crate) fn load_lookup_table<F: Field>(
         },
     )
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use halo2_proofs::circuit::SimpleFloorPlanner;
+    use halo2_proofs::dev::{CellValue, MockProver};
+    use halo2_proofs::halo2curves::bn256::Fr as F;
+    use halo2_proofs::plonk::{Circuit, ConstraintSystem};
+    use itertools::Itertools;
+    use std::iter::zip;
+
+    #[test]
+    fn normalize_table() {
+        normalize_table_impl(3, 10);
+        normalize_table_impl(4, 10);
+        normalize_table_impl(6, 10);
+        normalize_table_impl(6, 19);
+    }
+
+    fn normalize_table_impl(range: usize, log_height: usize) {
+        let table = build_table(&TableTestCircuit {
+            range,
+            log_height,
+            normalize_else_chi: true,
+        });
+
+        // On all rows, all inputs/outputs are correct, i.e. they have the same low bit.
+        assert_eq!(BIT_COUNT, 3, "this test assumes BIT_COUNT=3");
+        for (inp, out) in table.iter() {
+            for pos in (0..64).step_by(BIT_COUNT) {
+                assert_eq!((inp >> pos) & 1, (out >> pos) & 0b111);
+            }
+        }
+    }
+
+    #[test]
+    fn chi_table() {
+        // Check the base pattern for all combinations of bits.
+        for i in 0..16_usize {
+            let (a, b, c, d) = (i & 1, (i >> 1) & 1, (i >> 2) & 1, (i >> 3) & 1);
+            assert_eq!(
+                CHI_BASE_LOOKUP_TABLE[3 - 2 * a + b - c],
+                (a ^ ((!b) & c)) as u8
+            );
+            assert_eq!(
+                CHI_EXT_LOOKUP_TABLE[5 - 2 * a - b + c - 2 * d],
+                (a ^ ((!b) & c) ^ d) as u8
+            );
+        }
+
+        // Check the table with multiple parts per row.
+        chi_table_impl(10);
+        chi_table_impl(19);
+    }
+
+    fn chi_table_impl(log_height: usize) {
+        let range = 5; // CHI_BASE_LOOKUP_RANGE
+        let table = build_table(&TableTestCircuit {
+            range,
+            log_height,
+            normalize_else_chi: false,
+        });
+
+        // On all rows, all input/output pairs match the base table.
+        for (inp, out) in table.iter() {
+            for pos in (0..64).step_by(BIT_COUNT) {
+                let inp = ((inp >> pos) & 7) as usize;
+                let out = ((out >> pos) & 7) as u8;
+                assert_eq!(out, CHI_BASE_LOOKUP_TABLE[inp]);
+            }
+        }
+    }
+
+    // ---- Helpers ----
+
+    fn build_table(circuit: &TableTestCircuit) -> Vec<(u64, u64)> {
+        let prover = MockProver::<F>::run(circuit.log_height as u32, circuit, vec![]).unwrap();
+
+        let columns = prover.fixed();
+        assert_eq!(columns.len(), 2);
+        let unused_rows = 6; // What MockProver uses on this test circuit.
+        let used_rows = (1 << circuit.log_height) - unused_rows;
+
+        // Check the unused rows.
+        for io in zip(&columns[0], &columns[1]).skip(used_rows) {
+            assert_eq!(io, (&CellValue::Unassigned, &CellValue::Unassigned));
+        }
+
+        // Get the generated lookup table with the form: table[row] = (input, output).
+        let table = zip(&columns[0], &columns[1])
+            .take(used_rows)
+            .map(|(inp, out)| (unwrap_u64(inp), unwrap_u64(out)))
+            .collect::<Vec<_>>();
+
+        // All possible combinations of inputs are there.
+        let unique_rows = table.iter().unique().count();
+        assert_eq!(unique_rows, circuit.expected_num_entries());
+
+        table
+    }
+
+    #[derive(Clone)]
+    struct TableTestCircuit {
+        range: usize,
+        log_height: usize,
+        normalize_else_chi: bool,
+    }
+
+    impl TableTestCircuit {
+        fn expected_num_entries(&self) -> usize {
+            let num_bits = get_num_bits_per_lookup_impl(self.range, self.log_height);
+            self.range.pow(num_bits as u32)
+        }
+    }
+
+    impl Circuit<F> for TableTestCircuit {
+        type Config = [TableColumn; 2];
+        type FloorPlanner = SimpleFloorPlanner;
+
+        fn without_witnesses(&self) -> Self {
+            self.clone()
+        }
+
+        fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
+            array_init::array_init(|_| meta.lookup_table_column())
+        }
+
+        fn synthesize(
+            &self,
+            config: Self::Config,
+            mut layouter: impl Layouter<F>,
+        ) -> Result<(), Error> {
+            if self.normalize_else_chi {
+                load_normalize_table_impl(
+                    &mut layouter,
+                    "normalize",
+                    &config,
+                    self.range as u64,
+                    self.log_height,
+                )?;
+            } else {
+                let num_bits = get_num_bits_per_lookup_impl(self.range, self.log_height);
+                load_lookup_table(
+                    &mut layouter,
+                    "chi base",
+                    &config,
+                    num_bits,
+                    &CHI_BASE_LOOKUP_TABLE,
+                )?;
+            }
+            Ok(())
+        }
+    }
+
+    fn unwrap_u64<F: Field>(cv: &CellValue<F>) -> u64 {
+        match *cv {
+            CellValue::Assigned(f) => {
+                let f = f.get_lower_128();
+                assert_eq!(f >> 64, 0);
+                f as u64
+            }
+            _ => panic!("the cell should be assigned"),
+        }
+    }
+}

zkevm-circuits/src/keccak_circuit/util.rs

@@ -210,12 +210,18 @@ pub(crate) fn get_degree() -> usize {
 }
 
 /// Returns how many bits we can process in a single lookup given the range of
-/// values the bit can have and the height of the circuit.
+/// values the bit can have and the height of the circuit (via KECCAK_DEGREE).
 pub(crate) fn get_num_bits_per_lookup(range: usize) -> usize {
+    let log_height = get_degree();
+    get_num_bits_per_lookup_impl(range, log_height)
+}
+
+// Implementation of the above without environment dependency.
+pub(crate) fn get_num_bits_per_lookup_impl(range: usize, log_height: usize) -> usize {
     let num_unusable_rows = 31;
-    let degree = get_degree() as u32;
+    let height = 2usize.pow(log_height as u32);
     let mut num_bits = 1;
-    while range.pow(num_bits + 1) + num_unusable_rows <= 2usize.pow(degree) {
+    while range.pow(num_bits + 1) + num_unusable_rows <= height {
         num_bits += 1;
     }
     num_bits as usize
@@ -291,3 +297,31 @@ pub(crate) mod to_bytes {
         bytes
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use halo2_proofs::halo2curves::bn256::Fr as F;
+
+    #[test]
+    fn pack_into_bits() {
+        // The example number 128 in binary: |1|0|0|0|0|0|0|0|
+        // In packed form:                 |001|000|000|000|000|000|000|000|
+        let msb = 1 << (7 * BIT_COUNT);
+        for (idx, expected) in [(0, 0), (1, 1), (128, msb), (129, msb | 1)] {
+            let packed: F = pack(&into_bits(&[idx as u8]));
+            assert_eq!(packed, F::from(expected));
+        }
+    }
+
+    #[test]
+    fn num_bits_per_lookup() {
+        // Typical values.
+        assert_eq!(get_num_bits_per_lookup_impl(3, 19), 11);
+        assert_eq!(get_num_bits_per_lookup_impl(4, 19), 9);
+        assert_eq!(get_num_bits_per_lookup_impl(5, 19), 8);
+        assert_eq!(get_num_bits_per_lookup_impl(6, 19), 7);
+        // The largest possible value does not overflow u64.
+        assert_eq!(get_num_bits_per_lookup_impl(3, 32) * BIT_COUNT, 60);
+    }
+}