From 80d5e5375e18a7ea449168cecd7bfc31a46ff44e Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Fri, 3 Jun 2022 11:32:19 -0700
Subject: [PATCH] fixes, cleanup, and correctness test

---
 evm/Cargo.toml                 |  1 +
 evm/src/keccak/keccak_stark.rs | 71 +++++++++++++++++++++++-----------
 evm/src/keccak/registers.rs    | 14 +++----
 3 files changed, 56 insertions(+), 30 deletions(-)

diff --git a/evm/Cargo.toml b/evm/Cargo.toml
index d65fb30c..e282583e 100644
--- a/evm/Cargo.toml
+++ b/evm/Cargo.toml
@@ -14,3 +14,4 @@ log = "0.4.14"
 rayon = "1.5.1"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
+keccak-rust = { git = "https://github.com/npwardberkeley/keccak-rust" }
\ No newline at end of file
diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index 147d393c..87787e31 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -130,7 +130,7 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakStark<F, D> {
                     row[reg_a_prime(x, y, z)] = xor([
                         row[reg_a(x, y, z)],
                         row[reg_c((x + 4) % 5, z)],
-                        row[reg_c((x + 1) % 5, (z + 1) % 64)],
+                        row[reg_c((x + 1) % 5, (z + 64 - 1) % 64)],
                     ]);
                 }
             }
@@ -358,12 +358,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
     ) {
         let two = builder.two();
 
-        dbg!(builder.num_gates());
-
         eval_round_flags_recursively(builder, vars, yield_constr);
 
-        dbg!(builder.num_gates());
-
         // C_partial[x] = xor(A[x, 0], A[x, 1], A[x, 2])
         for x in 0..5 {
             for z in 0..64 {
@@ -378,8 +374,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
             }
         }
 
-        dbg!(builder.num_gates());
-
         // C[x] = xor(C_partial[x], A[x, 3], A[x, 4])
         for x in 0..5 {
             for z in 0..64 {
@@ -394,8 +388,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
             }
         }
 
-        dbg!(builder.num_gates());
-
         // A'[x, y] = xor(A[x, y], D[x])
         //          = xor(A[x, y], C[x - 1], ROT(C[x + 1], 1))
         for x in 0..5 {
@@ -414,8 +406,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
             }
         }
 
-        dbg!(builder.num_gates());
-
         // A''[x, y] = xor(B[x, y], andn(B[x + 1, y], B[x + 2, y])).
         for x in 0..5 {
             for y in 0..5 {
@@ -443,8 +433,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
             }
         }
 
-        dbg!(builder.num_gates());
-
         // A'''[0, 0] = A''[0, 0] XOR RC
         let a_prime_prime_0_0_bits = (0..64)
             .map(|i| vars.local_values[reg_a_prime_prime_0_0_bit(i)])
@@ -460,8 +448,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         let diff = builder.sub_extension(computed_a_prime_prime_0_0_hi, a_prime_prime_0_0_hi);
         yield_constr.constraint(builder, diff);
 
-        dbg!(builder.num_gates());
-
         let mut get_xored_bit = |i| {
             let mut rc_bit_i = builder.zero_extension();
             for r in 0..NUM_ROUNDS {
@@ -493,14 +479,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         );
         yield_constr.constraint(builder, diff);
 
-        dbg!(builder.num_gates());
-
         // Enforce that this round's output equals the next round's input.
         for x in 0..5 {
             for y in 0..5 {
-                dbg!(x);
-                dbg!(y);
-                dbg!(builder.num_gates());
                 let output_lo = vars.local_values[reg_a_prime_prime_prime(x, y)];
                 let output_hi = vars.local_values[reg_a_prime_prime_prime(x, y) + 1];
                 let input_bits = (0..64)
@@ -516,8 +497,6 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                 yield_constr.constraint_transition(builder, diff);
             }
         }
-
-        dbg!(builder.num_gates());
     }
 
     fn constraint_degree(&self) -> usize {
@@ -528,9 +507,13 @@ 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::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use plonky2::field::field_types::Field;
+    use rand::Rng;
 
-    use crate::keccak::keccak_stark::KeccakStark;
+    use crate::keccak::keccak_stark::{KeccakStark, INPUT_LIMBS};
+    use crate::keccak::registers::reg_a_prime_prime_prime;
     use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
 
     #[test]
@@ -559,4 +542,46 @@ mod tests {
         };
         test_stark_circuit_constraints::<F, C, S, D>(stark)
     }
+
+    #[test]
+    fn keccak_correctness_test() -> Result<()> {
+        let mut input: Vec<u64> = Vec::new();
+        let mut rng = rand::thread_rng();
+        for _ in 0..INPUT_LIMBS {
+            input.push(rng.gen());
+        }
+
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = KeccakStark<F, D>;
+
+        let stark = S {
+            f: Default::default(),
+        };
+
+        let rows = stark.generate_trace_rows(vec![input.clone().try_into().unwrap()]);
+        let last_row = rows[23];
+        let mut output = Vec::new();
+        let base = F::from_canonical_u64(1 << 32);
+        for x in 0..5 {
+            for y in 0..5 {
+                output.push(last_row[reg_a_prime_prime_prime(x, y)] + base * last_row[reg_a_prime_prime_prime(x, y) + 1]);
+            }
+        }
+
+        let mut expected: [[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 expected);
+        let expected_flattened: Vec<_> = expected.iter().flatten().map(|&x| F::from_canonical_u64(x)).collect();
+
+        assert_eq!(output, expected_flattened);
+
+        Ok(())
+    }
 }
diff --git a/evm/src/keccak/registers.rs b/evm/src/keccak/registers.rs
index f3e847b1..9a851eec 100644
--- a/evm/src/keccak/registers.rs
+++ b/evm/src/keccak/registers.rs
@@ -7,11 +7,11 @@ pub(crate) const fn reg_step(i: usize) -> usize {
 }
 
 const R: [[u8; 5]; 5] = [
-    [0, 18, 41, 3, 36],
-    [1, 2, 45, 10, 44],
-    [62, 61, 15, 43, 6],
-    [28, 56, 21, 25, 55],
-    [27, 14, 8, 39, 20],
+    [0, 36, 3, 41, 18],
+    [1, 44, 10, 45, 2],
+    [62, 6, 43, 15, 61],
+    [28, 55, 25, 21, 56],
+    [27, 20, 39, 8, 14],
 ];
 
 const RC: [u64; 24] = [
@@ -218,11 +218,10 @@ pub(crate) const fn reg_b(x: usize, y: usize, z: usize) -> usize {
     let a = (x + 3 * y) % 5;
     let b = x;
     let rot = R[a][b] as usize;
-    reg_a_prime(a, b, (z + rot) % 64)
+    reg_a_prime(a, b, (z + 64 - rot) % 64)
 }
 
 // A''[x, y] = xor(B[x, y], andn(B[x + 1, y], B[x + 2, y])).
-// A''[0, 0] is additionally xor'd with RC.
 const START_A_PRIME_PRIME: usize = START_A_PRIME + 5 * 5 * 64;
 pub(crate) const fn reg_a_prime_prime(x: usize, y: usize) -> usize {
     debug_assert!(x < 5);
@@ -239,6 +238,7 @@ pub(crate) const fn reg_a_prime_prime_0_0_bit(i: usize) -> usize {
 const REG_A_PRIME_PRIME_PRIME_0_0_LO: usize = START_A_PRIME_PRIME_0_0_BITS + 64;
 const REG_A_PRIME_PRIME_PRIME_0_0_HI: usize = REG_A_PRIME_PRIME_PRIME_0_0_LO + 1;
 
+// A'''[0, 0] is additionally xor'd with RC.
 pub(crate) const fn reg_a_prime_prime_prime(x: usize, y: usize) -> usize {
     debug_assert!(x < 5);
     debug_assert!(y < 5);