Range-check keccak sponge inputs to bytes (#1342)

* Range-check keccak sponge inputs to bytes

* Move outside of inner loop

* Apply review

evm/src/keccak_sponge/columns.rs

@@ -1,5 +1,6 @@
 use std::borrow::{Borrow, BorrowMut};
 use std::mem::{size_of, transmute};
+use std::ops::Range;
 
 use crate::util::{indices_arr, transmute_no_compile_time_size_checks};
 
@@ -72,12 +73,36 @@ pub(crate) struct KeccakSpongeColumnsView<T: Copy> {
     /// The first part of the state of the sponge, seen as bytes, after the permutation is applied.
     /// This also represents the output digest of the Keccak sponge during the squeezing phase.
     pub updated_digest_state_bytes: [T; KECCAK_DIGEST_BYTES],
+
+    /// The counter column (used for the range check) starts from 0 and increments.
+    pub range_counter: T,
+    /// The frequencies column used in logUp.
+    pub rc_frequencies: T,
 }
 
 // `u8` is guaranteed to have a `size_of` of 1.
 /// Number of columns in `KeccakSpongeStark`.
 pub const NUM_KECCAK_SPONGE_COLUMNS: usize = size_of::<KeccakSpongeColumnsView<u8>>();
 
+// Indices for LogUp range-check.
+// They are on the last registers of this table.
+pub(crate) const RC_FREQUENCIES: usize = NUM_KECCAK_SPONGE_COLUMNS - 1;
+pub(crate) const RANGE_COUNTER: usize = RC_FREQUENCIES - 1;
+
+pub(crate) const BLOCK_BYTES_START: usize =
+    6 + KECCAK_RATE_BYTES + KECCAK_RATE_U32S + KECCAK_CAPACITY_U32S;
+/// Indices for the range-checked values, i.e. the `block_bytes` section.
+// TODO: Find a better way to access those indices
+pub(crate) const fn get_block_bytes_range() -> Range<usize> {
+    BLOCK_BYTES_START..BLOCK_BYTES_START + KECCAK_RATE_BYTES
+}
+
+/// Return the index for the targeted `block_bytes` element.
+pub(crate) const fn get_single_block_bytes_value(i: usize) -> usize {
+    debug_assert!(i < KECCAK_RATE_BYTES);
+    get_block_bytes_range().start + i
+}
+
 impl<T: Copy> From<[T; NUM_KECCAK_SPONGE_COLUMNS]> for KeccakSpongeColumnsView<T> {
     fn from(value: [T; NUM_KECCAK_SPONGE_COLUMNS]) -> Self {
         unsafe { transmute_no_compile_time_size_checks(value) }

evm/src/keccak_sponge/keccak_sponge_stark.rs

@@ -12,6 +12,7 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
 use plonky2::timed;
 use plonky2::util::timing::TimingTree;
+use plonky2::util::transpose;
 use plonky2_util::ceil_div_usize;
 
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
@@ -19,10 +20,13 @@ use crate::cpu::kernel::keccak_util::keccakf_u32s;
 use crate::cross_table_lookup::Column;
 use crate::evaluation_frame::{StarkEvaluationFrame, StarkFrame};
 use crate::keccak_sponge::columns::*;
+use crate::lookup::Lookup;
 use crate::stark::Stark;
-use crate::util::trace_rows_to_poly_values;
 use crate::witness::memory::MemoryAddress;
 
+/// Strict upper bound for the individual bytes range-check.
+const BYTE_RANGE_MAX: usize = 256;
+
 /// Creates the vector of `Columns` corresponding to:
 /// - the address in memory of the inputs,
 /// - the length of the inputs,
@@ -246,13 +250,12 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
             self.generate_trace_rows(operations, min_rows)
         );
 
-        let trace_polys = timed!(
-            timing,
-            "convert to PolynomialValues",
-            trace_rows_to_poly_values(trace_rows)
-        );
+        let trace_row_vecs: Vec<_> = trace_rows.into_iter().map(|row| row.to_vec()).collect();
 
-        trace_polys
+        let mut trace_cols = transpose(&trace_row_vecs);
+        self.generate_range_checks(&mut trace_cols);
+
+        trace_cols.into_iter().map(PolynomialValues::new).collect()
     }
 
     /// Generates the trace rows given the vector of `KeccakSponge` operations.
@@ -477,6 +480,38 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
         // indicating that it's a dummy/padding row.
         KeccakSpongeColumnsView::default().into()
     }
+
+    /// Expects input in *column*-major layout
+    fn generate_range_checks(&self, cols: &mut Vec<Vec<F>>) {
+        debug_assert!(cols.len() == NUM_KECCAK_SPONGE_COLUMNS);
+
+        let n_rows = cols[0].len();
+        debug_assert!(cols.iter().all(|col| col.len() == n_rows));
+
+        for i in 0..BYTE_RANGE_MAX {
+            cols[RANGE_COUNTER][i] = F::from_canonical_usize(i);
+        }
+        for i in BYTE_RANGE_MAX..n_rows {
+            cols[RANGE_COUNTER][i] = F::from_canonical_usize(BYTE_RANGE_MAX - 1);
+        }
+
+        // For each column c in cols, generate the range-check
+        // permutations and put them in the corresponding range-check
+        // columns rc_c and rc_c+1.
+        for col in 0..KECCAK_RATE_BYTES {
+            let c = get_single_block_bytes_value(col);
+            for i in 0..n_rows {
+                let x = cols[c][i].to_canonical_u64() as usize;
+                assert!(
+                    x < BYTE_RANGE_MAX,
+                    "column value {} exceeds the max range value {}",
+                    x,
+                    BYTE_RANGE_MAX
+                );
+                cols[RC_FREQUENCIES][x] += F::ONE;
+            }
+        }
+    }
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeStark<F, D> {
@@ -733,6 +768,14 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS
     fn constraint_degree(&self) -> usize {
         3
     }
+
+    fn lookups(&self) -> Vec<Lookup> {
+        vec![Lookup {
+            columns: get_block_bytes_range().collect(),
+            table_column: RANGE_COUNTER,
+            frequencies_column: RC_FREQUENCIES,
+        }]
+    }
 }
 
 #[cfg(test)]