1 files changed, 12 insertions, 3 deletions
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();