diff --git a/Cargo.lock b/Cargo.lock index 871a392c..f3bfb5a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -174,7 +174,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -185,7 +185,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -652,7 +652,7 @@ dependencies = [ [[package]] name = "backend" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "mt-air", "mt-fiat-shamir", @@ -1940,7 +1940,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2010,6 +2010,7 @@ version = "0.1.0" dependencies = [ "clap", "ethlambda-blockchain", + "ethlambda-crypto", "ethlambda-network-api", "ethlambda-p2p", "ethlambda-rpc", @@ -2067,6 +2068,7 @@ dependencies = [ "leansig", "leansig_wrapper", "rand 0.10.1", + "rayon", "thiserror 2.0.18", ] @@ -3664,7 +3666,7 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "lean-multisig" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "clap", @@ -3682,7 +3684,7 @@ dependencies = [ [[package]] name = "lean_compiler" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "include_dir", @@ -3698,7 +3700,7 @@ dependencies = [ [[package]] name = "lean_prover" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "itertools 0.14.0", @@ -3716,7 +3718,7 @@ dependencies = [ [[package]] name = "lean_vm" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "itertools 0.14.0", @@ -3772,7 +3774,7 @@ dependencies = [ [[package]] name = "leansig_wrapper" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "ethereum_ssz", @@ -4807,7 +4809,7 @@ dependencies = [ [[package]] name = "mt-air" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "mt-field", "mt-poly", @@ -4816,7 +4818,7 @@ dependencies = [ [[package]] name = "mt-fiat-shamir" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "mt-field", "mt-koala-bear", @@ -4830,11 +4832,11 @@ dependencies = [ [[package]] name = "mt-field" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "itertools 0.14.0", "mt-utils", - "num-bigint 0.4.6", + "num-bigint 0.3.3", "paste", "rand 0.10.1", "rayon", @@ -4845,12 +4847,12 @@ dependencies = [ [[package]] name = "mt-koala-bear" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "itertools 0.14.0", "mt-field", "mt-utils", - "num-bigint 0.4.6", + "num-bigint 0.3.3", "paste", "rand 0.10.1", "rayon", @@ -4861,7 +4863,7 @@ dependencies = [ [[package]] name = "mt-poly" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "itertools 0.14.0", "mt-field", @@ -4875,7 +4877,7 @@ dependencies = [ [[package]] name = "mt-sumcheck" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "mt-air", "mt-fiat-shamir", @@ -4888,7 +4890,7 @@ dependencies = [ [[package]] name = "mt-symetric" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "mt-field", "mt-koala-bear", @@ -4898,7 +4900,7 @@ dependencies = [ [[package]] name = "mt-utils" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "serde", ] @@ -4906,7 +4908,7 @@ dependencies = [ [[package]] name = "mt-whir" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "itertools 0.14.0", "mt-fiat-shamir", @@ -5099,7 +5101,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -5302,7 +5304,7 @@ dependencies = [ [[package]] name = "p3-baby-bear" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-challenger 0.5.1", "p3-field 0.5.1", @@ -5346,7 +5348,7 @@ dependencies = [ [[package]] name = "p3-challenger" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-field 0.5.1", "p3-maybe-rayon 0.5.1", @@ -5372,14 +5374,14 @@ dependencies = [ [[package]] name = "p3-dft" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "itertools 0.14.0", "p3-field 0.5.1", "p3-matrix 0.5.1", "p3-maybe-rayon 0.5.1", "p3-util 0.5.1", - "spin 0.11.0", + "spin 0.10.0", "tracing", ] @@ -5400,7 +5402,7 @@ dependencies = [ [[package]] name = "p3-field" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "itertools 0.14.0", "num-bigint 0.4.6", @@ -5432,7 +5434,7 @@ dependencies = [ [[package]] name = "p3-koala-bear" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-challenger 0.5.1", "p3-field 0.5.1", @@ -5462,7 +5464,7 @@ dependencies = [ [[package]] name = "p3-matrix" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "itertools 0.14.0", "p3-field 0.5.1", @@ -5482,7 +5484,7 @@ checksum = "55ac1d2f102cf8c71dba1b449575c99697781fcc028831e83d2245787bd7a650" [[package]] name = "p3-maybe-rayon" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" [[package]] name = "p3-mds" @@ -5502,7 +5504,7 @@ dependencies = [ [[package]] name = "p3-mds" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-dft 0.5.1", "p3-field 0.5.1", @@ -5514,7 +5516,7 @@ dependencies = [ [[package]] name = "p3-monty-31" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "itertools 0.14.0", "num-bigint 0.4.6", @@ -5530,17 +5532,16 @@ dependencies = [ "paste", "rand 0.10.1", "serde", - "spin 0.11.0", + "spin 0.10.0", "tracing", ] [[package]] name = "p3-poseidon1" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-field 0.5.1", - "p3-mds 0.5.1", "p3-symmetric 0.5.1", "rand 0.10.1", ] @@ -5562,7 +5563,7 @@ dependencies = [ [[package]] name = "p3-poseidon2" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "p3-field 0.5.1", "p3-mds 0.5.1", @@ -5585,7 +5586,7 @@ dependencies = [ [[package]] name = "p3-symmetric" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "itertools 0.14.0", "p3-field 0.5.1", @@ -5605,7 +5606,7 @@ dependencies = [ [[package]] name = "p3-util" version = "0.5.1" -source = "git+https://github.com/Plonky3/Plonky3.git#2aeaa17557e7f54b0caa0add42e7d5b9aec4f564" +source = "git+https://github.com/Plonky3/Plonky3.git#3f67d136c71bec40f180c85d0bb2b654acddef22" dependencies = [ "serde", "transpose", @@ -6379,7 +6380,7 @@ dependencies = [ [[package]] name = "rec_aggregation" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "include_dir", @@ -6700,7 +6701,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -7212,7 +7213,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -7357,9 +7358,9 @@ dependencies = [ [[package]] name = "spin" -version = "0.11.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "783f3f6f6b01e295a669edfc402133a5f2553d1f0e81284b3ba4594e80bdd4a2" +checksum = "d5fe4ccb98d9c292d56fec89a5e07da7fc4cf0dc11e156b41793132775d3e591" dependencies = [ "lock_api", ] @@ -7428,7 +7429,7 @@ dependencies = [ [[package]] name = "sub_protocols" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "lean_vm", @@ -7522,7 +7523,7 @@ dependencies = [ [[package]] name = "system-info" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "libc", "rayon", @@ -7547,10 +7548,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -8072,7 +8073,7 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "utils" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "backend", "tracing", @@ -8371,7 +8372,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -9038,7 +9039,7 @@ dependencies = [ [[package]] name = "zk-alloc" version = "0.1.0" -source = "git+https://github.com/leanEthereum/leanMultisig.git?rev=0242c909#0242c909260c9e16893baae3004667230429808d" +source = "git+https://github.com/leanEthereum/leanVM.git?rev=8fcbd779#8fcbd77958a58666e828315de2d6ce7c93297117" dependencies = [ "libc", "system-info", diff --git a/Cargo.toml b/Cargo.toml index a6d37b86..016ccbb4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,4 +76,3 @@ eyre = "0.6" # Allocator + heap profiling tikv-jemallocator = { version = "0.6", features = ["stats", "unprefixed_malloc_on_supported_platforms", "profiling"] } jemalloc_pprof = { version = "0.8", features = ["flamegraph"] } - diff --git a/Makefile b/Makefile index 8568d137..fe68b374 100644 --- a/Makefile +++ b/Makefile @@ -24,8 +24,8 @@ docker-build: ## 🐳 Build the Docker image -t ghcr.io/lambdaclass/ethlambda:$(DOCKER_TAG) . @echo -# 2026-05-21 -LEAN_SPEC_COMMIT_HASH:=825bec6bf278920cfc56730d64a7c90522a0bb6c +# 2026-06-03 +LEAN_SPEC_COMMIT_HASH:=30ffb6cab54ca6d2e2e1c82e8e2713ebb9a8fa3f leanSpec: git clone https://github.com/leanEthereum/leanSpec.git --single-branch diff --git a/bin/ethlambda/Cargo.toml b/bin/ethlambda/Cargo.toml index e5e22ee9..e39605ef 100644 --- a/bin/ethlambda/Cargo.toml +++ b/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 diff --git a/bin/ethlambda/src/main.rs b/bin/ethlambda/src/main.rs index d361b969..58802d20 100644 --- a/bin/ethlambda/src/main.rs +++ b/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"); diff --git a/crates/common/crypto/Cargo.toml b/crates/common/crypto/Cargo.toml index dc5ba718..22e6ca69 100644 --- a/crates/common/crypto/Cargo.toml +++ b/crates/common/crypto/Cargo.toml @@ -12,13 +12,22 @@ version.workspace = true [dependencies] ethlambda-types.workspace = true -lean-multisig = { git = "https://github.com/leanEthereum/leanMultisig.git", rev = "0242c909" } +lean-multisig = { git = "https://github.com/leanEthereum/leanVM.git", rev = "8fcbd779" } # leansig_wrapper provides XmssPublicKey/XmssSignature types used by lean-multisig's public API -leansig_wrapper = { git = "https://github.com/leanEthereum/leanMultisig.git", rev = "0242c909" } +leansig_wrapper = { git = "https://github.com/leanEthereum/leanVM.git", rev = "8fcbd779" } 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 diff --git a/crates/common/crypto/src/lib.rs b/crates/common/crypto/src/lib.rs index bb37aa2b..9a4eb8e4 100644 --- a/crates/common/crypto/src/lib.rs +++ b/crates/common/crypto/src/lib.rs @@ -15,7 +15,7 @@ use thiserror::Error; /// log(1/rate) for the WHIR commitment scheme used inside lean-multisig. /// 2 matches the devnet-4 cross-client convention (zeam, ream, grandine, lantern -/// all use 2); the leanMultisig devnet5 examples also use 2 for recursion. +/// all use 2); the leanVM devnet5 examples also use 2 for recursion. const LOG_INV_RATE: usize = 2; // Lazy initialization for prover and verifier setup @@ -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(produce: F) -> Result +where + F: FnOnce() -> Result, +{ + // 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(produce: F) -> Result +where + F: FnOnce() -> Result, +{ + ensure_prover_ready(); + compress_type1_to_byte_list(&produce()?) +} + +/// Type-2 counterpart of [`prove_type1`]. +#[cfg(feature = "zk-alloc")] +fn prove_type2(produce: F) -> Result +where + F: FnOnce() -> Result, +{ + // 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(produce: F) -> Result +where + F: FnOnce() -> Result, +{ + 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 = children - .into_iter() - .enumerate() - .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) - .collect::>()?; - - 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 = children + .into_iter() + .enumerate() + .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) + .collect::>()?; + + 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 = children - .into_iter() - .enumerate() - .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) - .collect::>()?; - - let proof = aggregate_type_1(&children_native, vec![], message.0, slot, LOG_INV_RATE) - .map_err(|err| AggregationError::ProverFailure(err.to_string()))?; + prove_type1(move || { + let children_native: Vec = children + .into_iter() + .enumerate() + .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) + .collect::>()?; - 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 = type_1s - .into_iter() - .enumerate() - .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) - .collect::>()?; - - 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 = type_1s + .into_iter() + .enumerate() + .map(|(i, (pubkeys, proof_bytes))| decompress_type1(pubkeys, &proof_bytes, i)) + .collect::>()?; - 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>, message: &H256, ) -> Result { - ensure_prover_ready(); - - let pubkeys_per_info: Vec> = 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 = 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> = 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 = 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)] diff --git a/crates/common/crypto/src/zk_alloc.rs b/crates/common/crypto/src/zk_alloc.rs new file mode 100644 index 00000000..3ec3233a --- /dev/null +++ b/crates/common/crypto/src/zk_alloc.rs @@ -0,0 +1,188 @@ +//! Proving-scoped integration of leanVM's `zk-alloc` arena allocator. +//! +//! **Benchmark build only** (`zk-alloc` feature). The goal is to get leanVM's +//! bump-arena speedup for XMSS aggregation without destabilizing a long-running +//! node. +//! +//! # Why a dispatcher instead of installing `ZkAllocator` directly +//! +//! `ZkAllocator` is a *process-global* bump-arena: [`begin_phase`] flips a global +//! switch and resets every thread's slab. leanVM's own binary is safe because it +//! does nothing but prove between `begin_phase`/`end_phase`. ethlambda is not: +//! the tokio runtime, p2p, storage, and actor threads allocate continuously. Any +//! long-lived buffer one of them allocates during a phase would be silently +//! overwritten by the next phase's slab reset. +//! +//! [`ScopedZkAlloc`] is a `#[global_allocator]` that routes to the arena **only on +//! threads explicitly marked as proving**: the global rayon pool's workers, which +//! leanVM's prover uses exclusively (ethlambda itself never touches the global +//! pool). The pool is built at startup via [`init_arena_rayon_pool`] with a +//! `start_handler` that permanently flags each worker. Every other thread — +//! tokio, p2p, storage, the actor thread — always uses the system allocator, so +//! its allocations are never reset. The prover *caller* is also unflagged: only +//! the parallel work inside leanVM lands in the arena, and the assembled proof +//! itself lands in System. +//! +//! Two earlier designs failed and inform this one: +//! - Marking the global pool's workers with `rayon::broadcast` around each phase +//! deadlocked the node: broadcast left the pool's sleep/wakeup accounting in a +//! state where the prover's injected work was never stolen. +//! - A dedicated `install`-target pool crashed on the second proof: rayon's and +//! crossbeam's long-lived internals (epoch participants, work-stealing deque +//! buffers) were first allocated *inside* a phase, in arena memory the next +//! phase's slab reset corrupted (`crossbeam_epoch::try_advance` UAF, then +//! rayon `AbortIfPanic`). The global pool avoids this because `setup_prover` +//! runs heavy parallel work on it *before any phase exists*, growing those +//! internals to near-peak size in System memory — the same warmup leanVM's +//! own binaries rely on. +//! +//! Two conditions gate an arena allocation: the thread-local [`USE_ARENA`] flag +//! **and** leanVM's global `ARENA_ACTIVE` (checked inside `ZkAllocator::alloc`, +//! set by `begin_phase`). So even on a flagged thread, anything allocated outside +//! a phase — prover setup, input decompression, output serialization — lands in +//! the system allocator. +//! +//! # Accepted limitations +//! +//! - jemalloc and its `/debug/pprof` heap endpoints are gone in this build: +//! `ZkAllocator::dealloc` forwards non-arena frees to `std::alloc::System`, so +//! the non-proving path must be `System` (libc), not jemalloc, or we would free +//! a jemalloc pointer with libc `free`. +//! - All proving is serialized behind [`PROVING`]; concurrent aggregation (the +//! spawn_blocking worker vs. actor-thread block building) blocks rather than +//! running in parallel. Acceptable for a benchmark; it also prevents the +//! `begin_phase` nesting panic. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::sync::{Mutex, MutexGuard, Once}; + +use lean_multisig::{ZkAllocator, begin_phase, end_phase}; + +/// Re-exported so the binary can run leanVM's startup core-count assertion. +pub use lean_multisig::init_allocator; + +thread_local! { + /// Marks the current thread as a proving thread. While set *and* a phase is + /// active, this thread's allocations route to leanVM's arena; otherwise to + /// the system allocator. + static USE_ARENA: Cell = const { Cell::new(false) }; +} + +/// Serializes every entry into the leanVM prover so phases never nest and +/// `ARENA_ACTIVE` is only ever true for one thread group at a time. +pub static PROVING: Mutex<()> = Mutex::new(()); + +/// Global allocator that confines leanVM's arena to proving threads. +pub struct ScopedZkAlloc; + +// SAFETY: every returned pointer comes from either `ZkAllocator` (which yields +// arena memory only while a phase is active, System otherwise) or `System` +// directly. `dealloc` always defers to `ZkAllocator::dealloc`, which is +// address-based — arena pointers are a no-op, all others forward to System — +// so it correctly frees pointers produced by either path, from any thread. +unsafe impl GlobalAlloc for ScopedZkAlloc { + #[inline] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if USE_ARENA.get() { + unsafe { ZkAllocator.alloc(layout) } + } else { + unsafe { System.alloc(layout) } + } + } + + #[inline] + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + if USE_ARENA.get() { + unsafe { ZkAllocator.alloc_zeroed(layout) } + } else { + unsafe { System.alloc_zeroed(layout) } + } + } + + #[inline] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { ZkAllocator.dealloc(ptr, layout) }; + } + + #[inline] + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // Must NOT defer to `ZkAllocator::realloc`: on growth it calls its own + // `alloc`, which gates on the global `ARENA_ACTIVE` rather than our + // thread-local flag, which would leak arena memory onto non-proving + // threads. Route through our own thread-gated alloc/dealloc instead. + if new_size <= layout.size() { + // Shrink in place; the existing block is large enough. Matches + // `ZkAllocator`'s behaviour and is valid for arena and System + // pointers alike. + return ptr; + } + let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) }; + let new_ptr = unsafe { self.alloc(new_layout) }; + if !new_ptr.is_null() { + unsafe { std::ptr::copy_nonoverlapping(ptr, new_ptr, layout.size()) }; + unsafe { self.dealloc(ptr, layout) }; + } + new_ptr + } +} + +/// Build the **global** rayon pool with arena-flagged workers. Must run before +/// anything else initializes the global pool (leanVM's `setup_prover` is the +/// first rayon user in ethlambda). Idempotent. If the pool was already built by +/// someone else, its workers stay unflagged and proving simply falls back to the +/// system allocator — safe, just no arena speedup. +pub fn init_arena_rayon_pool() { + static INIT: Once = Once::new(); + INIT.call_once(|| { + rayon::ThreadPoolBuilder::new() + .thread_name(|i| format!("zk-prover-{i}")) + .start_handler(|_| USE_ARENA.set(true)) + .build_global() + .inspect_err(|err| { + eprintln!( + "zk-alloc: global rayon pool already initialized ({err}); \ + proving will not use the arena" + ); + }) + .ok(); + }); +} + +/// Holds the [`PROVING`] lock for one prover operation, serializing all proving so +/// phases never nest. The lock is released when the session is dropped, which must +/// be *after* the proof is serialized (the proof may reference arena memory that +/// the next phase would reset). +pub(crate) struct ArenaSession { + _lock: MutexGuard<'static, ()>, +} + +impl ArenaSession { + pub(crate) fn begin() -> Self { + let lock = PROVING.lock().unwrap_or_else(|poison| poison.into_inner()); + Self { _lock: lock } + } + + /// Run `produce` (the prover call) inside an arena phase. `produce` executes + /// on the calling thread (unflagged → System); the prover's internal rayon + /// work runs on the arena-flagged global pool workers. The returned value may + /// reference arena memory; it is safe to read until the next `begin_phase`, + /// which the held lock prevents. + pub(crate) fn prove(&self, produce: F) -> T + where + F: FnOnce() -> T, + { + begin_phase(); + // Guarantees `end_phase` runs even if the prover panics, so the global + // arena switch is never left stuck active. leanVM's `end_phase` also + // flushes the global pool's injector (its job blocks may live in arena). + struct EndOnDrop; + impl Drop for EndOnDrop { + fn drop(&mut self) { + end_phase(); + } + } + let _end = EndOnDrop; + produce() + } +}