Cleanup KeccakSpongeStark index accesses (#418)

0xPolygonZero · Jul 22, 2024 · c4f1289 · c4f1289
1 parent f531584
commit c4f1289
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 29 deletions.
diff --git a/evm_arithmetization/src/keccak_sponge/columns.rs b/evm_arithmetization/src/keccak_sponge/columns.rs
@@ -1,5 +1,4 @@
 use core::mem::{size_of, transmute};
-use core::ops::Range;
 
 use zk_evm_proc_macro::Columns;
 
@@ -84,8 +83,8 @@ pub(crate) struct KeccakSpongeColumnsView<T: Copy> {
     /// 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.
+    /// The counter column (used for the LogUp range check)
+    /// starts from 0 and increments.
     pub range_counter: T,
     /// The frequencies column used in logUp.
     pub rc_frequencies: T,
@@ -95,25 +94,6 @@ pub(crate) struct KeccakSpongeColumnsView<T: Copy> {
 /// Number of columns in `KeccakSpongeStark`.
 pub(crate) 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
-}
-
 const fn make_col_map() -> KeccakSpongeColumnsView<usize> {
     let indices_arr = indices_arr::<NUM_KECCAK_SPONGE_COLUMNS>();
     unsafe {

diff --git a/evm_arithmetization/src/keccak_sponge/keccak_sponge_stark.rs b/evm_arithmetization/src/keccak_sponge/keccak_sponge_stark.rs
@@ -507,17 +507,18 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
         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);
+            cols[KECCAK_SPONGE_COL_MAP.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);
+            cols[KECCAK_SPONGE_COL_MAP.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);
+            let c = KECCAK_SPONGE_COL_MAP.block_bytes[col];
             for i in 0..n_rows {
                 let x = cols[c][i].to_canonical_u64() as usize;
                 assert!(
@@ -526,7 +527,7 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
                     x,
                     BYTE_RANGE_MAX
                 );
-                cols[RC_FREQUENCIES][x] += F::ONE;
+                cols[KECCAK_SPONGE_COL_MAP.rc_frequencies][x] += F::ONE;
             }
         }
     }
@@ -943,9 +944,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS
 
     fn lookups(&self) -> Vec<Lookup<F>> {
         vec![Lookup {
-            columns: Column::singles(get_block_bytes_range()).collect(),
-            table_column: Column::single(RANGE_COUNTER),
-            frequencies_column: Column::single(RC_FREQUENCIES),
+            columns: Column::singles(KECCAK_SPONGE_COL_MAP.block_bytes).collect(),
+            table_column: Column::single(KECCAK_SPONGE_COL_MAP.range_counter),
+            frequencies_column: Column::single(KECCAK_SPONGE_COL_MAP.rc_frequencies),
             filter_columns: vec![Default::default(); KECCAK_RATE_BYTES],
         }]
     }