From fb34b098884d578195c742f4f5c5da5817aa398a Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 24 Aug 2022 15:35:52 -0700
Subject: [PATCH] Remove keccak_rust in favor of tiny-keccak

`keccak_rust` doesn't seem to have much usage, and it treats `x` as the major axis of its 5x5 input.  This is not exactly wrong, since Keccak itself doesn't have a notion of axis order. However, there is a convention for mapping bits of the cube to a flat list of bits, which is

> The mapping between the bits of `s` and those of `a` is `s[w(5y + x) + z] = a[x][y][z]`.

Obeying this convention would be awkward with `keccak_rust` - the words in memory would need to be transposed.
---
 evm/Cargo.toml                 |  1 -
 evm/src/keccak/columns.rs      | 14 ++++++++++----
 evm/src/keccak/keccak_stark.rs | 33 +++++++++++++--------------------
 3 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/evm/Cargo.toml b/evm/Cargo.toml
index e844da3a..0f34afed 100644
--- a/evm/Cargo.toml
+++ b/evm/Cargo.toml
@@ -21,7 +21,6 @@ maybe_rayon = { path = "../maybe_rayon" }
 rand = "0.8.5"
 rand_chacha = "0.3.1"
 rlp = "0.5.1"
-keccak-rust = { git = "https://github.com/npwardberkeley/keccak-rust" }
 keccak-hash = "0.9.0"
 
 [dev-dependencies]
diff --git a/evm/src/keccak/columns.rs b/evm/src/keccak/columns.rs
index 39116b4a..8313c676 100644
--- a/evm/src/keccak/columns.rs
+++ b/evm/src/keccak/columns.rs
@@ -15,8 +15,11 @@ pub const fn reg_step(i: usize) -> usize {
 pub fn reg_input_limb<F: Field>(i: usize) -> Column<F> {
     debug_assert!(i < 2 * NUM_INPUTS);
     let i_u64 = i / 2; // The index of the 64-bit chunk.
-    let x = i_u64 / 5;
-    let y = i_u64 % 5;
+
+    // The 5x5 state is treated as y-major, as per the Keccak spec.
+    let y = i_u64 / 5;
+    let x = i_u64 % 5;
+
     let reg_low_limb = reg_a(x, y);
     let is_high_limb = i % 2;
     Column::single(reg_low_limb + is_high_limb)
@@ -28,8 +31,11 @@ pub fn reg_input_limb<F: Field>(i: usize) -> Column<F> {
 pub const fn reg_output_limb(i: usize) -> usize {
     debug_assert!(i < 2 * NUM_INPUTS);
     let i_u64 = i / 2; // The index of the 64-bit chunk.
-    let x = i_u64 / 5;
-    let y = i_u64 % 5;
+
+    // The 5x5 state is treated as y-major, as per the Keccak spec.
+    let y = i_u64 / 5;
+    let x = i_u64 % 5;
+
     let is_high_limb = i % 2;
     reg_a_prime_prime_prime(x, y) + is_high_limb
 }
diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index 94fa795d..d405c0e8 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -76,7 +76,7 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakStark<F, D> {
 
         for x in 0..5 {
             for y in 0..5 {
-                let input_xy = input[x * 5 + y];
+                let input_xy = input[y * 5 + x];
                 let reg_lo = reg_a(x, y);
                 let reg_hi = reg_lo + 1;
                 rows[0][reg_lo] = F::from_canonical_u64(input_xy & 0xFFFFFFFF);
@@ -547,9 +547,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
-    use keccak_rust::{KeccakF, StateBitsWidth};
-    use plonky2::field::types::Field;
+    use plonky2::field::types::PrimeField64;
     use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use tiny_keccak::keccakf;
 
     use crate::keccak::columns::reg_output_limb;
     use crate::keccak::keccak_stark::{KeccakStark, NUM_INPUTS, NUM_ROUNDS};
@@ -596,26 +596,19 @@ mod tests {
 
         let rows = stark.generate_trace_rows(vec![input.try_into().unwrap()]);
         let last_row = rows[NUM_ROUNDS - 1];
-        let base = F::from_canonical_u64(1 << 32);
         let output = (0..NUM_INPUTS)
-            .map(|i| last_row[reg_output_limb(2 * i)] + base * last_row[reg_output_limb(2 * i + 1)])
+            .map(|i| {
+                let hi = last_row[reg_output_limb(2 * i + 1)].to_canonical_u64();
+                let lo = last_row[reg_output_limb(2 * i)].to_canonical_u64();
+                (hi << 32) | lo
+            })
             .collect::<Vec<_>>();
 
-        let mut keccak_input: [[u64; 5]; 5] = [
-            input[0..5].try_into().unwrap(),
-            input[5..10].try_into().unwrap(),
-            input[10..15].try_into().unwrap(),
-            input[15..20].try_into().unwrap(),
-            input[20..25].try_into().unwrap(),
-        ];
-
-        let keccak = KeccakF::new(StateBitsWidth::F1600);
-        keccak.permutations(&mut keccak_input);
-        let expected: Vec<_> = keccak_input
-            .iter()
-            .flatten()
-            .map(|&x| F::from_canonical_u64(x))
-            .collect();
+        let expected = {
+            let mut state = input;
+            keccakf(&mut state);
+            state
+        };
 
         assert_eq!(output, expected);