test for correctness of polygon keccak, separate script for running poseidon hash added

hash/plonky2/bench/build.sh

@@ -4,7 +4,6 @@
 rustup override set nightly
 
 # Build 
-# RUSTFLAGS=-Ctarget-cpu=native cargo build
 RUSTFLAGS=-Ctarget-cpu=native cargo build --release --bin plonky2_hash_benchmarks
 
 

hash/plonky2/bench/run.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 if [ -z ${ZKBENCH_HASH_TYPE} ]; then 
-ZKBENCH_HASH_TYPE="poseidon"    
+ZKBENCH_HASH_TYPE="sha256"    
 fi
 
 if [ -z ${ZKBENCH_INPUT_SIZE_BYTES} ]; then 
-ZKBENCH_INPUT_SIZE_BYTES=4
+ZKBENCH_INPUT_SIZE_BYTES=256
 fi
 
 echo "Running benchmarks with the following configurations:"

hash/plonky2/bench/run_tree.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+if [ -z ${ZKBENCH_HASH_TYPE_TREE} ]; then 
+ZKBENCH_HASH_TYPE_TREE="poseidon"    
+fi
+
+if [ -z ${ZKBENCH_TREE_DEPTH} ]; then 
+ZKBENCH_TREE_DEPTH=4
+fi
+
+echo "Running benchmarks with the following configurations:"
+echo "HASH            = $ZKBENCH_HASH_TYPE_TREE"
+echo "Tree Depth = $ZKBENCH_TREE_DEPTH"
+
+# Run the benchmarks
+./target/release/plonky2_hash_benchmarks $ZKBENCH_HASH_TYPE_TREE $ZKBENCH_TREE_DEPTH

hash/plonky2/bench/src/bench/keccak256/keccak_polygon.rs

@@ -25,9 +25,23 @@ use env_logger::{
     try_init_from_env
 };
 
-const NUM_INPUTS: usize = 25;
+#[allow(dead_code)]
+/// Number of rounds in a Keccak permutation.
+pub(crate) const NUM_ROUNDS: usize = 24;
 
-pub fn keccak_polygon_bench(num_perms: usize) -> Result<()> {
+/// Number of 64-bit elements in the Keccak permutation input.
+pub(crate) const NUM_INPUTS: usize = 25;
+
+fn ceil_div(a: usize, b: usize) -> usize {
+    (a + b - 1) / b
+}
+
+pub fn keccak_polygon_bench(size: usize) -> Result<()> {
+    // here input in in terms of block
+    let mut num_perms = 1;
+    if size > 25 * 64/8 {
+        num_perms = ceil_div(size, 25 * 64/8) as usize;
+    }
 
     const D: usize = 2;
     type C = PoseidonGoldilocksConfig;
@@ -99,4 +113,47 @@ pub fn keccak_polygon_bench(num_perms: usize) -> Result<()> {
 
 fn init_logger() {
     let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "debug"));
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::field::types::PrimeField64;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use tiny_keccak::keccakf;
+    use evm_arithmetization::keccak::columns::reg_output_limb;
+    use super::*;
+    const NUM_ROUNDS: usize = 24;
+
+    #[test]
+    fn keccak_correctness_test() -> Result<()> {
+        let input: [u64; NUM_INPUTS] = rand::random();
+
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = KeccakStark<F, D>;
+
+        let stark: KeccakStark<<PoseidonGoldilocksConfig as GenericConfig<2>>::F, 2> = S::default();
+
+        let rows = stark.generate_trace_rows(vec![(input, 0)], 8);
+        let last_row = rows[NUM_ROUNDS - 1];
+        let output = (0..NUM_INPUTS)
+            .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 expected = {
+            let mut state = input;
+            keccakf(&mut state);
+            state
+        };
+
+        assert_eq!(output, expected);
+
+        Ok(())
+    }
 }

hash/plonky2/bench/src/bench/poseidon.rs

@@ -41,7 +41,7 @@ pub fn poseidon_bench(depth: usize) -> Result<()> {
     let hash = builder.hash_or_noop::<PoseidonHash>(initial.clone());
 
     // Public inputs are the initial value (provided below) and the result (which is generated).
-    // builder.register_public_inputs(initial.clone().as_slice());
+    builder.register_public_inputs(initial.clone().as_slice());
     builder.register_public_input(hash.elements[0]);
     builder.register_public_input(hash.elements[1]);
     builder.register_public_input(hash.elements[2]);

hash/plonky2/bench/src/main.rs

@@ -58,7 +58,6 @@ fn main() {
 
         "keccak-polygon" => {
             println!("Running keccak of plolygon zk_evm: ");
-            eprintln!("number of permutation: {:?}", size);
             let _ = keccak_polygon_bench(size);
         }
 

hash/plonky2/external/zk_evm

@@ -1 +1 @@
-Subproject commit 29f8620338f2b3bb8be62def5666a2d9d7ba5331
+Subproject commit 5511569d166c29cc85211df31bb9485870ca53ab