feat(guests/keccak): Add ZKM precompile function call

4 files changed, 166 insertions, 9 deletions

diff --git a/guests/keccak/Cargo.lock b/guests/keccak/Cargo.lock
index 3d24abe..69063a0 100644
--- a/guests/keccak/Cargo.lock
+++ b/guests/keccak/Cargo.lock
@@ -3,6 +3,15 @@
 version = 3
 
 [[package]]
+name = "bincode"
+version = "1.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
+dependencies = [
+ "serde",
+]
+
+[[package]]
 name = "block-buffer"
 version = "0.10.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -12,6 +21,18 @@ dependencies = [
 ]
 
 [[package]]
+name = "bytemuck"
+version = "1.22.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6b1fc10dbac614ebc03540c9dbd60e83887fda27794998c6528f1782047d540"
+
+[[package]]
+name = "cfg-if"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
+
+[[package]]
 name = "cpufeatures"
 version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -57,6 +78,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "getrandom"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
+[[package]]
 name = "guests_macro"
 version = "0.1.0"
 dependencies = [
@@ -81,6 +113,15 @@ dependencies = [
 
 [[package]]
 name = "keccak"
+version = "0.1.0"
+dependencies = [
+ "guests_macro",
+ "sha3",
+ "zkm-runtime",
+]
+
+[[package]]
+name = "keccak"
 version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ecc2af9a1119c51f12a14607e783cb977bde58bc069ff0c3da1095e635d70654"
@@ -89,18 +130,39 @@ dependencies = [
 ]
 
 [[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
+[[package]]
 name = "libc"
 version = "0.2.171"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6"
 
 [[package]]
+name = "libm"
+version = "0.2.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"
+
+[[package]]
 name = "memchr"
 version = "2.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
 
 [[package]]
+name = "ppv-lite86"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
+dependencies = [
+ "zerocopy",
+]
+
+[[package]]
 name = "proc-macro2"
 version = "1.0.94"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -119,6 +181,36 @@ dependencies = [
 ]
 
 [[package]]
+name = "rand"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom",
+]
+
+[[package]]
 name = "serde"
 version = "1.0.219"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -149,10 +241,13 @@ dependencies = [
 
 [[package]]
 name = "sha2"
-version = "0.1.0"
+version = "0.10.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8"
 dependencies = [
- "guests_macro",
- "sha3",
+ "cfg-if",
+ "cpufeatures",
+ "digest",
 ]
 
 [[package]]
@@ -162,7 +257,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60"
 dependencies = [
  "digest",
- "keccak",
+ "keccak 0.1.5",
 ]
 
 [[package]]
@@ -229,6 +324,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
 
 [[package]]
+name = "wasi"
+version = "0.11.0+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
+
+[[package]]
 name = "winnow"
 version = "0.7.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -236,3 +337,49 @@ checksum = "63d3fcd9bba44b03821e7d699eeee959f3126dcc4aa8e4ae18ec617c2a5cea10"
 dependencies = [
  "memchr",
 ]
+
+[[package]]
+name = "zerocopy"
+version = "0.8.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "zkm-precompiles"
+version = "0.2.0"
+dependencies = [
+ "bincode",
+ "cfg-if",
+ "serde",
+ "sha2",
+]
+
+[[package]]
+name = "zkm-runtime"
+version = "0.2.0"
+dependencies = [
+ "bincode",
+ "bytemuck",
+ "cfg-if",
+ "getrandom",
+ "lazy_static",
+ "libm",
+ "rand",
+ "serde",
+ "sha2",
+ "zkm-precompiles",
+]
diff --git a/guests/keccak/Cargo.toml b/guests/keccak/Cargo.toml
index 5e4064a..beaa70e 100644
--- a/guests/keccak/Cargo.toml
+++ b/guests/keccak/Cargo.toml
@@ -7,11 +7,12 @@ edition = "2021"
 guests_macro = { version = "0.1.0", path = "../../guests_macro" }
 sha3 = { version = "0.10.8", default-features = false }
 
+zkm-runtime = { git = "https://github.com/zkMIPS/zkm.git", optional = true } # patched by the guest wrapper
 [features]
 no_std = []
 jolt = []
 nexus = []
 risc0 = []
 sp1 = []
-zkm = []
+zkm = ["zkm-runtime"]
 zkwasm = []
diff --git a/guests/keccak/default.env b/guests/keccak/default.env
index d09a649..1036f4a 100644
--- a/guests/keccak/default.env
+++ b/guests/keccak/default.env
@@ -1,2 +1,2 @@
 ## ZKM ##
-SEG_SIZE=3250
+SEG_SIZE=3400
diff --git a/guests/keccak/src/lib.rs b/guests/keccak/src/lib.rs
index 10815f4..4808442 100644
--- a/guests/keccak/src/lib.rs
+++ b/guests/keccak/src/lib.rs
@@ -5,13 +5,22 @@ extern crate alloc;
 #[cfg(feature = "no_std")]
 use alloc::vec::Vec;
 
+#[cfg(feature = "zkm")]
+use zkm_runtime::*;
+
 use sha3::{Digest, Keccak256};
 
 #[guests_macro::proving_entrypoint]
 pub fn main(secret: Vec<u8>, hash: Vec<u8>) -> bool {
-    let mut hasher = Keccak256::new();
-    hasher.update(secret);
-    let result = hasher.finalize();
+    #[cfg(feature = "zkm")]
+    let result = zkm_runtime::io::keccak(&secret.as_slice());
+
+    #[cfg(not(any(feature = "zkm")))]
+    let result = {
+        let mut hasher = Keccak256::new();
+        hasher.update(secret);
+        hasher.finalize()
+    };
 
     let output: [u8; 32] = result.into();