[NIT-4075] Extract keccak for ZKVM precompile substitution

Makefile

@@ -60,7 +60,7 @@ replay_wasm=$(output_latest)/replay.wasm
 arb_brotli_files = $(wildcard arbitrator/brotli/src/*.* arbitrator/brotli/src/*/*.* arbitrator/brotli/*.toml arbitrator/brotli/*.rs) .make/cbrotli-lib .make/cbrotli-wasm
 
 arbitrator_generated_header=$(output_root)/include/arbitrator.h
-arbitrator_wasm_libs=$(patsubst %, $(output_root)/machines/latest/%.wasm, forward wasi_stub host_io soft-float arbcompress user_host program_exec)
+arbitrator_wasm_libs=$(patsubst %, $(output_root)/machines/latest/%.wasm, forward wasi_stub host_io soft-float arbcompress arbkeccak user_host program_exec)
 arbitrator_stylus_lib=$(output_root)/lib/libstylus.a
 prover_bin=$(output_root)/bin/prover
 arbitrator_jit=$(output_root)/bin/jit
@@ -453,6 +453,11 @@ $(output_latest)/arbcompress.wasm: $(DEP_PREDICATE) $(call wasm_lib_deps,brotli)
 	install arbitrator/wasm-libraries/$(wasm32_wasi)/arbcompress.wasm $@
 	./scripts/remove_reference_types.sh $@
 
+$(output_latest)/arbkeccak.wasm: $(DEP_PREDICATE) $(call wasm_lib_deps) $(wasm_lib_go_abi)
+	cargo build --manifest-path arbitrator/wasm-libraries/Cargo.toml --release --target wasm32-wasip1 --config $(wasm_lib_cargo) --package arbkeccak
+	install arbitrator/wasm-libraries/$(wasm32_wasi)/arbkeccak.wasm $@
+	./scripts/remove_reference_types.sh $@
+
 $(output_latest)/forward.wasm: $(DEP_PREDICATE) $(wasm_lib_forward) .make/machines
 	cargo run --manifest-path $(forward_dir)/Cargo.toml -- --path $(forward_dir)/forward.wat
 	wat2wasm $(wasm_lib)/forward/forward.wat -o $@
@@ -463,7 +468,7 @@ $(output_latest)/forward_stub.wasm: $(DEP_PREDICATE) $(wasm_lib_forward) .make/m
 
 $(output_latest)/machine.wavm.br: $(DEP_PREDICATE) $(prover_bin) $(arbitrator_wasm_libs) $(replay_wasm)
 	$(prover_bin) $(replay_wasm) --generate-binaries $(output_latest) \
-	$(patsubst %,-l $(output_latest)/%.wasm, forward soft-float wasi_stub host_io user_host arbcompress program_exec)
+	$(patsubst %,-l $(output_latest)/%.wasm, forward soft-float wasi_stub host_io user_host arbcompress arbkeccak program_exec)
 
 $(arbitrator_cases)/%.wasm: $(arbitrator_cases)/%.wat
 	wat2wasm $< -o $@

arbitrator/caller-env/Cargo.toml

@@ -13,6 +13,7 @@ rust-version.workspace = true
 brotli = { workspace = true, optional = true }
 rand = { workspace = true }
 rand_pcg = { workspace = true }
+tiny-keccak = { workspace = true, features = ["keccak"] }
 wasmer = { workspace = true, optional = true }
 
 [features]

arbitrator/caller-env/src/arbkeccak.rs

@@ -0,0 +1,23 @@
+use crate::{ExecEnv, GuestPtr, MemAccess};
+use core::mem::MaybeUninit;
+use tiny_keccak::{Hasher, Keccak};
+
+pub fn keccak256<M: MemAccess, E: ExecEnv>(
+    mem: &mut M,
+    _env: &mut E,
+    in_buf_ptr: GuestPtr,
+    in_buf_len: u32,
+    out_buf_ptr: GuestPtr,
+) {
+    let input = mem.read_slice(in_buf_ptr, in_buf_len as usize);
+
+    let mut output = MaybeUninit::<[u8; 32]>::uninit();
+    let mut hasher = Keccak::v256();
+    hasher.update(input.as_ref());
+
+    // SAFETY: finalize() writes 32 bytes
+    unsafe {
+        hasher.finalize(&mut *output.as_mut_ptr());
+        mem.write_slice(out_buf_ptr, output.assume_init().as_slice());
+    }
+}

arbitrator/caller-env/src/lib.rs

@@ -20,6 +20,8 @@ pub mod wasmer_traits;
 #[cfg(feature = "brotli")]
 pub mod brotli;
 
+pub mod arbkeccak;
+
 mod guest_ptr;
 pub mod wasip1_stub;
 

arbitrator/jit/src/arbkeccak.rs

@@ -0,0 +1,21 @@
+use crate::caller_env::{JitEnv, JitExecEnv};
+use crate::machine::{MaybeEscape, WasmEnvMut};
+use caller_env::GuestPtr;
+
+pub fn keccak256(
+    mut src: WasmEnvMut,
+    in_buf_ptr: GuestPtr,
+    in_buf_len: u32,
+    out_buf_ptr: GuestPtr,
+) -> MaybeEscape {
+    let (mut mem, wenv) = src.jit_env();
+
+    caller_env::arbkeccak::keccak256(
+        &mut mem,
+        &mut JitExecEnv { wenv },
+        in_buf_ptr,
+        in_buf_len,
+        out_buf_ptr,
+    );
+    Ok(())
+}

arbitrator/jit/src/lib.rs

@@ -5,6 +5,7 @@ use std::path::PathBuf;
 use structopt::StructOpt;
 
 mod arbcompress;
+mod arbkeccak;
 mod caller_env;
 pub mod machine;
 mod prepare;

arbitrator/jit/src/machine.rs

@@ -2,7 +2,7 @@
 // For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE.md
 
 use crate::{
-    arbcompress, caller_env::GoRuntimeState, prepare::prepare_env, program, socket,
+    arbcompress, arbkeccak, caller_env::GoRuntimeState, prepare::prepare_env, program, socket,
     stylus_backend::CothreadHandler, wasip1_stub, wavmio, Opts,
 };
 use arbutil::{Bytes32, Color, PreimageType};
@@ -69,6 +69,9 @@ pub fn create(opts: &Opts, env: WasmEnv) -> (Instance, FunctionEnv<WasmEnv>, Sto
             "brotli_compress" => func!(arbcompress::brotli_compress),
             "brotli_decompress" => func!(arbcompress::brotli_decompress),
         },
+        "arbkeccak" => {
+            "keccak256" => func!(arbkeccak::keccak256),
+        },
         "wavmio" => {
             "getGlobalStateBytes32" => func!(wavmio::get_global_state_bytes32),
             "setGlobalStateBytes32" => func!(wavmio::set_global_state_bytes32),

arbitrator/wasm-libraries/Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
 	"arbcompress",
+	"arbkeccak",
 	"wasi-stub",
 	"host-io",
 	"user-host",

arbitrator/wasm-libraries/arbkeccak/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "arbkeccak"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+caller-env = { path = "../../caller-env/", features = ["static_caller"] }

arbitrator/wasm-libraries/arbkeccak/build.rs

@@ -0,0 +1,8 @@
+// Copyright 2025, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE.md
+
+fn main() {
+    // Tell Cargo that if the given file changes, to rerun this build script.
+    println!("cargo:rustc-link-search=../../target/lib-wasm/");
+    println!("cargo:rustc-link-search=../target/lib/");
+}

arbitrator/wasm-libraries/arbkeccak/src/lib.rs

@@ -0,0 +1,21 @@
+// Copyright 2025, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE.md
+
+#![allow(clippy::missing_safety_doc)]
+
+use caller_env::{self, GuestPtr};
+
+#[no_mangle]
+pub unsafe extern "C" fn arbkeccak__keccak256(
+    in_buf_ptr: GuestPtr,
+    in_buf_len: u32,
+    out_buf_ptr: GuestPtr,
+) {
+    caller_env::arbkeccak::keccak256(
+        &mut caller_env::static_caller::STATIC_MEM,
+        &mut caller_env::static_caller::STATIC_ENV,
+        in_buf_ptr,
+        in_buf_len,
+        out_buf_ptr,
+    )
+}