chore: use openvm guest library for zkvm target

sha3/Cargo.toml

@@ -19,6 +19,9 @@ categories = ["cryptography", "no-std"]
 digest = "0.10.4"
 keccak = "0.1.4"
 
+[target.'cfg(target_os = "zkvm")'.dependencies]
+openvm-keccak256 = { git = "https://github.com/openvm-org/openvm", branch = "develop-v1.6.0", default-features = false, features = [] }
+
 [dev-dependencies]
 digest = { version = "0.10.4", features = ["dev"] }
 hex-literal = "0.2.2"

sha3/src/lib.rs

@@ -99,7 +99,12 @@ const CSHAKE: u8 = 0x4;
 const TURBO_SHAKE_ROUND_COUNT: usize = 12;
 
 impl_sha3!(Keccak224Core, Keccak224, U28, U144, KECCAK, "Keccak-224");
+#[cfg(not(target_os = "zkvm"))]
 impl_sha3!(Keccak256Core, Keccak256, U32, U136, KECCAK, "Keccak-256");
+#[cfg(target_os = "zkvm")]
+mod zkvm_impl;
+#[cfg(target_os = "zkvm")]
+pub use zkvm_impl::Keccak256;
 impl_sha3!(Keccak384Core, Keccak384, U48, U104, KECCAK, "Keccak-384");
 impl_sha3!(Keccak512Core, Keccak512, U64, U72, KECCAK, "Keccak-512");
 

sha3/src/zkvm_impl.rs

@@ -0,0 +1,77 @@
+#[cfg(feature = "std")]
+extern crate std;
+
+use core::fmt;
+use digest::{
+    consts::{U136, U32},
+    core_api::BlockSizeUser,
+    FixedOutput, FixedOutputReset, HashMarker, Output, OutputSizeUser, Reset, Update,
+};
+use openvm_keccak256::Keccak256 as InnerKeccak256;
+
+/// Keccak-256 hasher backed by the openvm zkvm implementation.
+#[derive(Clone)]
+pub struct Keccak256 {
+    inner: InnerKeccak256,
+}
+
+impl Default for Keccak256 {
+    fn default() -> Self {
+        Self {
+            inner: InnerKeccak256::new(),
+        }
+    }
+}
+
+impl BlockSizeUser for Keccak256 {
+    type BlockSize = U136;
+}
+
+impl Update for Keccak256 {
+    fn update(&mut self, data: &[u8]) {
+        self.inner.update(data);
+    }
+}
+
+impl OutputSizeUser for Keccak256 {
+    type OutputSize = U32;
+}
+
+impl FixedOutput for Keccak256 {
+    fn finalize_into(self, out: &mut Output<Self>) {
+        self.inner.finalize(out);
+    }
+}
+
+impl HashMarker for Keccak256 {}
+
+impl Reset for Keccak256 {
+    fn reset(&mut self) {
+        *self = Self::default();
+    }
+}
+
+impl FixedOutputReset for Keccak256 {
+    fn finalize_into_reset(&mut self, out: &mut Output<Self>) {
+        FixedOutput::finalize_into(self.clone(), out);
+        Reset::reset(self);
+    }
+}
+
+impl fmt::Debug for Keccak256 {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("Keccak256 { ... }")
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::io::Write for Keccak256 {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        Update::update(self, buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}

sha3/zkvm_test/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "zkvm-host-test"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+openvm-sdk = { git = "https://github.com/openvm-org/openvm", branch = "develop-v1.6.0", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.3", default-features = false }
+openvm-build = { git = "https://github.com/openvm-org/openvm", branch = "develop-v1.6.0" }
+eyre = "0.6"
+
+[[bin]]
+name = "verify-keccak-chips"
+path = "src/main.rs"
+
+# Keep this package independent from the parent workspace
+[workspace]

sha3/zkvm_test/programs/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "zkvm-sha3-test"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+sha3 = { version = "0.10.8", default-features = false }
+openvm = { git = "https://github.com/openvm-org/openvm", branch = "develop-v1.6.0" }
+hex = { version = "0.4", default-features = false, features = ["alloc"] }
+
+# Keep this package independent from the parent workspace
+[workspace]
+
+[patch.crates-io]
+sha3 = { path = "../../" }

sha3/zkvm_test/programs/examples/keccak256.rs

@@ -0,0 +1,24 @@
+#![no_main]
+#![no_std]
+
+extern crate alloc;
+
+use alloc::vec::Vec;
+use hex::FromHex;
+use sha3::{Digest, Keccak256};
+
+openvm::entry!(main);
+
+fn main() {
+    let input = Vec::from_hex(
+        "E926AE8B0AF6E53176DBFFCC2A6B88C6BD765F939D3D178A9BDE9EF3AA131C61\
+         E31C1E42CDFAF4B4DCDE579A37E150EFBEF5555B4C1CB40439D835A724E2FAE7",
+    )
+    .unwrap();
+    let expected =
+        Vec::from_hex("574271CD13959E8DDEAE5BFBDB02A3FDF54F2BABFD0CBEB893082A974957D0C1").unwrap();
+
+    let output = Keccak256::digest(&input);
+
+    assert_eq!(output[..], expected[..]);
+}

sha3/zkvm_test/programs/openvm.toml

@@ -0,0 +1,3 @@
+[app_vm_config.rv32i]
+[app_vm_config.rv32m]
+[app_vm_config.keccak]

sha3/zkvm_test/src/main.rs

@@ -0,0 +1,83 @@
+use std::path::PathBuf;
+
+use eyre::{eyre, Result};
+use openvm_build::TargetFilter;
+use openvm_sdk::{config::SdkVmConfig, Sdk, StdIn};
+use openvm_stark_sdk::config::setup_tracing;
+
+const KECCAK_AIR_PREFIXES: &[&str] = &["Keccakf", "Xorin"];
+
+fn main() -> Result<()> {
+    setup_tracing();
+
+    let guest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("programs");
+    let openvm_toml_path = guest_dir.join("openvm.toml");
+    let openvm_toml = std::fs::read_to_string(&openvm_toml_path)
+        .map_err(|e| eyre!("Failed to read {openvm_toml_path:?}: {e}"))?;
+    let app_config = SdkVmConfig::from_toml(&openvm_toml)
+        .map_err(|e| eyre!("Failed to parse openvm.toml: {e}"))?;
+
+    let sdk = Sdk::new(app_config)?;
+    let target_filter = Some(TargetFilter {
+        name: "keccak256".to_string(),
+        kind: "example".to_string(),
+    });
+    let elf = sdk.build(Default::default(), &guest_dir, &target_filter, None)?;
+    let exe = sdk.convert_to_exe(elf)?;
+
+    // Create app_prover to get access to the VM (for metered execution)
+    // and the converted exe, without constructing them separately.
+    let app_prover = sdk.app_prover(exe)?;
+    let vm = app_prover.vm();
+    let exe = app_prover.exe();
+
+    let air_names: Vec<String> = vm.air_names().map(|s| s.to_string()).collect();
+
+    // Identify keccak chip AIRs by matching name prefixes
+    let keccak_airs: Vec<(usize, &str)> = air_names
+        .iter()
+        .enumerate()
+        .filter(|(_, name)| KECCAK_AIR_PREFIXES.iter().any(|p| name.contains(p)))
+        .map(|(idx, name)| (idx, name.as_str()))
+        .collect();
+
+    if keccak_airs.is_empty() {
+        return Err(eyre!(
+            "No keccak-related AIRs found. Is the keccak extension enabled in openvm.toml?"
+        ));
+    }
+
+    // Run metered execution to collect per-AIR trace heights
+    let ctx = vm.build_metered_ctx(&exe);
+    let interpreter = vm
+        .metered_interpreter(&exe)
+        .map_err(|e| eyre!("Failed to create metered interpreter: {e}"))?;
+    let (segments, _final_state) = interpreter
+        .execute_metered(StdIn::default(), ctx)
+        .map_err(|e| eyre!("Metered execution failed: {e}"))?;
+
+    // Verify that at least one keccak AIR has a non-zero trace height,
+    // which confirms the custom keccak opcodes were actually executed.
+    let mut any_keccak_used = false;
+
+    for (seg_idx, segment) in segments.iter().enumerate() {
+        println!("Segment {seg_idx} (num_insns: {}):", segment.num_insns);
+        for &(air_idx, air_name) in &keccak_airs {
+            let height = segment.trace_heights.get(air_idx).copied().unwrap_or(0);
+            println!("  {air_name} (idx {air_idx}): trace_height = {height}");
+            if height > 0 {
+                any_keccak_used = true;
+            }
+        }
+    }
+
+    if any_keccak_used {
+        println!("PASS: Keccak chips have non-zero trace heights.");
+        Ok(())
+    } else {
+        Err(eyre!(
+            "FAIL: All keccak chip trace heights are zero.\n\
+             The keccak chips are not being used; the patch may not be working correctly."
+        ))
+    }
+}

sha3/zkvm_test/test.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+echo "=== Running host test to verify keccak chip usage ==="
+cargo run --release --bin verify-keccak-chips
+echo "=== Host test passed: Keccak chips verified ==="