feat: add blake3 option for second public input

Cargo.toml

@@ -100,6 +100,7 @@ serde_json = "1.0.132"
 
 rand = "0.8.5"
 sha2 = { version = "0.10.8", default-features = false }
+blake3 = { version = "1.8.5", default-features = false }
 anyhow = "1.0.75"
 
 zkm-recursion-derive = { path = "crates/recursion/derive", default-features = false }

crates/build/src/build.rs

@@ -81,6 +81,11 @@ pub(crate) fn build_program_internal(path: &str, args: Option<BuildArgs>) {
     // Activate the build command if the dependencies change.
     cargo_rerun_if_changed(&metadata, program_dir);
 
+    // Also rebuild if `ZKM_IMM_WRAP_VK` changes, since it decides whether the guest is built with
+    // the `imm-wrap-vk` feature. Cargo only tracks what's declared here, so without this the guest
+    // would stay stale (built in the old mode) whenever the env var changes but no source changes.
+    println!("cargo:rerun-if-env-changed=ZKM_IMM_WRAP_VK");
+
     // Check if RUSTC_WORKSPACE_WRAPPER is set to clippy-driver (i.e. if `cargo clippy` is the
     // current compiler). If so, don't execute `cargo ziren build` because it breaks
     // rust-analyzer's `cargo clippy` feature.
@@ -99,18 +104,29 @@ pub(crate) fn build_program_internal(path: &str, args: Option<BuildArgs>) {
     }
 
     // Build the program with the given arguments.
-    let path_output = if let Some(args) = args {
-        execute_build_program(&args, Some(program_dir.to_path_buf()))
-    } else {
-        execute_build_program(&BuildArgs::default(), Some(program_dir.to_path_buf()))
-    };
+    let mut args = args.unwrap_or_default();
+    if imm_wrap_vk_mode() {
+        args.features.push("imm-wrap-vk".to_string());
+    }
+    let path_output = execute_build_program(&args, Some(program_dir.to_path_buf()));
     if let Err(err) = path_output {
         panic!("Failed to build Ziren program: {err}.");
     }
 
     println!("cargo:warning={} built at {}", root_package_name, current_datetime());
 }
 
+/// Returns true if the `ZKM_IMM_WRAP_VK` environment variable is enabled, mirroring
+/// `zkm_recursion_core::stark::zkm_imm_wrap_vk_mode`'s environment check. When enabled, the guest
+/// program is built with the `imm-wrap-vk` feature, which the guest program's `Cargo.toml` is
+/// expected to forward to `zkm-zkvm/imm-wrap-vk` so it hashes public values with BLAKE3 instead of
+/// SHA256, matching the Groth16 wrap circuit's immutable-vk mode.
+fn imm_wrap_vk_mode() -> bool {
+    std::env::var("ZKM_IMM_WRAP_VK")
+        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
+        .unwrap_or(false)
+}
+
 /// Collects the list of targets that would be built and their output ELF file paths.
 pub fn generate_elf_paths(
     metadata: &cargo_metadata::Metadata,

crates/primitives/Cargo.toml

@@ -21,3 +21,13 @@ p3-symmetric = { workspace = true }
 p3-monty-31 = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 sha2 = "0.10.8"
+blake3 = { version = "1.8.5", default-features = false }
+tracing = { workspace = true }
+
+[dev-dependencies]
+serial_test = "3.1.1"
+
+[features]
+# Hash `committed_values_digest` with BLAKE3 instead of SHA256, matching the guest's own switch
+# (see `zkm-zkvm`'s `imm-wrap-vk` feature) and the Groth16 wrap circuit's immutable-vk mode.
+imm-wrap-vk = []

crates/primitives/src/io.rs

@@ -3,6 +3,20 @@ use num_bigint::BigUint;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 
+/// Returns true if either the `ZKM_IMM_WRAP_VK` environment variable is set or the `imm-wrap-vk`
+/// feature is enabled.
+/// By default, the variable is disabled.
+pub fn zkm_imm_wrap_vk_mode() -> bool {
+    let value = std::env::var("ZKM_IMM_WRAP_VK").unwrap_or_else(|_| "false".to_string());
+    let enabled = value == "1" || value.to_lowercase() == "true" || cfg!(feature = "imm-wrap-vk");
+    if enabled {
+        tracing::warn!(
+            "`ZKM_IMM_WRAP_VK` environment variable or `imm-wrap-vk` feature is enabled."
+        );
+    }
+    enabled
+}
+
 /// Public values for the prover.
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct ZKMPublicValues {
@@ -53,10 +67,16 @@ impl ZKMPublicValues {
     }
 
     /// Hash the public values.
-    pub fn hash(&self) -> Vec<u8> {
-        let mut hasher = Sha256::new();
-        hasher.update(self.buffer.data.as_slice());
-        hasher.finalize().to_vec()
+    ///
+    /// Uses BLAKE3 in `imm-wrap-vk` mode, SHA256 otherwise (see [`zkm_imm_wrap_vk_mode`]).
+    pub fn hash(&self) -> [u8; 32] {
+        if zkm_imm_wrap_vk_mode() {
+            *blake3::hash(self.buffer.data.as_slice()).as_bytes()
+        } else {
+            let mut hasher = Sha256::new();
+            hasher.update(self.buffer.data.as_slice());
+            hasher.finalize().into()
+        }
     }
 
     /// Hash the public values, mask the top 3 bits and return a BigUint. Matches the implementation
@@ -67,10 +87,7 @@ impl ZKMPublicValues {
     /// ```
     pub fn hash_bn254(&self) -> BigUint {
         // Hash the public values.
-        let mut hasher = Sha256::new();
-        hasher.update(self.buffer.data.as_slice());
-        let hash_result = hasher.finalize();
-        let mut hash = hash_result.to_vec();
+        let mut hash = self.hash();
 
         // Mask the top 3 bits.
         hash[0] &= 0b00011111;
@@ -89,8 +106,13 @@ impl AsRef<[u8]> for ZKMPublicValues {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use serial_test::serial;
+
+    // `zkm_imm_wrap_vk_mode()` reads a process-wide env var, so any test that touches it is
+    // `#[serial]` to avoid racing with the others in this module.
 
     #[test]
+    #[serial]
     fn test_hash_public_values() {
         let test_hex = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
         let test_bytes = hex::decode(test_hex).unwrap();
@@ -104,4 +126,27 @@ mod tests {
 
         assert_eq!(hash, expected_hash_biguint);
     }
+
+    #[test]
+    #[serial]
+    fn test_hash_public_values_imm_wrap_vk() {
+        std::env::set_var("ZKM_IMM_WRAP_VK", "1");
+
+        let test_hex = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
+        let test_bytes = hex::decode(test_hex).unwrap();
+
+        let mut public_values = ZKMPublicValues::new();
+        public_values.write_slice(&test_bytes);
+
+        let hash = public_values.hash();
+        let expected_hash = *blake3::hash(&test_bytes).as_bytes();
+        assert_eq!(hash, expected_hash);
+
+        let hash_bn254 = public_values.hash_bn254();
+        let mut expected_masked = expected_hash;
+        expected_masked[0] &= 0b00011111;
+        assert_eq!(hash_bn254, BigUint::from_bytes_be(&expected_masked));
+
+        std::env::remove_var("ZKM_IMM_WRAP_VK");
+    }
 }

crates/prover/src/verify.rs

@@ -495,7 +495,7 @@ pub fn verify_groth16_bn254_public_inputs(
 }
 
 /// Compute the verification key hash committed into Groth16 public inputs.
-fn groth16_vk_hash(vk: &ZKMVerifyingKey) -> Result<BigUint> {
+pub fn groth16_vk_hash(vk: &ZKMVerifyingKey) -> Result<BigUint> {
     const PART_STARK_VK_BYTES: &[u8] = include_bytes!("../../verifier/bn254-vk/part_stark_vk.bin");
 
     let vk_hash = vk.hash_bn254();

crates/recursion/core/Cargo.toml

@@ -77,4 +77,6 @@ debug = ["zkm-core-machine/debug"]
 sys = ["zkm-core-machine/sys", "dep:glob", "dep:cc", "dep:cbindgen", "dep:pathdiff"]
 program_validation = ["dep:range-set-blaze", "dep:smallvec"]
 # The Groth16 verification key (vk) is not affected by the Ziren upgrade.
-imm-wrap-vk = []
+# Forwards to `zkm-primitives/imm-wrap-vk` so the two crates share a single canonical
+# `imm-wrap-vk` mode (see `zkm_imm_wrap_vk_mode` in `stark/utils.rs`).
+imm-wrap-vk = ["zkm-primitives/imm-wrap-vk"]

crates/recursion/core/src/stark/utils.rs

@@ -28,15 +28,12 @@ pub fn zkm_dev_mode() -> bool {
 /// inputs for verification.
 ///
 /// By default, the variable is disabled.
+///
+/// Delegates to `zkm_primitives::io::zkm_imm_wrap_vk_mode` so there is a single canonical source of
+/// truth: this crate's `imm-wrap-vk` feature forwards to `zkm-primitives/imm-wrap-vk` (see this
+/// crate's `Cargo.toml`), so enabling either crate's feature enables both.
 pub fn zkm_imm_wrap_vk_mode() -> bool {
-    let value = std::env::var("ZKM_IMM_WRAP_VK").unwrap_or_else(|_| "false".to_string());
-    let enabled = value == "1" || value.to_lowercase() == "true" || cfg!(feature = "imm-wrap-vk");
-    if enabled {
-        tracing::warn!(
-            "`ZKM_IMM_WRAP_VK` environment variable or `imm-wrap-vk` feature is enabled."
-        );
-    }
-    enabled
+    zkm_primitives::io::zkm_imm_wrap_vk_mode()
 }
 
 /// Combine the base vkey hash with `vk_commitment` and `pc_start` using a Poseidon2 permutation.

crates/sdk/Cargo.toml

@@ -48,7 +48,6 @@ zkm-primitives = { workspace = true }
 zkm-cuda = { workspace = true }
 itertools = { workspace = true }
 tonic = { version = "0.8.1", features = ["tls", "tls-roots", "transport"]}
-alloy-signer = { version = "1.0" }
 alloy-primitives = { version = "1.0", optional = true }
 num-bigint = "0.4.6"
 serde_json = "1.0.140"