feat: proving-scoped zk-alloc arena allocator (benchmark feature)

bin/ethlambda/Cargo.toml

@@ -33,5 +33,14 @@ eyre.workspace = true
 
 tikv-jemallocator.workspace = true
 
+# Only used by the `zk-alloc` feature, to install the proving-scoped global
+# allocator and run leanVM's startup core-count assertion.
+ethlambda-crypto = { workspace = true, optional = true }
+
+[features]
+# Benchmark-only: swap jemalloc for leanVM's arena allocator, scoped to proving
+# threads. Drops jemalloc + /debug/pprof heap profiling for this build.
+zk-alloc = ["dep:ethlambda-crypto", "ethlambda-crypto/zk-alloc"]
+
 [build-dependencies]
 vergen-git2.workspace = true

bin/ethlambda/src/main.rs

@@ -1,15 +1,22 @@
 mod checkpoint_sync;
 mod version;
 
-#[cfg(not(target_env = "msvc"))]
+// Default allocator: jemalloc with heap profiling. Under the `zk-alloc`
+// benchmark feature this is replaced by leanVM's proving-scoped arena allocator
+// (`ethlambda_crypto::ScopedZkAlloc`), which is incompatible with jemalloc.
+#[cfg(all(not(target_env = "msvc"), not(feature = "zk-alloc")))]
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-#[cfg(not(target_env = "msvc"))]
+#[cfg(all(not(target_env = "msvc"), not(feature = "zk-alloc")))]
 #[allow(non_upper_case_globals)]
 #[unsafe(export_name = "malloc_conf")]
 static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0";
 
+#[cfg(feature = "zk-alloc")]
+#[global_allocator]
+static ALLOC: ethlambda_crypto::ScopedZkAlloc = ethlambda_crypto::ScopedZkAlloc;
+
 use std::{
     collections::{BTreeMap, HashMap},
     net::{IpAddr, SocketAddr},
@@ -162,10 +169,21 @@ async fn main() -> eyre::Result<()> {
     })?;
     let p2p_socket = SocketAddr::new(IpAddr::from([0, 0, 0, 0]), options.gossipsub_port);
 
-    #[cfg(not(target_env = "msvc"))]
+    #[cfg(all(not(target_env = "msvc"), not(feature = "zk-alloc")))]
     info!("Using jemalloc allocator with heap profiling enabled");
-    #[cfg(target_env = "msvc")]
+    #[cfg(all(target_env = "msvc", not(feature = "zk-alloc")))]
     info!("Using system allocator (MSVC target)");
+    #[cfg(feature = "zk-alloc")]
+    {
+        // Asserts available_parallelism() == the core count this binary was
+        // built for; panics on mismatch. The binary must be built on (or for)
+        // the machine it runs on. See ethlambda_crypto::zk_alloc.
+        ethlambda_crypto::init_allocator();
+        // Build the global rayon pool with arena-flagged workers before any
+        // other rayon user (leanVM's setup_prover) creates it unflagged.
+        ethlambda_crypto::init_arena_rayon_pool();
+        info!("Using zk-alloc arena allocator (proving-scoped, benchmark build)");
+    }
 
     info!(node_key=?options.node_key, "got node key");
 

crates/common/crypto/Cargo.toml

@@ -20,5 +20,14 @@ leansig.workspace = true
 thiserror.workspace = true
 rand.workspace = true
 
+# Only needed by the `zk-alloc` feature, to mark the global rayon pool's prover
+# threads around an arena phase.
+rayon = { workspace = true, optional = true }
+
+[features]
+# Benchmark-only: route leanVM's bump-arena allocator to proving threads via a
+# scoped global allocator. See `src/zk_alloc.rs`. OFF by default.
+zk-alloc = ["dep:rayon"]
+
 [dev-dependencies]
 hex.workspace = true

crates/common/crypto/src/lib.rs

@@ -32,6 +32,72 @@ pub fn ensure_verifier_ready() {
     VERIFIER_INIT.call_once(setup_verifier);
 }
 
+#[cfg(feature = "zk-alloc")]
+mod zk_alloc;
+#[cfg(feature = "zk-alloc")]
+pub use zk_alloc::{ScopedZkAlloc, init_allocator, init_arena_rayon_pool};
+
+// Exercise the real arena path in this crate's test binary: the aggregate/verify
+// round-trip tests below then allocate through `ScopedZkAlloc`, so proving
+// allocations actually hit leanVM's arena and outputs must survive serialization.
+#[cfg(all(test, feature = "zk-alloc"))]
+#[global_allocator]
+static TEST_ALLOC: ScopedZkAlloc = ScopedZkAlloc;
+
+/// Run a Type-1 prover call, then serialize the proof to its on-wire bytes.
+///
+/// Under the `zk-alloc` feature the prover call runs inside an arena phase
+/// (serialized behind a global proving lock), and serialization happens *after*
+/// the phase ends so the returned bytes land in the system allocator and survive
+/// the next phase's slab reset. Without the feature this is just
+/// `ensure_prover_ready` + `produce` + serialize.
+#[cfg(feature = "zk-alloc")]
+fn prove_type1<F>(produce: F) -> Result<ByteList512KiB, AggregationError>
+where
+    F: FnOnce() -> Result<LMType1, AggregationError>,
+{
+    // Must precede `ensure_prover_ready`: `setup_prover` is the first rayon user
+    // and would otherwise build an unflagged global pool.
+    zk_alloc::init_arena_rayon_pool();
+    let session = zk_alloc::ArenaSession::begin();
+    ensure_prover_ready();
+    let proof = session.prove(produce);
+    compress_type1_to_byte_list(&proof?)
+}
+
+#[cfg(not(feature = "zk-alloc"))]
+fn prove_type1<F>(produce: F) -> Result<ByteList512KiB, AggregationError>
+where
+    F: FnOnce() -> Result<LMType1, AggregationError>,
+{
+    ensure_prover_ready();
+    compress_type1_to_byte_list(&produce()?)
+}
+
+/// Type-2 counterpart of [`prove_type1`].
+#[cfg(feature = "zk-alloc")]
+fn prove_type2<F>(produce: F) -> Result<ByteList512KiB, AggregationError>
+where
+    F: FnOnce() -> Result<LMType2, AggregationError>,
+{
+    // Must precede `ensure_prover_ready`: `setup_prover` is the first rayon user
+    // and would otherwise build an unflagged global pool.
+    zk_alloc::init_arena_rayon_pool();
+    let session = zk_alloc::ArenaSession::begin();
+    ensure_prover_ready();
+    let proof = session.prove(produce);
+    compress_type2_to_byte_list(&proof?)
+}
+
+#[cfg(not(feature = "zk-alloc"))]
+fn prove_type2<F>(produce: F) -> Result<ByteList512KiB, AggregationError>
+where
+    F: FnOnce() -> Result<LMType2, AggregationError>,
+{
+    ensure_prover_ready();
+    compress_type2_to_byte_list(&produce()?)
+}
+
 /// Error type for signature aggregation operations.
 #[derive(Debug, Error)]
 pub enum AggregationError {
@@ -164,18 +230,16 @@ pub fn aggregate_signatures(
         return Err(AggregationError::EmptyInput);
     }
 
-    ensure_prover_ready();
-
-    let raw_xmss: Vec<(LeanSigPubKey, LeanSigSignature)> = public_keys
-        .into_iter()
-        .zip(signatures)
-        .map(|(pk, sig)| (pk.into_inner(), sig.into_inner()))
-        .collect();
-
-    let proof = aggregate_type_1(&[], raw_xmss, message.0, slot, LOG_INV_RATE)
-        .map_err(|err| AggregationError::ProverFailure(err.to_string()))?;
+    prove_type1(move || {
+        let raw_xmss: Vec<(LeanSigPubKey, LeanSigSignature)> = public_keys
+            .into_iter()
+            .zip(signatures)
+            .map(|(pk, sig)| (pk.into_inner(), sig.into_inner()))
+            .collect();
 
-    compress_type1_to_byte_list(&proof)
+        aggregate_type_1(&[], raw_xmss, message.0, slot, LOG_INV_RATE)
+            .map_err(|err| AggregationError::ProverFailure(err.to_string()))
+    })
 }
 
 /// Aggregate both existing Type-1 proofs (children) and raw XMSS signatures.
@@ -202,24 +266,22 @@ pub fn aggregate_mixed(
         return Err(AggregationError::InsufficientChildren(children.len()));
     }
 
-    ensure_prover_ready();
-
-    let children_native: Vec<LMType1> = children
-        .into_iter()
-        .enumerate()
-        .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
-        .collect::<Result<_, _>>()?;
-
-    let raw_xmss: Vec<(LeanSigPubKey, LeanSigSignature)> = raw_public_keys
-        .into_iter()
-        .zip(raw_signatures)
-        .map(|(pk, sig)| (pk.into_inner(), sig.into_inner()))
-        .collect();
-
-    let proof = aggregate_type_1(&children_native, raw_xmss, message.0, slot, LOG_INV_RATE)
-        .map_err(|err| AggregationError::ProverFailure(err.to_string()))?;
-
-    compress_type1_to_byte_list(&proof)
+    prove_type1(move || {
+        let children_native: Vec<LMType1> = children
+            .into_iter()
+            .enumerate()
+            .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
+            .collect::<Result<_, _>>()?;
+
+        let raw_xmss: Vec<(LeanSigPubKey, LeanSigSignature)> = raw_public_keys
+            .into_iter()
+            .zip(raw_signatures)
+            .map(|(pk, sig)| (pk.into_inner(), sig.into_inner()))
+            .collect();
+
+        aggregate_type_1(&children_native, raw_xmss, message.0, slot, LOG_INV_RATE)
+            .map_err(|err| AggregationError::ProverFailure(err.to_string()))
+    })
 }
 
 /// Recursively aggregate two or more already-aggregated Type-1 proofs into one.
@@ -235,18 +297,16 @@ pub fn aggregate_proofs(
         return Err(AggregationError::InsufficientChildren(children.len()));
     }
 
-    ensure_prover_ready();
-
-    let children_native: Vec<LMType1> = children
-        .into_iter()
-        .enumerate()
-        .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
-        .collect::<Result<_, _>>()?;
+    prove_type1(move || {
+        let children_native: Vec<LMType1> = children
+            .into_iter()
+            .enumerate()
+            .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
+            .collect::<Result<_, _>>()?;
 
-    let proof = aggregate_type_1(&children_native, vec![], message.0, slot, LOG_INV_RATE)
-        .map_err(|err| AggregationError::ProverFailure(err.to_string()))?;
-
-    compress_type1_to_byte_list(&proof)
+        aggregate_type_1(&children_native, vec![], message.0, slot, LOG_INV_RATE)
+            .map_err(|err| AggregationError::ProverFailure(err.to_string()))
+    })
 }
 
 /// Verify a Type-1 aggregated signature proof.
@@ -299,18 +359,16 @@ pub fn merge_type_1s_into_type_2(
         return Err(AggregationError::EmptyInput);
     }
 
-    ensure_prover_ready();
-
-    let type_1s_native: Vec<LMType1> = type_1s
-        .into_iter()
-        .enumerate()
-        .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
-        .collect::<Result<_, _>>()?;
-
-    let merged = merge_many_type_1(type_1s_native, LOG_INV_RATE)
-        .map_err(|err| AggregationError::ProverFailure(err.to_string()))?;
+    prove_type2(move || {
+        let type_1s_native: Vec<LMType1> = type_1s
+            .into_iter()
+            .enumerate()
+            .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i))
+            .collect::<Result<_, _>>()?;
 
-    compress_type2_to_byte_list(&merged)
+        merge_many_type_1(type_1s_native, LOG_INV_RATE)
+            .map_err(|err| AggregationError::ProverFailure(err.to_string()))
+    })
 }
 
 /// Verify a Type-2 merged proof against the per-component expected bindings.
@@ -380,32 +438,30 @@ pub fn split_type_2_by_message(
     pubkeys_per_component: Vec<Vec<ValidatorPublicKey>>,
     message: &H256,
 ) -> Result<ByteList512KiB, AggregationError> {
-    ensure_prover_ready();
-
-    let pubkeys_per_info: Vec<Vec<LeanSigPubKey>> = pubkeys_per_component
-        .into_iter()
-        .map(into_lean_pubkeys)
-        .collect();
-
-    let type_2 = LMType2::decompress_without_pubkeys(proof_data, pubkeys_per_info)
-        .ok_or(AggregationError::DeserializationFailed)?;
-
-    let matches: Vec<usize> = type_2
-        .info
-        .iter()
-        .enumerate()
-        .filter_map(|(i, info)| (info.without_pubkeys.message == message.0).then_some(i))
-        .collect();
-    let index = match matches.as_slice() {
-        [i] => *i,
-        [] => return Err(AggregationError::UnknownMessage),
-        _ => return Err(AggregationError::MultipleMessages),
-    };
-
-    let component = split_type_2(type_2, index, LOG_INV_RATE)
-        .map_err(|err| AggregationError::ProverFailure(err.to_string()))?;
-
-    compress_type1_to_byte_list(&component)
+    prove_type1(move || {
+        let pubkeys_per_info: Vec<Vec<LeanSigPubKey>> = pubkeys_per_component
+            .into_iter()
+            .map(into_lean_pubkeys)
+            .collect();
+
+        let type_2 = LMType2::decompress_without_pubkeys(proof_data, pubkeys_per_info)
+            .ok_or(AggregationError::DeserializationFailed)?;
+
+        let matches: Vec<usize> = type_2
+            .info
+            .iter()
+            .enumerate()
+            .filter_map(|(i, info)| (info.without_pubkeys.message == message.0).then_some(i))
+            .collect();
+        let index = match matches.as_slice() {
+            [i] => *i,
+            [] => return Err(AggregationError::UnknownMessage),
+            _ => return Err(AggregationError::MultipleMessages),
+        };
+
+        split_type_2(type_2, index, LOG_INV_RATE)
+            .map_err(|err| AggregationError::ProverFailure(err.to_string()))
+    })
 }
 
 #[cfg(test)]