tooling(recursion): count keccak hashes to verify a proof (excl. grin…

crypto/crypto/src/hash/platform_keccak.rs

@@ -60,7 +60,49 @@ mod imp {
 
 #[cfg(not(target_arch = "riscv64"))]
 mod imp {
-    pub type PlatformKeccak256 = sha3::Keccak256;
+    use digest::{
+        FixedOutput, FixedOutputReset, HashMarker, Output, OutputSizeUser, Reset, Update,
+    };
+
+    /// Host keccak-256: `sha3::Keccak256` plus a finalize counter for
+    /// [`crate::hash_metrics`] (the total "all hashes" verify metric). The
+    /// counter is a PURE SIDE EFFECT — every method forwards to the inner
+    /// `sha3::Keccak256`, so the digest is byte-identical to the bare hasher.
+    /// Only compiled on the host; the guest keeps the syscall passthrough above.
+    #[derive(Clone, Default)]
+    pub struct PlatformKeccak256(sha3::Keccak256);
+
+    impl HashMarker for PlatformKeccak256 {}
+
+    impl OutputSizeUser for PlatformKeccak256 {
+        type OutputSize = digest::typenum::U32;
+    }
+
+    impl Update for PlatformKeccak256 {
+        fn update(&mut self, data: &[u8]) {
+            Update::update(&mut self.0, data);
+        }
+    }
+
+    impl FixedOutput for PlatformKeccak256 {
+        fn finalize_into(self, out: &mut Output<Self>) {
+            crate::hash_metrics::count_total();
+            FixedOutput::finalize_into(self.0, out);
+        }
+    }
+
+    impl Reset for PlatformKeccak256 {
+        fn reset(&mut self) {
+            Reset::reset(&mut self.0);
+        }
+    }
+
+    impl FixedOutputReset for PlatformKeccak256 {
+        fn finalize_into_reset(&mut self, out: &mut Output<Self>) {
+            crate::hash_metrics::count_total();
+            FixedOutputReset::finalize_into_reset(&mut self.0, out);
+        }
+    }
 }
 
 pub use imp::PlatformKeccak256;

crypto/crypto/src/hash_metrics.rs

@@ -0,0 +1,119 @@
+//! Host-only keccak-hash counters — a metric for the cost of VERIFYING a proof
+//! (a proxy for the recursion guest's dominant work: keccak hashing).
+//!
+//! [`count_total`] fires on EVERY keccak-256 finalize (the host
+//! [`crate::hash::platform_keccak::PlatformKeccak256`] wrapper) — Merkle trees,
+//! the Fiat-Shamir transcript, the program-id/ELF fold, everything. The Merkle
+//! backend additionally splits its share via [`count_merkle`] (every Merkle
+//! finalize) and [`count_merkle_node`] (auth-path parent compressions), so a
+//! caller can report `nodes`, `leaves = merkle − nodes`, and
+//! `transcript+other = total − merkle`.
+//!
+//! GRINDING IS EXCLUDED at its call site: the verifier wraps `is_valid_nonce`
+//! with [`disable`]/re-[`enable`] (guarded by [`is_enabled`]), so the proof-of-
+//! work check's finalizes are not counted.
+//!
+//! Compiled OUT on the riscv64 guest (`#[cfg]`) so the guest — whose cycles we
+//! actually care about — pays nothing. On the host each counted site is one
+//! relaxed atomic load (negligible vs prove time; disabled by default).
+
+#[cfg(not(target_arch = "riscv64"))]
+mod host {
+    use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+
+    static ENABLED: AtomicBool = AtomicBool::new(false);
+    static TOTAL: AtomicU64 = AtomicU64::new(0);
+    static MERKLE: AtomicU64 = AtomicU64::new(0);
+    static MERKLE_NODES: AtomicU64 = AtomicU64::new(0);
+
+    /// Every keccak-256 finalize, from any site (called by the host
+    /// `PlatformKeccak256` wrapper). The headline "all hashes" number.
+    #[inline(always)]
+    pub fn count_total() {
+        if ENABLED.load(Ordering::Relaxed) {
+            TOTAL.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    /// A Merkle finalize (leaf or node). Subset of [`count_total`].
+    #[inline(always)]
+    pub fn count_merkle() {
+        if ENABLED.load(Ordering::Relaxed) {
+            MERKLE.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    /// A Merkle parent (auth-path) compression. Subset of [`count_merkle`].
+    #[inline(always)]
+    pub fn count_merkle_node() {
+        if ENABLED.load(Ordering::Relaxed) {
+            MERKLE_NODES.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    /// Whether counting is currently on — so the grinding exclusion can restore
+    /// the prior state instead of blindly re-enabling.
+    #[inline(always)]
+    pub fn is_enabled() -> bool {
+        ENABLED.load(Ordering::Relaxed)
+    }
+
+    pub fn enable() {
+        ENABLED.store(true, Ordering::Relaxed);
+    }
+
+    pub fn disable() {
+        ENABLED.store(false, Ordering::Relaxed);
+    }
+
+    /// Zero the counters (does not change the enabled state).
+    pub fn reset() {
+        TOTAL.store(0, Ordering::Relaxed);
+        MERKLE.store(0, Ordering::Relaxed);
+        MERKLE_NODES.store(0, Ordering::Relaxed);
+    }
+
+    /// `(total, merkle, merkle_nodes)`. leaves = merkle − nodes;
+    /// transcript+other = total − merkle.
+    pub fn snapshot() -> (u64, u64, u64) {
+        (
+            TOTAL.load(Ordering::Relaxed),
+            MERKLE.load(Ordering::Relaxed),
+            MERKLE_NODES.load(Ordering::Relaxed),
+        )
+    }
+}
+
+#[cfg(not(target_arch = "riscv64"))]
+pub use host::{
+    count_merkle, count_merkle_node, count_total, disable, enable, is_enabled, reset, snapshot,
+};
+
+// Guest stubs — compiled to nothing; the guest must not pay for measurement.
+// `enable`/`disable`/`is_enabled` exist too so the shared verifier code (which
+// wraps the grinding check) compiles for the guest without `#[cfg]` noise.
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn count_total() {}
+
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn count_merkle() {}
+
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn count_merkle_node() {}
+
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn enable() {}
+
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn disable() {}
+
+#[cfg(target_arch = "riscv64")]
+#[inline(always)]
+pub fn is_enabled() -> bool {
+    false
+}

crypto/crypto/src/lib.rs

@@ -9,6 +9,7 @@ extern crate alloc;
 
 pub mod fiat_shamir;
 pub mod hash;
+pub mod hash_metrics;
 pub mod merkle_tree;
 #[cfg(feature = "disk-spill")]
 pub mod mmap_util;

crypto/crypto/src/merkle_tree/backends/field_element_vector.rs

@@ -43,6 +43,9 @@ use lambda_vm_syscalls::keccak::Keccak256 as SyscallKeccak256;
 fn hash_streamed<D: Digest + 'static, const NUM_BYTES: usize>(
     feed: impl Fn(&mut dyn FnMut(&[u8])),
 ) -> [u8; NUM_BYTES] {
+    // Metric: a Merkle finalize (leaf or node) — the total keccak count is taken
+    // at the primitive; this is the Merkle sub-count. No-op on guest / disabled.
+    crate::hash_metrics::count_merkle();
     #[cfg(target_arch = "riscv64")]
     if NUM_BYTES == 32 && TypeId::of::<D>() == TypeId::of::<PlatformKeccak256>() {
         let mut hasher = SyscallKeccak256::new();
@@ -75,6 +78,10 @@ fn hash_new_parent_bytes<D: Digest + 'static, const NUM_BYTES: usize>(
     left: &[u8; NUM_BYTES],
     right: &[u8; NUM_BYTES],
 ) -> [u8; NUM_BYTES] {
+    // Metric: a Merkle parent (auth-path) compression. On the host this also
+    // flows through `hash_streamed` (one `count_merkle`), so merkle − nodes =
+    // leaves. No-op on the guest / when disabled.
+    crate::hash_metrics::count_merkle_node();
     #[cfg(target_arch = "riscv64")]
     if NUM_BYTES == 32 && TypeId::of::<D>() == TypeId::of::<PlatformKeccak256>() {
         let l: &[u8; 32] = left[..].try_into().unwrap();

crypto/stark/src/verifier.rs

@@ -1664,9 +1664,17 @@ pub trait IsStarkVerifier<
         // verify grinding
         let grinding_factor = air.context().proof_options.grinding_factor;
         if grinding_factor > 0 {
+            // Exclude the proof-of-work check from the host hash metric
+            // (`crypto::hash_metrics`): the metric asks for "all hashes except
+            // grinding". No-op on the guest and whenever counting is off.
+            let hm_was = crypto::hash_metrics::is_enabled();
+            crypto::hash_metrics::disable();
             let nonce_is_valid = proof.nonce().is_some_and(|nonce_value| {
                 grinding::is_valid_nonce(&challenges.grinding_seed, nonce_value, grinding_factor)
             });
+            if hm_was {
+                crypto::hash_metrics::enable();
+            }
 
             if !nonce_is_valid {
                 #[cfg(not(feature = "test_fiat_shamir"))]

prover/src/tests/recursion_smoke_test.rs

@@ -1052,6 +1052,61 @@ fn test_dump_recursion_input() {
     }
 }
 
+/// Count the keccak hashes done to VERIFY the dumped recursion blob — a
+/// prover-change metric: fewer hashes ⇒ a cheaper recursion guest. Runs the exact
+/// guest verify (`verify_continuation_and_attest`) on `/tmp/recursion_input.bin`
+/// (override with `RECURSION_INPUT_PATH`) with `crypto::hash_metrics` counting
+/// EVERY keccak-256 finalize EXCEPT the grinding proof-of-work check. The preset
+/// MUST match the dump's `RECURSION_DUMP_PRESET`, or the verify fails.
+///
+/// Loop: change the prover → re-run `test_dump_recursion_input` (re-proves +
+/// dumps the new blob) → run this (fast, verify-only) → compare the counts.
+///
+///   RECURSION_DUMP_PRESET=blowup4 cargo test --release \
+///     -p lambda-vm-prover --lib test_count_recursion_hashes -- --ignored --nocapture
+#[test]
+#[ignore = "diagnostic: counts keccak hashes verifying the dumped recursion blob"]
+fn test_count_recursion_hashes() {
+    let preset_name =
+        std::env::var("RECURSION_DUMP_PRESET").unwrap_or_else(|_| "blowup4".to_string());
+    let preset = Preset::ALL
+        .into_iter()
+        .find(|p| p.name() == preset_name)
+        .unwrap_or_else(|| panic!("unknown RECURSION_DUMP_PRESET '{preset_name}'"));
+    let path = std::env::var("RECURSION_INPUT_PATH")
+        .unwrap_or_else(|_| "/tmp/recursion_input.bin".to_string());
+    let blob = std::fs::read(&path)
+        .unwrap_or_else(|e| panic!("read {path} (run test_dump_recursion_input first): {e}"));
+
+    crypto::hash_metrics::reset();
+    crypto::hash_metrics::enable();
+    let attestation = recursion::verify_continuation_and_attest(&blob, &preset.options())
+        .expect("verify_continuation_and_attest errored");
+    crypto::hash_metrics::disable();
+    let (total, merkle, nodes) = crypto::hash_metrics::snapshot();
+
+    assert!(
+        attestation.is_some(),
+        "the blob must verify under preset '{}' — does it match the dump's RECURSION_DUMP_PRESET?",
+        preset.name()
+    );
+    // `total` = every keccak-256 finalize EXCEPT the grinding PoW check (excluded
+    // at its call site). `merkle` is its Merkle subset; `total - merkle` is the
+    // transcript + program-id/ELF fold + anything else.
+    println!(
+        "[hash-count] preset={} blob={}B fri_queries={} | total(excl. grinding)={} | \
+         merkle={} (nodes={} leaves={}) | transcript+other={}",
+        preset.name(),
+        blob.len(),
+        preset.options().fri_number_of_queries,
+        total,
+        merkle,
+        nodes,
+        merkle.saturating_sub(nodes),
+        total.saturating_sub(merkle),
+    );
+}
+
 /// Cycle count only of the recursion guest verifying a 1-query inner proof.
 #[test]
 #[ignore = "diagnostic: fast; recursion guest cycle count (1 query)"]