Merge pull request #691 from mir-protocol/remove_keccak_rust

Remove keccak_rust in favor of tiny-keccak

evm/Cargo.toml

@@ -21,8 +21,8 @@ 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"
+tiny-keccak = "2.0.2"
 
 [dev-dependencies]
 criterion = "0.3.5"

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
 }

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::PrimeField64;
-    use plonky2::field::types::Field;
     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] = [
+        let expected = {
-            input[0..5].try_into().unwrap(),
+            let mut state = input;
-            input[5..10].try_into().unwrap(),
+            keccakf(&mut state);
-            input[10..15].try_into().unwrap(),
+            state
-            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();
 
         assert_eq!(output, expected);