Merge pull request #1958 from o1-labs/keccak-interpreter/tests

[easy] Add tests for zkVM Keccak

Cargo.toml

@@ -68,6 +68,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0.79"
 serde_with = "3.6.0"
 sha2 = "0.10.0"
+sha3 = "0.10.8"
 strum = "0.26.1"
 strum_macros = "0.26.1"
 syn = { version = "1.0.109", features = ["full"] }

optimism/Cargo.toml

@@ -46,3 +46,4 @@ os_pipe.workspace = true
 rand.workspace = true
 libc.workspace = true
 rayon.workspace = true
+sha3.workspace = true

optimism/src/keccak/mod.rs

@@ -10,6 +10,8 @@ pub mod column;
 pub mod constraints;
 pub mod environment;
 pub mod interpreter;
+#[cfg(test)]
+pub mod tests;
 pub mod witness;
 
 /// Desired output length of the hash in bits

optimism/src/keccak/tests.rs

@@ -0,0 +1,53 @@
+use crate::keccak::{environment::KeccakEnv, interpreter::KeccakInterpreter};
+use kimchi::o1_utils::{self, FieldHelpers};
+use mina_curves::pasta::Fp;
+use rand::Rng;
+use sha3::{Digest, Keccak256};
+
+#[test]
+fn test_pad_blocks() {
+    let blocks_1 = crate::keccak::pad_blocks::<Fp>(1);
+    assert_eq!(blocks_1[0], Fp::from(0x00));
+    assert_eq!(blocks_1[1], Fp::from(0x00));
+    assert_eq!(blocks_1[2], Fp::from(0x00));
+    assert_eq!(blocks_1[3], Fp::from(0x00));
+    assert_eq!(blocks_1[4], Fp::from(0x81));
+
+    let blocks_136 = crate::keccak::pad_blocks::<Fp>(136);
+    assert_eq!(blocks_136[0], Fp::from(0x010000000000000000000000u128));
+    assert_eq!(blocks_136[1], Fp::from(0x00));
+    assert_eq!(blocks_136[2], Fp::from(0x00));
+    assert_eq!(blocks_136[3], Fp::from(0x00));
+    assert_eq!(blocks_136[4], Fp::from(0x80));
+}
+
+#[test]
+fn test_keccak_witness_satisfies_constraints() {
+    let mut rng = o1_utils::tests::make_test_rng();
+
+    // Generate random bytelength and preimage for Keccak
+    let bytelength = rng.gen_range(1..1000);
+    let preimage: Vec<u8> = (0..bytelength).map(|_| rng.gen()).collect();
+    // Use an external library to compute the hash
+    let mut hasher = Keccak256::new();
+    hasher.update(&preimage);
+    let hash = hasher.finalize();
+
+    // Initialize the environment and run the interpreter
+    let mut keccak_env = KeccakEnv::<Fp>::new(0, &preimage);
+    while keccak_env.keccak_step.is_some() {
+        keccak_env.step();
+        // Simulate the constraints for each row (still checks nothing about lookups)
+        keccak_env.witness_env.constraints();
+    }
+    // Extract the hash from the witness
+    let output = keccak_env.witness_env.sponge_bytes()[0..32]
+        .iter()
+        .map(|byte| byte.to_bytes()[0])
+        .collect::<Vec<_>>();
+
+    // Check that the hash matches
+    for (i, byte) in output.iter().enumerate() {
+        assert_eq!(*byte, hash[i]);
+    }
+}