diff --git a/.github/workflows/axvm-sdk.yml b/.github/workflows/axvm-sdk.yml index 6972ce52de..c57de86288 100644 --- a/.github/workflows/axvm-sdk.yml +++ b/.github/workflows/axvm-sdk.yml @@ -2,9 +2,9 @@ name: axVM SDK Tests on: push: - branches: [ "main" ] + branches: ["main"] pull_request: - branches: [ "**" ] + branches: ["**"] paths: - "crates/stark-backend/**" - "crates/circuits/primitives/**" @@ -23,7 +23,8 @@ jobs: tests: runs-on: - runs-on=${{ github.run_id }} - - runner=64cpu-linux-arm64 + - family=m7a.48xlarge + - disk=large steps: - uses: actions/checkout@v4 @@ -41,9 +42,11 @@ jobs: case $arch in arm64|aarch64) S5CMD_BIN="s5cmd_2.2.2_linux_arm64.deb" + rustup component add rust-src --toolchain nightly-2024-10-30-aarch64-unknown-linux-gnu ;; x86_64|amd64) S5CMD_BIN="s5cmd_2.2.2_linux_amd64.deb" + rustup component add rust-src --toolchain nightly-2024-10-30-x86_64-unknown-linux-gnu ;; *) echo "Unsupported architecture: $arch" @@ -71,6 +74,5 @@ jobs: - name: Run axvm-sdk crate tests working-directory: crates/axvm-sdk run: | - rustup component add rust-src --toolchain nightly-2024-10-30-aarch64-unknown-linux-gnu export RUST_BACKTRACE=1 cargo nextest run --cargo-profile=fast --test-threads=2 --features parallel diff --git a/Cargo.lock b/Cargo.lock index 9c55e6600e..4af4a16f43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,9 +76,9 @@ dependencies = [ [[package]] name = "alloy-primitives" -version = "0.8.14" +version = "0.8.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9db948902dfbae96a73c2fbf1f7abec62af034ab883e4c777c3fd29702bd6e2c" +checksum = "6259a506ab13e1d658796c31e6e39d2e2ee89243bcc505ddc613b35732e0a430" dependencies = [ "alloy-rlp", "bytes", @@ -104,9 +104,9 @@ dependencies = [ [[package]] name = "alloy-rlp" -version = "0.3.9" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0822426598f95e45dd1ea32a738dac057529a709ee645fcc516ffa4cbde08f" +checksum = "f542548a609dca89fcd72b3b9f355928cf844d4363c5eed9c5273a3dd225e097" dependencies = [ "alloy-rlp-derive", "arrayvec", @@ -115,9 +115,9 @@ dependencies = [ [[package]] name = "alloy-rlp-derive" -version = "0.3.9" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b09cae092c27b6f1bde952653a22708691802e57bfef4a2973b80bea21efd3f" +checksum = "5a833d97bf8a5f0f878daf2c8451fff7de7f9de38baa5a45d936ec718d81255a" dependencies = [ "proc-macro2", "quote", @@ -475,7 +475,7 @@ dependencies = [ "rayon", "serde", "serde_json", - "thiserror", + "thiserror 1.0.69", "tikv-jemallocator", "tracing", "tracing-forest", @@ -549,12 +549,11 @@ dependencies = [ "derive-new", "derive_more 1.0.0", "halo2curves-axiom 0.7.0", - "hex-literal", "itertools 0.13.0", "num-bigint-dig", "num-traits", - "once_cell", "rand", + "serde", "strum", "tracing", ] @@ -639,6 +638,7 @@ dependencies = [ "num-bigint-dig", "pprof", "rand_chacha", + "serde", "tempfile", "tiny-keccak", "tokio", @@ -662,15 +662,12 @@ dependencies = [ "axvm-rv32im-transpiler", "derive-new", "derive_more 1.0.0", - "eyre", "hex", - "hex-literal", - "itertools 0.13.0", "lazy_static", "rand", + "serde", "test-case", "test-log", - "tracing", ] [[package]] @@ -764,7 +761,7 @@ dependencies = [ "static_assertions", "test-case", "test-log", - "thiserror", + "thiserror 1.0.69", "toml", "tracing", "zkhash", @@ -805,6 +802,7 @@ dependencies = [ "num-integer", "num-traits", "once_cell", + "serde", "strum", ] @@ -907,6 +905,7 @@ dependencies = [ "lazy_static", "p3-keccak-air", "rand", + "serde", "strum", "test-case", "test-log", @@ -1032,6 +1031,7 @@ dependencies = [ "axvm-native-circuit", "axvm-native-compiler", "axvm-native-recursion", + "bson", "cfg-if", "itertools 0.13.0", "lazy_static", @@ -1077,6 +1077,7 @@ dependencies = [ "num-traits", "once_cell", "rand", + "serde", "strum", ] @@ -1174,6 +1175,7 @@ dependencies = [ "num-integer", "parking_lot", "rand", + "serde", "strum", "test-case", "test-log", @@ -1230,6 +1232,7 @@ dependencies = [ "axvm-transpiler", "bincode 1.3.3", "bon", + "bson", "derivative", "derive_more 1.0.0", "eyre", @@ -1303,7 +1306,7 @@ dependencies = [ "strum", "tempfile", "test-case", - "thiserror", + "thiserror 1.0.69", "tracing", ] @@ -1477,9 +1480,9 @@ dependencies = [ [[package]] name = "bon" -version = "3.2.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9276fe602371cd8a7f70fe68c4db55b2d3e92c570627d6ed0427646edfa5cf47" +checksum = "f265cdb2e8501f1c952749e78babe8f1937be92c98120e5f78fc72d634682bad" dependencies = [ "bon-macros", "rustversion", @@ -1487,9 +1490,9 @@ dependencies = [ [[package]] name = "bon-macros" -version = "3.2.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94828b84b32b4f3ac3865f692fcdbc46c7d0dd87b29658a391d58a244e1ce45a" +checksum = "38aa5c627cd7706490e5b003d685f8b9d69bc343b1a00b9fdd01e75fdf6827cf" dependencies = [ "darling", "ident_case", @@ -1500,6 +1503,27 @@ dependencies = [ "syn 2.0.90", ] +[[package]] +name = "bson" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "068208f2b6fcfa27a7f1ee37488d2bb8ba2640f68f5475d08e1d9130696aba59" +dependencies = [ + "ahash", + "base64", + "bitvec", + "hex", + "indexmap", + "js-sys", + "once_cell", + "rand", + "serde", + "serde_bytes", + "serde_json", + "time", + "uuid", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -1628,7 +1652,7 @@ dependencies = [ "semver 1.0.23", "serde", "serde_json", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1639,9 +1663,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.2" +version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f34d93e62b03caf570cccc334cbc6c2fceca82f39211051345108adcba3eebdc" +checksum = "27f657647bcff5394bf56c7317665bbf790a137a50eaaa5c6bfbb9e27a518f2d" dependencies = [ "jobserver", "libc", @@ -1662,9 +1686,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chrono" -version = "0.4.38" +version = "0.4.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825" dependencies = [ "num-traits", "serde", @@ -1699,9 +1723,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.22" +version = "4.5.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69371e34337c4c984bbe322360c2547210bf632eb2814bbe78a6e87a2935bd2b" +checksum = "3135e7ec2ef7b10c6ed8950f0f792ed96ee093fa088608f1c76e569722700c84" dependencies = [ "clap_builder", "clap_derive", @@ -1709,9 +1733,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.22" +version = "4.5.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e24c1b4099818523236a8ca881d2b45db98dadfb4625cf6608c12069fcbbde1" +checksum = "30582fc632330df2bd26877bde0c1f4470d57c582bbc070376afcd04d8cb4838" dependencies = [ "anstream", "anstyle", @@ -1733,9 +1757,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.3" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "afb84c814227b90d6895e01398aee0d8033c00e7466aca416fb6a8e0eb19d8a7" +checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6" [[package]] name = "colorchoice" @@ -2336,9 +2360,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "486f806e73c5707928240ddc295403b1b93c96a02038563881c4a2fd84b81ac4" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "fastrlp" @@ -3057,9 +3081,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.74" +version = "0.3.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a865e038f7f6ed956f788f0d7d60c541fff74c7bd74272c5d4cf15c63743e705" +checksum = "6717b6b5b077764fb5966237269cb3c64edddde4b14ce42647430a78ced9e7b7" dependencies = [ "once_cell", "wasm-bindgen", @@ -3122,9 +3146,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.167" +version = "0.2.168" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09d6582e104315a817dff97f75133544b2e094ee22447d2acf4a74e189ba06fc" +checksum = "5aaeb2981e0606ca11d79718f8bb01164f1d6ed75080182d3abf017e6d244b6d" [[package]] name = "libgit2-sys" @@ -3973,12 +3997,12 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pest" -version = "2.7.14" +version = "2.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879952a81a83930934cbf1786752d6dedc3b1f29e8f8fb2ad1d0a36f377cf442" +checksum = "8b7cafe60d6cf8e62e1b9b2ea516a089c008945bb5a275416789e7db0bc199dc" dependencies = [ "memchr", - "thiserror", + "thiserror 2.0.6", "ucd-trie", ] @@ -4056,7 +4080,7 @@ dependencies = [ "log", "rand", "rand_xorshift", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -4084,7 +4108,7 @@ dependencies = [ "smallvec", "symbolic-demangle", "tempfile", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -4340,7 +4364,7 @@ checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" dependencies = [ "getrandom", "libredox", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -4585,15 +4609,15 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.41" +version = "0.38.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7f649912bc1495e167a6edee79151c84b1bad49748cb4f1f1167f459f6224f6" +checksum = "f93dc38ecbab2eb790ff964bb77fa94faf256fd3e73285fd7ba0903b76bedb85" dependencies = [ "bitflags 2.6.0", "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -4742,6 +4766,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_bytes" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "387cc504cb06bb40a96c8e04e951fe01854cf6bc921053c954e4a606d9675c6a" +dependencies = [ + "serde", +] + [[package]] name = "serde_derive" version = "1.0.215" @@ -5167,7 +5200,16 @@ version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec2a1820ebd077e2b90c4df007bebf344cd394098a13c563957d0afc83ea47" +dependencies = [ + "thiserror-impl 2.0.6", ] [[package]] @@ -5181,6 +5223,17 @@ dependencies = [ "syn 2.0.90", ] +[[package]] +name = "thiserror-impl" +version = "2.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d65750cab40f4ff1929fb1ba509e9914eb756131cef4210da8d5d700d26f6312" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.90", +] + [[package]] name = "thread_local" version = "1.1.8" @@ -5378,7 +5431,7 @@ checksum = "ee40835db14ddd1e3ba414292272eddde9dad04d3d4b65509656414d1c42592f" dependencies = [ "ansi_term", "smallvec", - "thiserror", + "thiserror 1.0.69", "tracing", "tracing-subscriber", ] @@ -5520,6 +5573,10 @@ name = "uuid" version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" +dependencies = [ + "getrandom", + "serde", +] [[package]] name = "valuable" @@ -5585,9 +5642,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasm-bindgen" -version = "0.2.97" +version = "0.2.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d15e63b4482863c109d70a7b8706c1e364eb6ea449b201a76c5b89cedcec2d5c" +checksum = "a474f6281d1d70c17ae7aa6a613c87fce69a127e2624002df63dcb39d6cf6396" dependencies = [ "cfg-if", "once_cell", @@ -5596,13 +5653,12 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.97" +version = "0.2.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d36ef12e3aaca16ddd3f67922bc63e48e953f126de60bd33ccc0101ef9998cd" +checksum = "5f89bb38646b4f81674e8f5c3fb81b562be1fd936d84320f3264486418519c79" dependencies = [ "bumpalo", "log", - "once_cell", "proc-macro2", "quote", "syn 2.0.90", @@ -5611,9 +5667,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.97" +version = "0.2.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "705440e08b42d3e4b36de7d66c944be628d579796b8090bfa3471478a2260051" +checksum = "2cc6181fd9a7492eef6fef1f33961e3695e4579b9872a6f7c83aee556666d4fe" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -5621,9 +5677,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.97" +version = "0.2.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98c9ae5a76e46f4deecd0f0255cc223cfa18dc9b261213b8aa0c7b36f61b3f1d" +checksum = "30d7a95b763d3c45903ed6c81f156801839e5ee968bb07e534c44df0fcd330c2" dependencies = [ "proc-macro2", "quote", @@ -5634,15 +5690,15 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.97" +version = "0.2.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ee99da9c5ba11bd675621338ef6fa52296b76b83305e9b6e5c77d4c286d6d49" +checksum = "943aab3fdaaa029a6e0271b35ea10b72b943135afe9bffca82384098ad0e06a6" [[package]] name = "web-sys" -version = "0.3.74" +version = "0.3.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a98bc3c33f0fe7e59ad7cd041b89034fa82a7c2d4365ca538dda6cdaf513863c" +checksum = "04dd7223427d52553d3702c004d3b2fe07c148165faa56313cb00211e31c12bc" dependencies = [ "js-sys", "wasm-bindgen", diff --git a/Cargo.toml b/Cargo.toml index d1c7ba0ad3..05cc3bcdfa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -207,6 +207,7 @@ static_assertions = "1.1.0" async-trait = "0.1.83" getset = "0.1.3" rrs-lib = "0.1.0" +bson = "2.13.0" # cryptography, default-features = false for no_std tiny-keccak = { version = "2.0.2", features = ["keccak"] } diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 3c8c47e919..1e4aae2aae 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -40,7 +40,8 @@ k256 = { workspace = true, features = ["ecdsa"] } tiny-keccak.workspace = true derive-new.workspace = true derive_more = { workspace = true, features = ["from"] } -num-bigint-dig = { workspace = true, features = ["std"] } +num-bigint-dig = { workspace = true, features = ["std", "serde"] } +serde.workspace = true [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } diff --git a/benchmarks/src/bin/base64_json.rs b/benchmarks/src/bin/base64_json.rs index 9e8f5e35f3..493a31ef27 100644 --- a/benchmarks/src/bin/base64_json.rs +++ b/benchmarks/src/bin/base64_json.rs @@ -40,22 +40,21 @@ fn main() -> Result<()> { )?; run_with_metric_collection("OUTPUT_PATH", || -> Result<()> { - let vdata = - info_span!("Base64 Json Program", group = "base64_json_program").in_scope(|| { - let engine = BabyBearPoseidon2Engine::new( - FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), - ); + let vdata = info_span!("Base64 Json Program").in_scope(|| { + let engine = BabyBearPoseidon2Engine::new( + FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), + ); - let data = include_str!("../../programs/base64_json/json_payload_encoded.txt"); + let data = include_str!("../../programs/base64_json/json_payload_encoded.txt"); - let fe_bytes = data.to_owned().into_bytes(); - bench_from_exe( - engine, - Keccak256Rv32Config::default(), - exe, - StdIn::from_bytes(&fe_bytes), - ) - })?; + let fe_bytes = data.to_owned().into_bytes(); + bench_from_exe( + engine, + Keccak256Rv32Config::default(), + exe, + StdIn::from_bytes(&fe_bytes), + ) + })?; #[cfg(feature = "aggregation")] { diff --git a/benchmarks/src/bin/ecrecover.rs b/benchmarks/src/bin/ecrecover.rs index 7169991063..9cd8fc4d06 100644 --- a/benchmarks/src/bin/ecrecover.rs +++ b/benchmarks/src/bin/ecrecover.rs @@ -47,6 +47,7 @@ use eyre::Result; use k256::ecdsa::{RecoveryId, Signature, SigningKey, VerifyingKey}; use num_bigint_dig::BigUint; use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng}; +use serde::{Deserialize, Serialize}; use tiny_keccak::{Hasher, Keccak}; use tracing::info_span; @@ -66,7 +67,7 @@ fn make_input(signing_key: &SigningKey, msg: &[u8]) -> Vec { input.into_iter().map(BabyBear::from_canonical_u8).collect() } -#[derive(Clone, Debug, VmConfig, derive_new::new)] +#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)] pub struct Rv32ImEcRecoverConfig { #[system] pub system: SystemConfig, @@ -122,37 +123,36 @@ fn main() -> Result<()> { let vm_config = Rv32ImEcRecoverConfig::for_curves(vec![SECP256K1_CONFIG.clone()]); run_with_metric_collection("OUTPUT_PATH", || -> Result<()> { - let vdata = - info_span!("ECDSA Recover Program", group = "ecrecover_program").in_scope(|| { - let mut rng = ChaCha8Rng::seed_from_u64(12345); - let signing_key: SigningKey = SigningKey::random(&mut rng); - let verifying_key = VerifyingKey::from(&signing_key); - let mut hasher = Keccak::v256(); - let mut expected_address = [0u8; 32]; - hasher.update( - &verifying_key - .to_encoded_point(/* compress = */ false) - .as_bytes()[1..], - ); - hasher.finalize(&mut expected_address); - expected_address[..12].fill(0); // 20 bytes as the address. - let mut input_stream = vec![expected_address - .into_iter() - .map(BabyBear::from_canonical_u8) - .collect::>()]; + let vdata = info_span!("ECDSA Recover Program").in_scope(|| { + let mut rng = ChaCha8Rng::seed_from_u64(12345); + let signing_key: SigningKey = SigningKey::random(&mut rng); + let verifying_key = VerifyingKey::from(&signing_key); + let mut hasher = Keccak::v256(); + let mut expected_address = [0u8; 32]; + hasher.update( + &verifying_key + .to_encoded_point(/* compress = */ false) + .as_bytes()[1..], + ); + hasher.finalize(&mut expected_address); + expected_address[..12].fill(0); // 20 bytes as the address. + let mut input_stream = vec![expected_address + .into_iter() + .map(BabyBear::from_canonical_u8) + .collect::>()]; - let msg = ["Elliptic", "Curve", "Digital", "Signature", "Algorithm"]; - input_stream.extend( - msg.iter() - .map(|s| make_input(&signing_key, s.as_bytes())) - .collect::>(), - ); + let msg = ["Elliptic", "Curve", "Digital", "Signature", "Algorithm"]; + input_stream.extend( + msg.iter() + .map(|s| make_input(&signing_key, s.as_bytes())) + .collect::>(), + ); - let engine = BabyBearPoseidon2Engine::new( - FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), - ); - bench_from_exe(engine, vm_config, exe, input_stream.into()) - })?; + let engine = BabyBearPoseidon2Engine::new( + FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), + ); + bench_from_exe(engine, vm_config, exe, input_stream.into()) + })?; Ok(()) }) diff --git a/benchmarks/src/bin/fib_e2e.rs b/benchmarks/src/bin/fib_e2e.rs index dd9d4e809e..8740cc8498 100644 --- a/benchmarks/src/bin/fib_e2e.rs +++ b/benchmarks/src/bin/fib_e2e.rs @@ -5,60 +5,70 @@ use ax_stark_sdk::{ use axvm_benchmarks::utils::{build_bench_program, BenchmarkCli}; use axvm_circuit::arch::instructions::{exe::AxVmExe, program::DEFAULT_MAX_NUM_PUBLIC_VALUES}; use axvm_native_compiler::conversion::CompilerOptions; -#[cfg(feature = "static-verifier")] -use axvm_native_compiler::prelude::Witness; -#[cfg(feature = "static-verifier")] -use axvm_native_recursion::witness::Witnessable; use axvm_rv32im_circuit::Rv32ImConfig; use axvm_rv32im_transpiler::{ Rv32ITranspilerExtension, Rv32IoTranspilerExtension, Rv32MTranspilerExtension, }; use axvm_sdk::{ - config::{AggConfig, AppConfig}, - keygen::{AggProvingKey, AppProvingKey}, - prover::{commit_app_exe, generate_leaf_committed_exe, StarkProver}, - StdIn, + commit::commit_app_exe, + config::{AggConfig, AppConfig, FullAggConfig, Halo2Config}, + Sdk, StdIn, }; -use axvm_transpiler::{axvm_platform::bincode, transpiler::Transpiler, FromElf}; +use axvm_transpiler::{transpiler::Transpiler, FromElf}; use clap::Parser; use eyre::Result; -#[cfg(feature = "static-verifier")] -use tracing::info_span; const NUM_PUBLIC_VALUES: usize = DEFAULT_MAX_NUM_PUBLIC_VALUES; #[tokio::main] async fn main() -> Result<()> { let cli_args = BenchmarkCli::parse(); - let app_log_blowup = cli_args.app_log_blowup.unwrap_or(2); - let agg_log_blowup = cli_args.agg_log_blowup.unwrap_or(2); - let root_log_blowup = cli_args.root_log_blowup.unwrap_or(2); - let internal_log_blowup = cli_args.internal_log_blowup.unwrap_or(2); + let app_fri_params = standard_fri_params_with_100_bits_conjectured_security( + cli_args.app_log_blowup.unwrap_or(2), + ); + let leaf_fri_params = standard_fri_params_with_100_bits_conjectured_security( + cli_args.agg_log_blowup.unwrap_or(2), + ); + let internal_fri_params = standard_fri_params_with_100_bits_conjectured_security( + cli_args.internal_log_blowup.unwrap_or(2), + ); + let root_fri_params = standard_fri_params_with_100_bits_conjectured_security( + cli_args.root_log_blowup.unwrap_or(2), + ); + let compiler_options = CompilerOptions { + // For metric collection + enable_cycle_tracker: true, + ..Default::default() + }; + // Must be larger than RangeTupleCheckerAir.height == 524288 let max_segment_length = cli_args.max_segment_length.unwrap_or(1_000_000); let app_config = AppConfig { - app_fri_params: standard_fri_params_with_100_bits_conjectured_security(app_log_blowup), + app_fri_params, app_vm_config: Rv32ImConfig::with_public_values_and_segment_len( NUM_PUBLIC_VALUES, max_segment_length, ), + leaf_fri_params, + compiler_options, }; - let agg_config = AggConfig { - max_num_user_public_values: NUM_PUBLIC_VALUES, - leaf_fri_params: standard_fri_params_with_100_bits_conjectured_security(agg_log_blowup), - internal_fri_params: standard_fri_params_with_100_bits_conjectured_security( - internal_log_blowup, - ), - root_fri_params: standard_fri_params_with_100_bits_conjectured_security(root_log_blowup), - compiler_options: CompilerOptions { - enable_cycle_tracker: true, - ..Default::default() + let full_agg_config = FullAggConfig { + agg_config: AggConfig { + max_num_user_public_values: NUM_PUBLIC_VALUES, + leaf_fri_params, + internal_fri_params, + root_fri_params, + compiler_options, + }, + halo2_config: Halo2Config { + verifier_k: 24, + wrapper_k: None, }, }; - let app_pk = AppProvingKey::keygen(app_config.clone()); - let agg_pk = AggProvingKey::keygen(agg_config); + let app_pk = Sdk.app_keygen(app_config, None::<&str>)?; + let full_agg_pk = Sdk.agg_keygen(full_agg_config, None::<&str>)?; let elf = build_bench_program("fibonacci")?; let exe = AxVmExe::from_elf( elf, @@ -67,32 +77,15 @@ async fn main() -> Result<()> { .with_extension(Rv32MTranspilerExtension) .with_extension(Rv32IoTranspilerExtension), )?; - let app_committed_exe = commit_app_exe(&app_config, exe); - let leaf_committed_exe = generate_leaf_committed_exe(&agg_config, &app_pk); + let app_committed_exe = commit_app_exe(app_fri_params, exe); - let prover = StarkProver::new(app_pk, app_committed_exe) - .with_agg_pk_and_leaf_committed_exe(agg_pk, leaf_committed_exe); + let e2e_prover = Sdk.create_e2e_prover(app_pk, app_committed_exe, full_agg_pk)?; let n = 800_000u64; - let app_input: Vec<_> = bincode::serde::encode_to_vec(n, bincode::config::standard())?; + let mut stdin = StdIn::default(); + stdin.write(&n); run_with_metric_collection("OUTPUT_PATH", || { - #[allow(unused_variables)] - let root_proof = prover.generate_e2e_proof_with_metric_spans( - StdIn::from_bytes(&app_input), - "Fibonacci Continuation Program", - ); - #[cfg(feature = "static-verifier")] - #[allow(unused_variables)] - let static_verifier_snark = info_span!("static verifier", group = "static_verifier") - .in_scope(|| { - let static_verifier = prover - .agg_pk() - .root_verifier_pk - .keygen_static_verifier(23, root_proof.clone()); - let mut witness = Witness::default(); - root_proof.write(&mut witness); - static_verifier.prove(witness) - }); + e2e_prover.generate_proof_for_evm(stdin); }); Ok(()) diff --git a/benchmarks/src/bin/fibonacci.rs b/benchmarks/src/bin/fibonacci.rs index bdbe270835..638c50d278 100644 --- a/benchmarks/src/bin/fibonacci.rs +++ b/benchmarks/src/bin/fibonacci.rs @@ -37,20 +37,19 @@ fn main() -> Result<()> { .with_extension(Rv32IoTranspilerExtension), )?; run_with_metric_collection("OUTPUT_PATH", || -> Result<()> { - let vdata = - info_span!("Fibonacci Program", group = "fibonacci_program").in_scope(|| { - let engine = BabyBearPoseidon2Engine::new( - FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), - ); - let n = 100_000u64; - let input = bincode::serde::encode_to_vec(n, bincode::config::standard())?; - bench_from_exe( - engine, - Rv32ImConfig::default(), - exe, - StdIn::from_bytes(&input), - ) - })?; + let vdata = info_span!("Fibonacci Program").in_scope(|| { + let engine = BabyBearPoseidon2Engine::new( + FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), + ); + let n = 100_000u64; + let input = bincode::serde::encode_to_vec(n, bincode::config::standard())?; + bench_from_exe( + engine, + Rv32ImConfig::default(), + exe, + StdIn::from_bytes(&input), + ) + })?; #[cfg(feature = "aggregation")] { diff --git a/benchmarks/src/bin/regex.rs b/benchmarks/src/bin/regex.rs index b65a1cd62d..151ce2a020 100644 --- a/benchmarks/src/bin/regex.rs +++ b/benchmarks/src/bin/regex.rs @@ -39,7 +39,7 @@ fn main() -> Result<()> { .with_extension(Keccak256TranspilerExtension), )?; run_with_metric_collection("OUTPUT_PATH", || -> Result<()> { - let vdata = info_span!("Regex Program", group = "regex_program").in_scope(|| { + let vdata = info_span!("Regex Program").in_scope(|| { let engine = BabyBearPoseidon2Engine::new( FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), ); diff --git a/benchmarks/src/bin/revm_transfer.rs b/benchmarks/src/bin/revm_transfer.rs index 27c4faaa0d..8f4a86edd6 100644 --- a/benchmarks/src/bin/revm_transfer.rs +++ b/benchmarks/src/bin/revm_transfer.rs @@ -37,18 +37,17 @@ fn main() -> Result<()> { .with_extension(Rv32IoTranspilerExtension), )?; run_with_metric_collection("OUTPUT_PATH", || -> Result<()> { - let vdata = - info_span!("revm 100 transfers", group = "revm_100_transfers").in_scope(|| { - let engine = BabyBearPoseidon2Engine::new( - FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), - ); - bench_from_exe( - engine, - Keccak256Rv32Config::default(), - exe, - StdIn::default(), - ) - })?; + let vdata = info_span!("revm 100 transfers").in_scope(|| { + let engine = BabyBearPoseidon2Engine::new( + FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup), + ); + bench_from_exe( + engine, + Keccak256Rv32Config::default(), + exe, + StdIn::default(), + ) + })?; Ok(()) }) } diff --git a/benchmarks/src/bin/verify_fibair.rs b/benchmarks/src/bin/verify_fibair.rs index 5f72ee28ef..0fbdc2d2a2 100644 --- a/benchmarks/src/bin/verify_fibair.rs +++ b/benchmarks/src/bin/verify_fibair.rs @@ -42,7 +42,7 @@ fn main() -> Result<()> { enable_cycle_tracker: true, ..Default::default() }; - info_span!("Verify Fibonacci AIR", group = "verify_fibair",).in_scope(|| { + info_span!("Verify Fibonacci AIR").in_scope(|| { let (program, input_stream) = build_verification_program(vdata, compiler_options); let engine = BabyBearPoseidon2Engine::new( FriParameters::standard_with_100_bits_conjectured_security(agg_log_blowup), diff --git a/benchmarks/src/utils.rs b/benchmarks/src/utils.rs index 20e138d79a..c766ae98e1 100644 --- a/benchmarks/src/utils.rs +++ b/benchmarks/src/utils.rs @@ -9,10 +9,7 @@ use ax_stark_sdk::{ use axvm_build::{build_guest_package, get_package, guest_methods, GuestOptions}; use axvm_circuit::arch::{instructions::exe::AxVmExe, VirtualMachine, VmConfig}; use axvm_sdk::{ - config::AppConfig, - keygen::AppProvingKey, - prover::{commit_app_exe, StarkProver}, - StdIn, + commit::commit_app_exe, config::AppConfig, keygen::AppProvingKey, prover::AppProver, StdIn, }; use axvm_transpiler::{axvm_platform::memory::MEM_SIZE, elf::Elf}; use clap::{command, Parser}; @@ -92,6 +89,9 @@ where let app_config = AppConfig { app_vm_config: config.clone(), app_fri_params: engine.fri_params(), + // leaf_fri_params/compiler_options don't matter for this benchmark. + leaf_fri_params: engine.fri_params(), + compiler_options: Default::default(), }; let vm = VirtualMachine::new(engine, config); // 1. Generate proving key from config. @@ -100,7 +100,7 @@ where }); // 2. Commit to the exe by generating cached trace for program. let committed_exe = time(gauge!("commit_exe_time_ms"), || { - commit_app_exe(&app_config, exe) + commit_app_exe(app_config.app_fri_params, exe) }); // 3. Executes runtime again without metric collection and generate trace. time(gauge!("execute_and_trace_gen_time_ms"), || { @@ -109,10 +109,11 @@ where // 4. Executes runtime once with full metric collection for flamegraphs (slow). // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with timer. // generate_app_proof will emit metrics for proof time of each - let prover = StarkProver::new(app_pk, committed_exe); + let vk = app_pk.app_vm_pk.vm_pk.get_vk(); + let mut prover = AppProver::new(app_pk.app_vm_pk, committed_exe); + prover.profile = true; let proofs = prover.generate_app_proof(input_stream).per_segment; // 6. Verify STARK proofs. - let vk = prover.app_pk.app_vm_pk.vm_pk.get_vk(); vm.verify(&vk, proofs.clone()).expect("Verification failed"); let vdata = proofs .into_iter() diff --git a/crates/axvm-sdk/Cargo.toml b/crates/axvm-sdk/Cargo.toml index 377bbdbf7d..67cbb7c7a1 100644 --- a/crates/axvm-sdk/Cargo.toml +++ b/crates/axvm-sdk/Cargo.toml @@ -30,6 +30,7 @@ axvm-circuit = { workspace = true } bincode = { workspace = true } bon = { workspace = true } +bson = { workspace = true } derivative = { workspace = true } derive_more = { workspace = true } serde = { workspace = true } diff --git a/crates/axvm-sdk/src/commit.rs b/crates/axvm-sdk/src/commit.rs index 790d218dc8..61fab2304a 100644 --- a/crates/axvm-sdk/src/commit.rs +++ b/crates/axvm-sdk/src/commit.rs @@ -1,10 +1,17 @@ +use std::sync::Arc; + +use ax_stark_backend::{config::StarkGenericConfig, p3_field::PrimeField32}; use ax_stark_sdk::{ - ax_stark_backend::{config::Val, p3_field::AbstractField}, - config::baby_bear_poseidon2::BabyBearPoseidon2Config, + ax_stark_backend::p3_field::AbstractField, + config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters}, + engine::StarkFriEngine, + p3_baby_bear::BabyBear, + p3_bn254_fr::Bn254Fr, }; use axvm_circuit::{ arch::{ hasher::{poseidon2::vm_poseidon2_hasher, Hasher}, + instructions::exe::AxVmExe, VmConfig, }, system::{ @@ -12,9 +19,11 @@ use axvm_circuit::{ program::trace::AxVmCommittedExe, }, }; -use axvm_native_compiler::ir::DIGEST_SIZE; +use axvm_native_compiler::{conversion::CompilerOptions, ir::DIGEST_SIZE}; -type SC = BabyBearPoseidon2Config; +use crate::{ + keygen::AppProvingKey, verifier::leaf::LeafVmVerifierConfig, NonRootCommittedExe, F, SC, +}; /// `AppExecutionCommit` has all the commitments users should check against the final proof. pub struct AppExecutionCommit { @@ -33,22 +42,22 @@ pub struct AppExecutionCommit { pub exe_commit: [T; DIGEST_SIZE], } -impl AppExecutionCommit> { +impl AppExecutionCommit { /// Users should use this function to compute `AppExecutionCommit` and check it against the final /// proof. - pub fn compute>>( + pub fn compute>( app_vm_config: &VC, - app_exe: &AxVmCommittedExe, - leaf_vm_verifier_exe: &AxVmCommittedExe, + app_exe: &NonRootCommittedExe, + leaf_vm_verifier_exe: &NonRootCommittedExe, ) -> Self { assert!( app_exe.exe.program.max_num_public_values <= app_vm_config.system().num_public_values ); let hasher = vm_poseidon2_hasher(); let memory_dimensions = app_vm_config.system().memory_config.memory_dimensions(); - let app_program_commit: [Val; DIGEST_SIZE] = + let app_program_commit: [F; DIGEST_SIZE] = app_exe.committed_program.prover_data.commit.into(); - let leaf_verifier_program_commit: [Val; DIGEST_SIZE] = leaf_vm_verifier_exe + let leaf_verifier_program_commit: [F; DIGEST_SIZE] = leaf_vm_verifier_exe .committed_program .prover_data .commit @@ -60,8 +69,8 @@ impl AppExecutionCommit> { &hasher, ) .hash(); - let mut padded_pc_start = [Val::::ZERO; DIGEST_SIZE]; - padded_pc_start[0] = Val::::from_canonical_u32(app_exe.exe.pc_start); + let mut padded_pc_start = [F::ZERO; DIGEST_SIZE]; + padded_pc_start[0] = F::from_canonical_u32(app_exe.exe.pc_start); let app_hash = hasher.hash(&app_program_commit); let init_memory_hash = hasher.hash(&init_memory_commit); let pc_start_hash = hasher.hash(&padded_pc_start); @@ -74,3 +83,42 @@ impl AppExecutionCommit> { } } } + +pub(crate) fn babybear_digest_to_bn254(digest: &[F; DIGEST_SIZE]) -> Bn254Fr { + let mut ret = Bn254Fr::ZERO; + let order = Bn254Fr::from_canonical_u32(BabyBear::ORDER_U32); + let mut base = Bn254Fr::ONE; + digest.iter().for_each(|&x| { + ret += base * Bn254Fr::from_canonical_u32(x.as_canonical_u32()); + base *= order; + }); + ret +} + +pub fn generate_leaf_committed_exe>( + leaf_fri_params: FriParameters, + compiler_options: CompilerOptions, + app_pk: &AppProvingKey, +) -> Arc { + let app_vm_vk = app_pk.app_vm_pk.vm_pk.get_vk(); + let leaf_engine = BabyBearPoseidon2Engine::new(leaf_fri_params); + let leaf_program = LeafVmVerifierConfig { + app_fri_params: app_pk.app_vm_pk.fri_params, + app_system_config: app_pk.app_vm_pk.vm_config.system().clone(), + compiler_options, + } + .build_program(&app_vm_vk); + Arc::new(AxVmCommittedExe::commit( + leaf_program.into(), + leaf_engine.config.pcs(), + )) +} + +pub fn commit_app_exe( + app_fri_params: FriParameters, + app_exe: impl Into>, +) -> Arc { + let exe: AxVmExe<_> = app_exe.into(); + let app_engine = BabyBearPoseidon2Engine::new(app_fri_params); + Arc::new(AxVmCommittedExe::::commit(exe, app_engine.config.pcs())) +} diff --git a/crates/axvm-sdk/src/config.rs b/crates/axvm-sdk/src/config.rs index f54c7d47d4..f30ac0cb07 100644 --- a/crates/axvm-sdk/src/config.rs +++ b/crates/axvm-sdk/src/config.rs @@ -9,7 +9,8 @@ use axvm_bigint_circuit::{Int256, Int256Executor, Int256Periphery}; use axvm_bigint_transpiler::Int256TranspilerExtension; use axvm_circuit::{ arch::{ - SystemConfig, SystemExecutor, SystemPeriphery, VmChipComplex, VmConfig, VmInventoryError, + instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES, SystemConfig, SystemExecutor, + SystemPeriphery, VmChipComplex, VmConfig, VmInventoryError, }, circuit_derive::{Chip, ChipUsageGetter}, derive::{AnyEnum, InstructionExecutor}, @@ -38,10 +39,17 @@ use serde::{Deserialize, Serialize}; use crate::F; -#[derive(Clone, Debug)] -pub struct AppConfig> { +const DEFAULT_LEAF_BLOWUP: usize = 2; +const DEFAULT_INTERNAL_BLOWUP: usize = 2; +const DEFAULT_ROOT_BLOWUP: usize = 3; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct AppConfig { pub app_fri_params: FriParameters, pub app_vm_config: VC, + pub leaf_fri_params: FriParameters, + /// Only for AggVM debugging. App VM users should not need this in regular flow. + pub compiler_options: CompilerOptions, } #[derive(Clone, Copy, Debug, Serialize, Deserialize)] @@ -50,6 +58,7 @@ pub struct AggConfig { pub leaf_fri_params: FriParameters, pub internal_fri_params: FriParameters, pub root_fri_params: FriParameters, + /// Only for AggVM debugging. pub compiler_options: CompilerOptions, } @@ -69,7 +78,7 @@ pub struct FullAggConfig { pub halo2_config: Halo2Config, } -#[derive(Builder, Clone, Debug)] +#[derive(Builder, Clone, Debug, Serialize, Deserialize)] pub struct SdkVmConfig { pub system: SystemConfig, pub rv32i: Option, @@ -170,6 +179,61 @@ impl SdkVmConfig { } } +impl Default for FullAggConfig { + fn default() -> Self { + Self { + agg_config: AggConfig::default(), + halo2_config: Halo2Config { + verifier_k: 24, + wrapper_k: None, + }, + } + } +} + +impl Default for AggConfig { + fn default() -> Self { + Self { + max_num_user_public_values: DEFAULT_MAX_NUM_PUBLIC_VALUES, + leaf_fri_params: FriParameters::standard_with_100_bits_conjectured_security( + DEFAULT_LEAF_BLOWUP, + ), + internal_fri_params: FriParameters::standard_with_100_bits_conjectured_security( + DEFAULT_INTERNAL_BLOWUP, + ), + root_fri_params: FriParameters::standard_with_100_bits_conjectured_security( + DEFAULT_ROOT_BLOWUP, + ), + compiler_options: Default::default(), + } + } +} + +impl AppConfig { + pub fn new(app_fri_params: FriParameters, app_vm_config: VC) -> Self { + Self { + app_fri_params, + app_vm_config, + leaf_fri_params: FriParameters::standard_with_100_bits_conjectured_security( + DEFAULT_LEAF_BLOWUP, + ), + compiler_options: Default::default(), + } + } + pub fn new_with_leaf_fri_params( + app_fri_params: FriParameters, + app_vm_config: VC, + leaf_fri_params: FriParameters, + ) -> Self { + Self { + app_fri_params, + app_vm_config, + leaf_fri_params, + compiler_options: Default::default(), + } + } +} + impl VmConfig for SdkVmConfig { type Executor = SdkVmConfigExecutor; type Periphery = SdkVmConfigPeriphery; diff --git a/crates/axvm-sdk/src/keygen/dummy.rs b/crates/axvm-sdk/src/keygen/dummy.rs index f5a10e0b2f..ba808ddfe0 100644 --- a/crates/axvm-sdk/src/keygen/dummy.rs +++ b/crates/axvm-sdk/src/keygen/dummy.rs @@ -36,7 +36,7 @@ use crate::{ leaf::{types::LeafVmVerifierInput, LeafVmVerifierConfig}, root::types::RootVmVerifierInput, }, - F, SC, + NonRootCommittedExe, F, SC, }; /// Returns: @@ -68,7 +68,7 @@ pub(super) fn compute_root_proof_heights( pub(super) fn dummy_internal_proof( internal_vm_pk: VmProvingKey, - internal_exe: Arc>, + internal_exe: Arc, leaf_proof: Proof, ) -> Proof { let mut internal_inputs = InternalVmVerifierInput::chunk_leaf_or_internal_proofs( @@ -87,7 +87,7 @@ pub(super) fn dummy_internal_proof( pub(super) fn dummy_internal_proof_riscv_app_vm( leaf_vm_pk: VmProvingKey, internal_vm_pk: VmProvingKey, - internal_exe: Arc>, + internal_exe: Arc, num_public_values: usize, ) -> Proof { let fri_params = standard_fri_params_with_100_bits_conjectured_security(1); @@ -126,7 +126,7 @@ fn dummy_leaf_proof_impl>( ) -> Proof { let leaf_program = LeafVmVerifierConfig { app_fri_params: app_vm_pk.fri_params, - app_vm_config: app_vm_pk.vm_config.clone(), + app_system_config: app_vm_pk.vm_config.system().clone(), compiler_options: Default::default(), } .build_program(&app_vm_pk.vm_pk.get_vk()); @@ -196,7 +196,7 @@ where ContinuationVmProver::prove(&app_prover, vec![]) } -fn dummy_app_committed_exe(fri_params: FriParameters) -> Arc> { +fn dummy_app_committed_exe(fri_params: FriParameters) -> Arc { let program = dummy_app_program(); let e = BabyBearPoseidon2Engine::new(fri_params); Arc::new(AxVmCommittedExe::::commit( diff --git a/crates/axvm-sdk/src/keygen/mod.rs b/crates/axvm-sdk/src/keygen/mod.rs index 7a2151ae94..8ac91a6e65 100644 --- a/crates/axvm-sdk/src/keygen/mod.rs +++ b/crates/axvm-sdk/src/keygen/mod.rs @@ -9,9 +9,10 @@ use ax_stark_sdk::{ }, config::{ baby_bear_poseidon2::BabyBearPoseidon2Engine, - baby_bear_poseidon2_outer::BabyBearPoseidon2OuterEngine, + baby_bear_poseidon2_outer::BabyBearPoseidon2OuterEngine, FriParameters, }, engine::StarkFriEngine, + p3_bn254_fr::Bn254Fr, }; use axvm_circuit::{ arch::{VirtualMachine, VmConfig}, @@ -28,16 +29,19 @@ use dummy::{compute_root_proof_heights, dummy_internal_proof_riscv_app_vm}; use serde::{Deserialize, Serialize}; use crate::{ + commit::babybear_digest_to_bn254, config::{AggConfig, AppConfig, FullAggConfig}, keygen::perm::AirIdPermutation, - verifier::{internal::InternalVmVerifierConfig, root::RootVmVerifierConfig}, - OuterSC, F, SC, + verifier::{ + internal::InternalVmVerifierConfig, leaf::LeafVmVerifierConfig, root::RootVmVerifierConfig, + }, + NonRootCommittedExe, OuterSC, F, SC, }; pub(crate) mod dummy; pub mod perm; -#[derive(Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize)] pub struct FullAggProvingKey { pub agg_vm_pk: AggProvingKey, pub halo2_pk: Halo2ProvingKey, @@ -45,20 +49,22 @@ pub struct FullAggProvingKey { #[derive(Clone, Serialize, Deserialize)] pub struct AppProvingKey { + pub leaf_committed_exe: Arc, + pub leaf_fri_params: FriParameters, pub app_vm_pk: VmProvingKey, } pub type AppVerifyingKey = MultiStarkVerifyingKey; -#[derive(Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize)] pub struct AggProvingKey { pub leaf_vm_pk: VmProvingKey, pub internal_vm_pk: VmProvingKey, - pub internal_committed_exe: Arc>, + pub internal_committed_exe: Arc, pub root_verifier_pk: RootVerifierProvingKey, } /// Attention: the size of this struct is VERY large, usually >10GB. -#[derive(Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize)] pub struct Halo2ProvingKey { /// Static verifier to verify a stark proof of the root verifier. pub verifier: Halo2VerifierProvingKey, @@ -84,7 +90,24 @@ where vm_pk, } }; - Self { app_vm_pk } + let leaf_committed_exe = { + let leaf_engine = BabyBearPoseidon2Engine::new(config.leaf_fri_params); + let leaf_program = LeafVmVerifierConfig { + app_fri_params: config.app_fri_params, + app_system_config: config.app_vm_config.system().clone(), + compiler_options: config.compiler_options, + } + .build_program(&app_vm_pk.vm_pk.get_vk()); + Arc::new(AxVmCommittedExe::commit( + leaf_program.into(), + leaf_engine.config.pcs(), + )) + }; + Self { + leaf_committed_exe, + leaf_fri_params: config.leaf_fri_params, + app_vm_pk, + } } pub fn num_public_values(&self) -> usize { @@ -94,6 +117,15 @@ where pub fn get_vk(&self) -> AppVerifyingKey { self.app_vm_pk.vm_pk.get_vk() } + pub fn app_fri_params(&self) -> FriParameters { + self.app_vm_pk.fri_params + } + pub fn commit_in_bn254(&self) -> Bn254Fr { + babybear_digest_to_bn254(&self.commit_in_babybear()) + } + pub fn commit_in_babybear(&self) -> [F; DIGEST_SIZE] { + self.leaf_committed_exe.get_program_commit().into() + } } impl AggProvingKey { @@ -255,6 +287,7 @@ impl FullAggProvingKey { let dummy_root_proof = agg_vm_pk .root_verifier_pk .generate_dummy_root_proof(dummy_internal_proof); + // FIXME: Halo2VerifierProvingKey is not Send + Sync because Array/Usize use Rc. let verifier = agg_vm_pk .root_verifier_pk .keygen_static_verifier(halo2_config.verifier_k, dummy_root_proof); @@ -264,9 +297,10 @@ impl FullAggProvingKey { } else { Halo2WrapperProvingKey::keygen_auto_tune(dummy_snark) }; + let halo2_pk = Halo2ProvingKey { verifier, wrapper }; Self { agg_vm_pk, - halo2_pk: Halo2ProvingKey { verifier, wrapper }, + halo2_pk, } } } diff --git a/crates/axvm-sdk/src/lib.rs b/crates/axvm-sdk/src/lib.rs index d071b92955..0321a50287 100644 --- a/crates/axvm-sdk/src/lib.rs +++ b/crates/axvm-sdk/src/lib.rs @@ -2,16 +2,21 @@ extern crate core; use std::{ fs::{create_dir_all, read, write}, + io::Write, + panic::catch_unwind, path::Path, sync::Arc, }; +use ax_stark_backend::engine::StarkEngine; use ax_stark_sdk::{ - ax_stark_backend::{prover::types::Proof, verifier::VerificationError, Chip}, + ax_stark_backend::{verifier::VerificationError, Chip}, config::{ - baby_bear_poseidon2::BabyBearPoseidon2Config, + baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine}, baby_bear_poseidon2_outer::BabyBearPoseidon2OuterConfig, + FriParameters, }, + engine::StarkFriEngine, p3_baby_bear::BabyBear, }; use axvm_build::{build_guest_package, get_package, get_target_dir, GuestOptions}; @@ -21,7 +26,10 @@ use axvm_circuit::{ system::program::trace::AxVmCommittedExe, }; use axvm_native_recursion::{ - halo2::{verifier::Halo2VerifierProvingKey, EvmProof}, + halo2::{ + wrapper::{EvmVerifier, Halo2WrapperProvingKey}, + EvmProof, + }, types::InnerConfig, }; use axvm_transpiler::{ @@ -30,12 +38,11 @@ use axvm_transpiler::{ transpiler::{Transpiler, TranspilerError}, FromElf, }; -use bincode::{deserialize, serialize}; -use config::{AggConfig, AppConfig}; +use commit::commit_app_exe; +use config::AppConfig; use eyre::{bail, Result}; use itertools::Itertools; -use keygen::{AggProvingKey, AppProvingKey, AppVerifyingKey}; -use prover::{generate_leaf_committed_exe, StarkProver}; +use keygen::AppProvingKey; pub mod commit; pub mod config; @@ -48,10 +55,17 @@ pub mod verifier; mod io; pub use io::*; +use crate::{ + config::FullAggConfig, + keygen::FullAggProvingKey, + prover::{AppProver, ContinuationProver}, +}; + pub(crate) type SC = BabyBearPoseidon2Config; pub(crate) type C = InnerConfig; pub(crate) type F = BabyBear; pub(crate) type OuterSC = BabyBearPoseidon2OuterConfig; +pub(crate) type NonRootCommittedExe = AxVmCommittedExe; pub struct Sdk; @@ -97,161 +111,130 @@ impl Sdk { todo!() } - pub fn app_keygen>(&self, _config: AppConfig) -> Result> - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - todo!() + pub fn commit_app_exe( + &self, + app_fri_params: FriParameters, + exe: AxVmExe, + ) -> Result> { + Ok(commit_app_exe(app_fri_params, exe)) } - pub fn commit_app_exe>( + pub fn app_keygen, P: AsRef>( &self, - _config: AppConfig, - _exe: AxVmExe, - ) -> Result>> + config: AppConfig, + output_path: Option

, + ) -> Result> where VC::Executor: Chip, VC::Periphery: Chip, { - todo!() + let app_pk = AppProvingKey::keygen(config); + if let Some(output_path) = output_path { + if let Some(parent) = output_path.as_ref().parent() { + create_dir_all(parent)?; + } + let output: Vec = bson::to_vec(&app_pk)?; + write(output_path, output)?; + } + Ok(app_pk) } - - pub fn app_keygen_and_commit_exe>( + pub fn load_app_pk_from_file, P: AsRef>( &self, - _config: AppConfig, - _exe: AxVmExe, - ) -> Result<(AppProvingKey, Arc>)> - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - todo!() + app_pk_path: P, + ) -> Result> { + let ret = bson::from_reader(std::fs::File::open(app_pk_path)?)?; + Ok(ret) } - pub fn generate_app_proof>( + pub fn create_app_prover>( &self, - app_pk: AppProvingKey, - app_exe: Arc>, - inputs: StdIn, - ) -> Result> + app_vm_pk: AppProvingKey, + app_committed_exe: Arc, + ) -> Result> where VC::Executor: Chip, VC::Periphery: Chip, { - let prover = StarkProver::new(app_pk, app_exe); - let proof = prover.generate_app_proof(inputs); - Ok(proof) + Ok(AppProver::new(app_vm_pk.app_vm_pk, app_committed_exe)) } - pub fn verify_app_proof( + pub fn verify_app_proof>( &self, - _proof: Vec>, - _app_vk: &AppVerifyingKey, + app_pk: &AppProvingKey, + proof: &ContinuationVmProof, ) -> Result<(), VerificationError> { - todo!() + let e = BabyBearPoseidon2Engine::new(app_pk.app_vm_pk.fri_params); + for seg_proof in &proof.per_segment { + e.verify(&app_pk.app_vm_pk.vm_pk.get_vk(), seg_proof)? + } + // TODO: verify continuation. + Ok(()) } pub fn agg_keygen>( &self, - config: AggConfig, + config: FullAggConfig, output_path: Option

, - ) -> Result<(AggConfig, AggProvingKey)> { - let agg_pk: AggProvingKey = AggProvingKey::keygen(config); - let ret = (config, agg_pk); + ) -> Result { + let agg_pk = FullAggProvingKey::keygen(config); if let Some(output_path) = output_path { if let Some(parent) = output_path.as_ref().parent() { create_dir_all(parent)?; } - let output: Vec = serialize(&ret)?; + let output: Vec = bson::to_vec(&agg_pk)?; write(output_path, output)?; } - Ok(ret) + Ok(agg_pk) } pub fn load_agg_pk_from_file>( &self, agg_pk_path: P, - ) -> Result<(AggConfig, AggProvingKey)> { - let ret = deserialize(&read(agg_pk_path)?)?; + ) -> Result { + let ret = bson::from_reader(std::fs::File::open(agg_pk_path)?)?; Ok(ret) } - pub fn generate_leaf_committed_exe>( + pub fn create_e2e_prover>( &self, - config: AggConfig, - app_pk: &AppProvingKey, - ) -> Result>> - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - let leaf_exe = generate_leaf_committed_exe(&config, app_pk); - Ok(leaf_exe) - } - - pub fn agg_keygen_and_generate_leaf_committed_exe>( - &self, - config: AggConfig, - app_pk: &AppProvingKey, - ) -> Result<(AggProvingKey, Arc>)> - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - let (_, agg_pk) = self.agg_keygen(config, None::<&Path>)?; - let leaf_exe = self.generate_leaf_committed_exe(config, app_pk)?; - Ok((agg_pk, leaf_exe)) + app_pk: AppProvingKey, + app_exe: Arc, + agg_pk: FullAggProvingKey, + ) -> Result> { + Ok(ContinuationProver::new(app_pk, app_exe, agg_pk)) } - pub fn generate_static_verifier_circuit( - &self, - _agg_pk: AggProvingKey, - ) -> Result { - todo!() + pub fn load_evm_proof_from_file>(&self, evm_proof_path: P) -> Result { + let ret = bson::from_reader(std::fs::File::open(evm_proof_path)?)?; + Ok(ret) } - #[allow(clippy::too_many_arguments)] - pub fn generate_e2e_proof, P: AsRef>( + pub fn generate_snark_verifier_contract>( &self, - _app_pk: AppProvingKey, - _app_exe: Arc>, - _agg_pk: AggProvingKey, - _leaf_exe: Arc>, - _static_verifier: Halo2VerifierProvingKey, - _inputs: StdIn, - _output_path: Option

, - ) -> Result - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - todo!() - } - - pub fn load_e2e_proof_from_file>(&self, _e2e_proof_path: P) -> Result { - todo!() + full_agg_proving_key: &FullAggProvingKey, + output_path: Option

, + ) -> Result { + let evm_verifier = full_agg_proving_key + .halo2_pk + .wrapper + .generate_evm_verifier(); + if let Some(output_path) = output_path { + let mut f = std::fs::File::create(output_path)?; + f.write_all(&bson::to_vec(&evm_verifier)?)?; + } + Ok(evm_verifier) } - pub fn generate_snark_verifier_contract, P: AsRef>( - &self, - _evm_proof: EvmProof, - _app_pk: AppProvingKey, - _agg_pk: AggProvingKey, - _output_path: P, - ) -> Result<()> - where - VC::Executor: Chip, - VC::Periphery: Chip, - { - todo!() + pub fn load_snark_verifier_contract>(&self, path: P) -> Result { + let ret = bson::from_reader(std::fs::File::open(path)?)?; + Ok(ret) } - pub fn evm_verify_snark>( - &self, - _evm_proof: EvmProof, - _contract_path: P, - ) -> Result<(), VerificationError> { - todo!() + pub fn verify_evm_proof(&self, evm_verifier: &EvmVerifier, evm_proof: &EvmProof) -> bool { + // FIXME: we should return the concrete error. + catch_unwind(|| { + Halo2WrapperProvingKey::evm_verify(evm_verifier, evm_proof); + }) + .is_ok() } } diff --git a/crates/axvm-sdk/src/prover/agg.rs b/crates/axvm-sdk/src/prover/agg.rs new file mode 100644 index 0000000000..613aa6c498 --- /dev/null +++ b/crates/axvm-sdk/src/prover/agg.rs @@ -0,0 +1,208 @@ +use std::sync::Arc; + +use ax_stark_sdk::{ + ax_stark_backend::prover::types::Proof, config::baby_bear_poseidon2::BabyBearPoseidon2Engine, + engine::StarkFriEngine, +}; +#[cfg(feature = "bench-metrics")] +use axvm_circuit::arch::SingleSegmentVmExecutor; +use axvm_circuit::{ + arch::Streams, + prover::{local::VmLocalProver, ContinuationVmProof, SingleSegmentVmProver}, +}; +use axvm_native_circuit::NativeConfig; +use axvm_native_recursion::hints::Hintable; +use tracing::info_span; + +use crate::{ + keygen::AggProvingKey, + prover::RootVerifierLocalProver, + verifier::{ + internal::types::InternalVmVerifierInput, leaf::types::LeafVmVerifierInput, + root::types::RootVmVerifierInput, + }, + NonRootCommittedExe, OuterSC, F, SC, +}; + +const DEFAULT_NUM_CHILDREN_LEAF: usize = 2; +const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 2; +const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4; + +pub struct AggStarkProver { + leaf_prover: VmLocalProver, + internal_prover: VmLocalProver, + root_prover: RootVerifierLocalProver, + + pub num_children_leaf: usize, + pub num_children_internal: usize, + pub max_internal_wrapper_layers: usize, +} +impl AggStarkProver { + pub fn new(agg_pk: AggProvingKey, leaf_committed_exe: Arc) -> Self { + let leaf_prover = VmLocalProver::::new( + agg_pk.leaf_vm_pk.clone(), + leaf_committed_exe.clone(), + ); + let internal_prover = VmLocalProver::::new( + agg_pk.internal_vm_pk.clone(), + agg_pk.internal_committed_exe.clone(), + ); + let root_prover = RootVerifierLocalProver::new(agg_pk.root_verifier_pk.clone()); + Self { + leaf_prover, + internal_prover, + root_prover, + num_children_leaf: DEFAULT_NUM_CHILDREN_LEAF, + num_children_internal: DEFAULT_NUM_CHILDREN_INTERNAL, + max_internal_wrapper_layers: DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS, + } + } + pub fn with_num_children_leaf(mut self, num_children_leaf: usize) -> Self { + self.num_children_leaf = num_children_leaf; + self + } + + pub fn with_num_children_internal(mut self, num_children_internal: usize) -> Self { + self.num_children_internal = num_children_internal; + self + } + + pub fn with_max_internal_wrapper_layers(mut self, max_internal_wrapper_layers: usize) -> Self { + self.max_internal_wrapper_layers = max_internal_wrapper_layers; + self + } + + /// Generate a proof to aggregate app proofs. + pub fn generate_agg_proof(&self, app_proofs: ContinuationVmProof) -> Proof { + let leaf_proofs = info_span!("leaf verifier", group = "leaf_verifier").in_scope(|| { + #[cfg(feature = "bench-metrics")] + metrics::counter!("fri.log_blowup") + .absolute(self.leaf_prover.pk.fri_params.log_blowup as u64); + self.generate_leaf_proof_impl(&app_proofs) + }); + let public_values = app_proofs.user_public_values.public_values; + let internal_proof = self.generate_internal_proof_impl(leaf_proofs, &public_values); + info_span!("root verifier", group = "root_verifier").in_scope(|| { + #[cfg(feature = "bench-metrics")] + metrics::counter!("fri.log_blowup").absolute( + self.root_prover + .root_verifier_pk + .vm_pk + .fri_params + .log_blowup as u64, + ); + self.generate_root_proof_impl(RootVmVerifierInput { + proofs: vec![internal_proof], + public_values, + }) + }) + } + + fn generate_leaf_proof_impl(&self, app_proofs: &ContinuationVmProof) -> Vec> { + let leaf_inputs = + LeafVmVerifierInput::chunk_continuation_vm_proof(app_proofs, self.num_children_leaf); + leaf_inputs + .into_iter() + .enumerate() + .map(|(leaf_node_idx, input)| { + info_span!("leaf verifier proof", index = leaf_node_idx) + .in_scope(|| single_segment_prove(&self.leaf_prover, input.write_to_stream())) + }) + .collect::>() + } + + fn generate_internal_proof_impl( + &self, + leaf_proofs: Vec>, + public_values: &[F], + ) -> Proof { + let mut internal_node_idx = -1; + let mut internal_node_height = 0; + let mut proofs = leaf_proofs; + let mut wrapper_layers = 0; + loop { + // TODO: what's a good test case for the wrapping logic? + if proofs.len() == 1 { + // TODO: record execution time as a part of root verifier execution time. + let actual_air_heights = + self.root_prover + .execute_for_air_heights(RootVmVerifierInput { + proofs: vec![proofs[0].clone()], + public_values: public_values.to_vec(), + }); + // Root verifier can handle the internal proof. We can stop here. + if heights_le( + &actual_air_heights, + &self.root_prover.root_verifier_pk.air_heights, + ) { + break; + } + if wrapper_layers >= self.max_internal_wrapper_layers { + panic!("The heights of the root verifier still exceed the required heights after {} wrapper layers", self.max_internal_wrapper_layers); + } + wrapper_layers += 1; + } + let internal_inputs = InternalVmVerifierInput::chunk_leaf_or_internal_proofs( + self.internal_prover + .committed_exe + .get_program_commit() + .into(), + &proofs, + self.num_children_internal, + ); + let group = format!("internal_verifier_height_{}", internal_node_height); + proofs = info_span!("internal verifier", group = group).in_scope(|| { + #[cfg(feature = "bench-metrics")] + metrics::counter!("fri.log_blowup") + .absolute(self.internal_prover.pk.fri_params.log_blowup as u64); + internal_inputs + .into_iter() + .map(|input| { + internal_node_idx += 1; + info_span!( + "Internal verifier proof", + index = internal_node_idx, + height = internal_node_height + ) + .in_scope(|| single_segment_prove(&self.internal_prover, input.write())) + }) + .collect() + }); + internal_node_height += 1; + } + proofs.pop().unwrap() + } + + fn generate_root_proof_impl(&self, root_input: RootVmVerifierInput) -> Proof { + let input = root_input.write(); + let root_prover = &self.root_prover; + #[cfg(feature = "bench-metrics")] + { + let mut vm_config = root_prover.root_verifier_pk.vm_pk.vm_config.clone(); + vm_config.system.collect_metrics = true; + let vm = SingleSegmentVmExecutor::new(vm_config); + let exe = root_prover.root_verifier_pk.root_committed_exe.exe.clone(); + vm.execute(exe, input.clone()).unwrap(); + } + SingleSegmentVmProver::prove(root_prover, input) + } +} + +fn single_segment_prove>( + prover: &VmLocalProver, + input: impl Into> + Clone, +) -> Proof { + #[cfg(feature = "bench-metrics")] + { + let mut vm_config = prover.pk.vm_config.clone(); + vm_config.system.collect_metrics = true; + let vm = SingleSegmentVmExecutor::new(vm_config); + vm.execute(prover.committed_exe.exe.clone(), input.clone()) + .unwrap(); + } + SingleSegmentVmProver::prove(prover, input) +} +fn heights_le(a: &[usize], b: &[usize]) -> bool { + assert_eq!(a.len(), b.len()); + a.iter().zip(b.iter()).all(|(a, b)| a <= b) +} diff --git a/crates/axvm-sdk/src/prover/app.rs b/crates/axvm-sdk/src/prover/app.rs new file mode 100644 index 0000000000..cf990b2870 --- /dev/null +++ b/crates/axvm-sdk/src/prover/app.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use ax_stark_backend::Chip; +use ax_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine; +#[cfg(feature = "bench-metrics")] +use axvm_circuit::arch::{instructions::exe::AxVmExe, VmExecutor}; +use axvm_circuit::{ + arch::VmConfig, + prover::{ + local::VmLocalProver, types::VmProvingKey, ContinuationVmProof, ContinuationVmProver, + }, +}; +use tracing::info_span; + +use crate::{NonRootCommittedExe, StdIn, F, SC}; + +pub struct AppProver { + /// If true, will run execution once with full metric collection for + /// flamegraphs (WARNING: this degrades performance). + pub profile: bool, + app_prover: VmLocalProver, +} + +impl AppProver { + pub fn new(app_vm_pk: VmProvingKey, app_committed_exe: Arc) -> Self + where + VC: VmConfig, + { + Self { + profile: false, + app_prover: VmLocalProver::::new( + app_vm_pk, + app_committed_exe, + ), + } + } + pub fn generate_app_proof(&self, input: StdIn) -> ContinuationVmProof + where + VC: VmConfig, + VC::Executor: Chip, + VC::Periphery: Chip, + { + info_span!("app proof", group = "app_proof").in_scope(|| { + #[cfg(feature = "bench-metrics")] + if self.profile { + emit_app_execution_metrics( + self.app_prover.pk.vm_config.clone(), + self.app_prover.committed_exe.exe.clone(), + input.clone(), + ); + } + #[cfg(feature = "bench-metrics")] + metrics::counter!("fri.log_blowup") + .absolute(self.app_prover.pk.fri_params.log_blowup as u64); + ContinuationVmProver::prove(&self.app_prover, input) + }) + } +} + +#[cfg(feature = "bench-metrics")] +fn emit_app_execution_metrics>(mut vm_config: VC, exe: AxVmExe, input: StdIn) +where + VC::Executor: Chip, + VC::Periphery: Chip, +{ + vm_config.system_mut().collect_metrics = true; + let vm = VmExecutor::new(vm_config); + vm.execute_segments(exe, input).unwrap(); +} diff --git a/crates/axvm-sdk/src/prover/exe.rs b/crates/axvm-sdk/src/prover/exe.rs deleted file mode 100644 index 2b70f65fc6..0000000000 --- a/crates/axvm-sdk/src/prover/exe.rs +++ /dev/null @@ -1,45 +0,0 @@ -use std::sync::Arc; - -use ax_stark_sdk::{ - ax_stark_backend::config::StarkGenericConfig, - config::baby_bear_poseidon2::BabyBearPoseidon2Engine, engine::StarkFriEngine, -}; -use axvm_circuit::{ - arch::{instructions::exe::AxVmExe, VmConfig}, - system::program::trace::AxVmCommittedExe, -}; - -use crate::{ - config::{AggConfig, AppConfig}, - keygen::AppProvingKey, - verifier::leaf::LeafVmVerifierConfig, - F, SC, -}; - -pub fn commit_app_exe>( - app_config: &AppConfig, - app_exe: impl Into>, -) -> Arc> { - let mut exe: AxVmExe<_> = app_exe.into(); - exe.program.max_num_public_values = app_config.app_vm_config.system().num_public_values; - let app_engine = BabyBearPoseidon2Engine::new(app_config.app_fri_params); - Arc::new(AxVmCommittedExe::::commit(exe, app_engine.config.pcs())) -} - -pub fn generate_leaf_committed_exe>( - agg_config: &AggConfig, - app_pk: &AppProvingKey, -) -> Arc> { - let app_vm_vk = app_pk.app_vm_pk.vm_pk.get_vk(); - let leaf_engine = BabyBearPoseidon2Engine::new(agg_config.leaf_fri_params); - let leaf_program = LeafVmVerifierConfig { - app_fri_params: app_pk.app_vm_pk.fri_params, - app_vm_config: app_pk.app_vm_pk.vm_config.clone(), - compiler_options: agg_config.compiler_options, - } - .build_program(&app_vm_vk); - Arc::new(AxVmCommittedExe::commit( - leaf_program.into(), - leaf_engine.config.pcs(), - )) -} diff --git a/crates/axvm-sdk/src/prover/halo2.rs b/crates/axvm-sdk/src/prover/halo2.rs index 173b5067b5..fda0362227 100644 --- a/crates/axvm-sdk/src/prover/halo2.rs +++ b/crates/axvm-sdk/src/prover/halo2.rs @@ -6,6 +6,7 @@ use axvm_native_recursion::{ halo2::{utils::read_params, EvmProof, Halo2Params}, witness::Witnessable, }; +use tracing::info_span; use crate::{keygen::Halo2ProvingKey, OuterSC}; pub struct Halo2Prover { @@ -33,12 +34,15 @@ impl Halo2Prover { pub fn prove_for_evm(&self, root_proof: &Proof) -> EvmProof { let mut witness = Witness::default(); root_proof.write(&mut witness); - let snark = self - .halo2_pk - .verifier - .prove_with_loaded_params(&self.verifier_srs, witness); - self.halo2_pk - .wrapper - .prove_for_evm_with_loaded_params(&self.wrapper_srs, snark) + let snark = info_span!("halo2 verifier", group = "halo2_verifier").in_scope(|| { + self.halo2_pk + .verifier + .prove_with_loaded_params(&self.verifier_srs, witness) + }); + info_span!("halo2 wrapper", group = "halo2_wrapper").in_scope(|| { + self.halo2_pk + .wrapper + .prove_for_evm_with_loaded_params(&self.wrapper_srs, snark) + }) } } diff --git a/crates/axvm-sdk/src/prover/mod.rs b/crates/axvm-sdk/src/prover/mod.rs index 6e3fcd7752..48107481de 100644 --- a/crates/axvm-sdk/src/prover/mod.rs +++ b/crates/axvm-sdk/src/prover/mod.rs @@ -1,335 +1,60 @@ use std::sync::Arc; -use ax_stark_sdk::{ - ax_stark_backend::{prover::types::Proof, Chip}, - config::baby_bear_poseidon2::BabyBearPoseidon2Engine, - engine::StarkFriEngine, -}; -#[cfg(feature = "bench-metrics")] -use axvm_circuit::arch::{SingleSegmentVmExecutor, VmExecutor}; -use axvm_circuit::{ - arch::{Streams, VmConfig}, - prover::{ - local::VmLocalProver, ContinuationVmProof, ContinuationVmProver, SingleSegmentVmProver, - }, - system::program::trace::AxVmCommittedExe, -}; -use axvm_native_circuit::NativeConfig; -use axvm_native_recursion::hints::Hintable; -use metrics::counter; -use tracing::info_span; - -use crate::{ - config::AggConfig, - io::StdIn, - keygen::{AggProvingKey, AppProvingKey}, - verifier::{ - internal::types::InternalVmVerifierInput, leaf::types::LeafVmVerifierInput, - root::types::RootVmVerifierInput, - }, - OuterSC, F, SC, -}; - -mod exe; -pub use exe::*; -pub mod halo2; +use ax_stark_sdk::ax_stark_backend::Chip; +use axvm_circuit::arch::VmConfig; +use axvm_native_recursion::halo2::EvmProof; + +use crate::{io::StdIn, keygen::AppProvingKey, NonRootCommittedExe, F, SC}; + +mod agg; +pub use agg::*; +mod app; +pub use app::*; +mod halo2; +#[allow(unused_imports)] +pub use halo2::*; mod root; - pub use root::*; +mod stark; +#[allow(unused_imports)] +pub use stark::*; -const DEFAULT_NUM_CHILDREN_LEAF: usize = 2; -const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 2; -const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4; - -pub struct StarkProver { - pub app_pk: AppProvingKey, - pub app_committed_exe: Arc>, - app_prover: VmLocalProver, - - pub agg_pk: Option, - pub leaf_committed_exe: Option>>, - leaf_prover: Option>, - internal_prover: Option>, - root_prover: Option, +use crate::{ + keygen::FullAggProvingKey, + prover::{halo2::Halo2Prover, stark::StarkProver}, +}; - pub num_children_leaf: usize, - pub num_children_internal: usize, - pub max_internal_wrapper_layers: usize, +pub struct ContinuationProver { + stark_prover: StarkProver, + halo2_prover: Halo2Prover, } -impl> StarkProver -where - VC::Executor: Chip, - VC::Periphery: Chip, -{ - pub fn new(app_pk: AppProvingKey, app_committed_exe: Arc>) -> Self { - let app_prover = VmLocalProver::::new( - app_pk.app_vm_pk.clone(), - app_committed_exe.clone(), - ); +impl ContinuationProver { + pub fn new( + app_pk: AppProvingKey, + app_committed_exe: Arc, + full_agg_pk: FullAggProvingKey, + ) -> Self + where + VC: VmConfig, + { + let FullAggProvingKey { + agg_vm_pk, + halo2_pk, + } = full_agg_pk; + let stark_prover = StarkProver::new(app_pk, app_committed_exe, agg_vm_pk); Self { - app_pk, - agg_pk: None, - app_committed_exe, - leaf_committed_exe: None, - num_children_leaf: DEFAULT_NUM_CHILDREN_LEAF, - num_children_internal: DEFAULT_NUM_CHILDREN_INTERNAL, - max_internal_wrapper_layers: DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS, - app_prover, - leaf_prover: None, - internal_prover: None, - root_prover: None, + stark_prover, + halo2_prover: Halo2Prover::new(halo2_pk), } } - - pub fn with_agg_config(self, agg_config: AggConfig) -> Self { - let leaf_committed_exe = generate_leaf_committed_exe(&agg_config, &self.app_pk); - let agg_pk = AggProvingKey::keygen(agg_config); - self.with_agg_pk_and_leaf_committed_exe(agg_pk, leaf_committed_exe) - } - - pub fn with_agg_pk_and_leaf_committed_exe( - mut self, - agg_pk: AggProvingKey, - leaf_committed_exe: Arc>, - ) -> Self { - assert_eq!(self.app_pk.num_public_values(), agg_pk.num_public_values()); - - let leaf_prover = VmLocalProver::::new( - agg_pk.leaf_vm_pk.clone(), - leaf_committed_exe.clone(), - ); - let internal_prover = VmLocalProver::::new( - agg_pk.internal_vm_pk.clone(), - agg_pk.internal_committed_exe.clone(), - ); - let root_prover = RootVerifierLocalProver::new(agg_pk.root_verifier_pk.clone()); - - self.agg_pk = Some(agg_pk); - self.leaf_committed_exe = Some(leaf_committed_exe); - self.leaf_prover = Some(leaf_prover); - self.internal_prover = Some(internal_prover); - self.root_prover = Some(root_prover); - self - } - - pub fn with_num_children_leaf(mut self, num_children_leaf: usize) -> Self { - self.num_children_leaf = num_children_leaf; - self - } - - pub fn with_num_children_internal(mut self, num_children_internal: usize) -> Self { - self.num_children_internal = num_children_internal; - self - } - - pub fn with_max_internal_wrapper_layers(mut self, max_internal_wrapper_layers: usize) -> Self { - self.max_internal_wrapper_layers = max_internal_wrapper_layers; - self - } - - pub fn agg_pk(&self) -> &AggProvingKey { - assert!(self.agg_pk.is_some(), "Aggregation has not been configured"); - self.agg_pk.as_ref().unwrap() - } - - pub fn leaf_committed_exe(&self) -> &Arc> { - assert!( - self.leaf_committed_exe.is_some(), - "Aggregation has not been configured" - ); - self.leaf_committed_exe.as_ref().unwrap() - } - - pub fn generate_e2e_proof(&self, input: StdIn) -> Proof { - assert!(self.agg_pk.is_some(), "Aggregation has not been configured"); - let app_proofs = self.generate_app_proof(input); - let leaf_proofs = self.generate_leaf_proof_impl(&app_proofs); - let public_values = app_proofs.user_public_values.public_values; - let internal_proof = self.generate_internal_proof_impl(leaf_proofs, &public_values); - self.generate_root_proof_impl(RootVmVerifierInput { - proofs: vec![internal_proof], - public_values, - }) - } - - pub fn generate_app_proof(&self, input: StdIn) -> ContinuationVmProof { - #[cfg(feature = "bench-metrics")] - { - execute_app_exe_for_metrics_collection( - &self.app_pk, - &self.app_committed_exe, - input.clone(), - ); - } - ContinuationVmProver::prove(&self.app_prover, input) - } - - pub fn generate_e2e_proof_with_metric_spans( - &self, - input: StdIn, - program_name: &str, - ) -> Proof { - assert!(self.agg_pk.is_some(), "Aggregation has not been configured"); - let app_proofs = self.generate_app_proof_with_metric_spans(input, program_name); - let leaf_proofs = info_span!("leaf verifier", group = "leaf_verifier").in_scope(|| { - counter!("fri.log_blowup") - .absolute(self.agg_pk().leaf_vm_pk.fri_params.log_blowup as u64); - self.generate_leaf_proof_impl(&app_proofs) - }); - let public_values = app_proofs.user_public_values.public_values; - let internal_proof = self.generate_internal_proof_impl(leaf_proofs, &public_values); - info_span!("root verifier", group = "root_verifier").in_scope(|| { - counter!("fri.log_blowup") - .absolute(self.agg_pk().root_verifier_pk.vm_pk.fri_params.log_blowup as u64); - self.generate_root_proof_impl(RootVmVerifierInput { - proofs: vec![internal_proof], - public_values, - }) - }) - } - - pub fn generate_app_proof_with_metric_spans( - &self, - input: StdIn, - program_name: &str, - ) -> ContinuationVmProof { - let group_name = program_name.replace(" ", "_").to_lowercase(); - info_span!("App Continuation Program", group = group_name).in_scope(|| { - counter!("fri.log_blowup").absolute(self.app_pk.app_vm_pk.fri_params.log_blowup as u64); - self.generate_app_proof(input) - }) - } - - fn generate_leaf_proof_impl(&self, app_proofs: &ContinuationVmProof) -> Vec> { - let leaf_inputs = - LeafVmVerifierInput::chunk_continuation_vm_proof(app_proofs, self.num_children_leaf); - leaf_inputs - .into_iter() - .enumerate() - .map(|(leaf_node_idx, input)| { - info_span!("leaf verifier proof", index = leaf_node_idx).in_scope(|| { - single_segment_prove( - self.leaf_prover.as_ref().unwrap(), - input.write_to_stream(), - ) - }) - }) - .collect::>() - } - - fn generate_internal_proof_impl( - &self, - leaf_proofs: Vec>, - public_values: &[F], - ) -> Proof { - let mut internal_node_idx = -1; - let mut internal_node_height = 0; - let mut proofs = leaf_proofs; - let mut wrapper_layers = 0; - loop { - // TODO: what's a good test case for the wrapping logic? - if proofs.len() == 1 { - let root_prover = self.root_prover.as_ref().unwrap(); - // TODO: record execution time as a part of root verifier execution time. - let actual_air_heights = root_prover.execute_for_air_heights(RootVmVerifierInput { - proofs: vec![proofs[0].clone()], - public_values: public_values.to_vec(), - }); - // Root verifier can handle the internal proof. We can stop here. - if heights_le( - &actual_air_heights, - &root_prover.root_verifier_pk.air_heights, - ) { - break; - } - if wrapper_layers >= self.max_internal_wrapper_layers { - panic!("The heights of the root verifier still exceed the required heights after {} wrapper layers", self.max_internal_wrapper_layers); - } - wrapper_layers += 1; - } - let internal_inputs = InternalVmVerifierInput::chunk_leaf_or_internal_proofs( - self.agg_pk() - .internal_committed_exe - .get_program_commit() - .into(), - &proofs, - self.num_children_internal, - ); - let group = format!("internal_verifier_height_{}", internal_node_height); - proofs = info_span!("internal verifier", group = group).in_scope(|| { - counter!("fri.log_blowup") - .absolute(self.agg_pk().internal_vm_pk.fri_params.log_blowup as u64); - internal_inputs - .into_iter() - .map(|input| { - internal_node_idx += 1; - info_span!( - "Internal verifier proof", - index = internal_node_idx, - height = internal_node_height - ) - .in_scope(|| { - single_segment_prove( - self.internal_prover.as_ref().unwrap(), - input.write(), - ) - }) - }) - .collect() - }); - internal_node_height += 1; - } - proofs.pop().unwrap() - } - - fn generate_root_proof_impl(&self, root_input: RootVmVerifierInput) -> Proof { - let input = root_input.write(); - let root_prover = self.root_prover.as_ref().unwrap(); - #[cfg(feature = "bench-metrics")] - { - let mut vm_config = root_prover.root_verifier_pk.vm_pk.vm_config.clone(); - vm_config.system.collect_metrics = true; - let vm = SingleSegmentVmExecutor::new(vm_config); - let exe = root_prover.root_verifier_pk.root_committed_exe.exe.clone(); - vm.execute(exe, input.clone()).unwrap(); - } - SingleSegmentVmProver::prove(root_prover, input) - } -} - -fn single_segment_prove>( - prover: &VmLocalProver, - input: impl Into> + Clone, -) -> Proof { - #[cfg(feature = "bench-metrics")] + pub fn generate_proof_for_evm(&self, input: StdIn) -> EvmProof + where + VC: VmConfig, + VC::Executor: Chip, + VC::Periphery: Chip, { - let mut vm_config = prover.pk.vm_config.clone(); - vm_config.system.collect_metrics = true; - let vm = SingleSegmentVmExecutor::new(vm_config); - vm.execute(prover.committed_exe.exe.clone(), input.clone()) - .unwrap(); + let root_proof = self.stark_prover.generate_proof_for_outer_recursion(input); + self.halo2_prover.prove_for_evm(&root_proof) } - SingleSegmentVmProver::prove(prover, input) -} - -#[cfg(feature = "bench-metrics")] -fn execute_app_exe_for_metrics_collection>( - app_pk: &AppProvingKey, - app_committed_exe: &Arc>, - input: StdIn, -) where - VC::Executor: Chip, - VC::Periphery: Chip, -{ - let mut vm_config = app_pk.app_vm_pk.vm_config.clone(); - vm_config.system_mut().collect_metrics = true; - let vm = VmExecutor::new(vm_config); - vm.execute_segments(app_committed_exe.exe.clone(), input) - .unwrap(); -} - -fn heights_le(a: &[usize], b: &[usize]) -> bool { - assert_eq!(a.len(), b.len()); - a.iter().zip(b.iter()).all(|(a, b)| a <= b) } diff --git a/crates/axvm-sdk/src/prover/stark.rs b/crates/axvm-sdk/src/prover/stark.rs new file mode 100644 index 0000000000..64383ecddf --- /dev/null +++ b/crates/axvm-sdk/src/prover/stark.rs @@ -0,0 +1,54 @@ +use std::sync::Arc; + +use ax_stark_backend::{prover::types::Proof, Chip}; +use axvm_circuit::arch::VmConfig; + +use crate::{ + keygen::{AggProvingKey, AppProvingKey}, + prover::{agg::AggStarkProver, app::AppProver}, + NonRootCommittedExe, OuterSC, StdIn, F, SC, +}; + +pub struct StarkProver { + app_prover: AppProver, + agg_prover: AggStarkProver, +} +impl StarkProver { + pub fn new( + app_pk: AppProvingKey, + app_committed_exe: Arc, + agg_pk: AggProvingKey, + ) -> Self + where + VC: VmConfig, + { + let AppProvingKey { + leaf_committed_exe, + leaf_fri_params, + app_vm_pk, + } = app_pk; + assert_eq!( + leaf_fri_params, agg_pk.leaf_vm_pk.fri_params, + "App VM is incompatible with Agg VM because of leaf FRI parameters" + ); + assert_eq!( + app_vm_pk.vm_config.system().num_public_values, + agg_pk.num_public_values(), + "App VM is incompatible with Agg VM because of the number of public values" + ); + + Self { + app_prover: AppProver::new(app_vm_pk, app_committed_exe), + agg_prover: AggStarkProver::new(agg_pk, leaf_committed_exe), + } + } + pub fn generate_proof_for_outer_recursion(&self, input: StdIn) -> Proof + where + VC: VmConfig, + VC::Executor: Chip, + VC::Periphery: Chip, + { + let app_proof = self.app_prover.generate_app_proof(input); + self.agg_prover.generate_agg_proof(app_proof) + } +} diff --git a/crates/axvm-sdk/src/static_verifier/mod.rs b/crates/axvm-sdk/src/static_verifier/mod.rs index e780ba8d6c..2937c8f0ed 100644 --- a/crates/axvm-sdk/src/static_verifier/mod.rs +++ b/crates/axvm-sdk/src/static_verifier/mod.rs @@ -1,5 +1,3 @@ -mod utils; - use ax_stark_sdk::{ ax_stark_backend::{ p3_field::{AbstractField, PrimeField32}, diff --git a/crates/axvm-sdk/src/static_verifier/utils.rs b/crates/axvm-sdk/src/static_verifier/utils.rs deleted file mode 100644 index 8b13789179..0000000000 --- a/crates/axvm-sdk/src/static_verifier/utils.rs +++ /dev/null @@ -1 +0,0 @@ - diff --git a/crates/axvm-sdk/src/verifier/leaf/mod.rs b/crates/axvm-sdk/src/verifier/leaf/mod.rs index 37217d2057..827b6bd8b6 100644 --- a/crates/axvm-sdk/src/verifier/leaf/mod.rs +++ b/crates/axvm-sdk/src/verifier/leaf/mod.rs @@ -6,7 +6,7 @@ use ax_stark_sdk::{ config::{baby_bear_poseidon2::BabyBearPoseidon2Config, FriParameters}, }; use axvm_circuit::{ - arch::{instructions::program::Program, VmConfig}, + arch::{instructions::program::Program, SystemConfig}, system::memory::tree::public_values::PUBLIC_VALUES_ADDRESS_SPACE_OFFSET, }; use axvm_native_compiler::{conversion::CompilerOptions, prelude::*}; @@ -33,13 +33,13 @@ pub mod types; mod vars; /// Config to generate leaf VM verifier program. -pub struct LeafVmVerifierConfig> { +pub struct LeafVmVerifierConfig { pub app_fri_params: FriParameters, - pub app_vm_config: VC, + pub app_system_config: SystemConfig, pub compiler_options: CompilerOptions, } -impl> LeafVmVerifierConfig { +impl LeafVmVerifierConfig { pub fn build_program( &self, app_vm_vk: &MultiStarkVerifyingKey, @@ -111,17 +111,12 @@ impl> LeafVmVerifierConfig { &self, builder: &mut Builder, ) -> ([Felt; DIGEST_SIZE], [Felt; DIGEST_SIZE]) { - let memory_dimensions = self - .app_vm_config - .system() - .memory_config - .memory_dimensions(); + let memory_dimensions = self.app_system_config.memory_config.memory_dimensions(); let pv_as = F::from_canonical_usize( PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset, ); let pv_start_idx = memory_dimensions.label_to_index((pv_as, 0)); - let pv_height = - log2_strict_usize(self.app_vm_config.system().num_public_values / DIGEST_SIZE); + let pv_height = log2_strict_usize(self.app_system_config.num_public_values / DIGEST_SIZE); let proof_len = memory_dimensions.overall_height() - pv_height; let idx_prefix = pv_start_idx >> pv_height; diff --git a/crates/axvm-sdk/tests/integration_test.rs b/crates/axvm-sdk/tests/integration_test.rs index 1e0c5dca9c..2279ba8ab7 100644 --- a/crates/axvm-sdk/tests/integration_test.rs +++ b/crates/axvm-sdk/tests/integration_test.rs @@ -1,7 +1,7 @@ -use std::{borrow::Borrow, path::PathBuf}; +use std::{borrow::Borrow, path::PathBuf, sync::Arc}; use ax_stark_sdk::{ - ax_stark_backend::{p3_field::AbstractField, prover::types::Proof, Chip}, + ax_stark_backend::{p3_field::AbstractField, Chip}, config::{ baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine}, fri_params::standard_fri_params_with_100_bits_conjectured_security, @@ -15,55 +15,56 @@ use axvm_circuit::{ hasher::poseidon2::vm_poseidon2_hasher, ExecutionError, SingleSegmentVmExecutor, SystemConfig, VmConfig, VmExecutor, }, - system::memory::tree::public_values::UserPublicValuesProof, + system::{ + memory::tree::public_values::UserPublicValuesProof, program::trace::AxVmCommittedExe, + }, }; use axvm_native_circuit::{Native, NativeConfig}; use axvm_native_compiler::{conversion::CompilerOptions, prelude::*}; use axvm_native_recursion::types::InnerConfig; use axvm_rv32im_transpiler::{Rv32ITranspilerExtension, Rv32MTranspilerExtension}; use axvm_sdk::{ - commit::AppExecutionCommit, - config::{AggConfig, AppConfig, SdkVmConfig}, - keygen::{AggProvingKey, AppProvingKey}, - prover::{commit_app_exe, generate_leaf_committed_exe, StarkProver}, + config::{AggConfig, AppConfig, FullAggConfig, Halo2Config}, + keygen::AppProvingKey, verifier::{ common::types::VmVerifierPvs, leaf::types::{LeafVmVerifierInput, UserPublicValuesRootProof}, - root::types::RootVmVerifierPvs, }, Sdk, StdIn, }; use axvm_transpiler::transpiler::Transpiler; -use utils::{assert_agg_config_eq, assert_agg_pk_eq}; - -mod utils; type SC = BabyBearPoseidon2Config; type C = InnerConfig; type F = BabyBear; const NUM_PUB_VALUES: usize = 16; +const LEAF_LOG_BLOWUP: usize = 2; +const INTERNAL_LOG_BLOWUP: usize = 3; +const ROOT_LOG_BLOWUP: usize = 4; -// TODO: keygen agg_pk once for all IT tests and store in a file -fn load_agg_pk_into_e2e_prover>( - app_config: AppConfig, -) -> (StarkProver, Proof) +fn run_leaf_verifier>( + leaf_vm: &SingleSegmentVmExecutor, + leaf_committed_exe: Arc>, + verifier_input: LeafVmVerifierInput, +) -> Result, ExecutionError> where VC::Executor: Chip, VC::Periphery: Chip, { - let agg_config = AggConfig { - max_num_user_public_values: NUM_PUB_VALUES, - leaf_fri_params: standard_fri_params_with_100_bits_conjectured_security(4), - internal_fri_params: standard_fri_params_with_100_bits_conjectured_security(3), - root_fri_params: standard_fri_params_with_100_bits_conjectured_security(2), - compiler_options: CompilerOptions { - enable_cycle_tracker: true, - compile_prints: true, - ..Default::default() - }, - }; + let exe_result = leaf_vm.execute( + leaf_committed_exe.exe.clone(), + verifier_input.write_to_stream(), + )?; + let runtime_pvs: Vec<_> = exe_result + .public_values + .iter() + .map(|v| v.unwrap()) + .collect(); + Ok(runtime_pvs) +} +fn app_committed_exe_for_test(app_log_blowup: usize) -> Arc> { let program = { let n = 200; let mut builder = Builder::::default(); @@ -78,42 +79,42 @@ where builder.halt(); builder.compile_isa() }; - - let app_committed_exe = commit_app_exe(&app_config, program); - let app_pk = AppProvingKey::keygen(app_config); - let leaf_committed_exe = generate_leaf_committed_exe(&agg_config, &app_pk); - let (agg_pk, dummy) = AggProvingKey::dummy_proof_and_keygen(agg_config); - ( - StarkProver::new(app_pk, app_committed_exe) - .with_agg_pk_and_leaf_committed_exe(agg_pk, leaf_committed_exe), - dummy, + Sdk.commit_app_exe( + standard_fri_params_with_100_bits_conjectured_security(app_log_blowup), + program.into(), ) + .unwrap() } -fn run_leaf_verifier>( - verifier_input: LeafVmVerifierInput, - e2e_prover: &StarkProver, -) -> Result, ExecutionError> -where - VC::Executor: Chip, - VC::Periphery: Chip, -{ - let leaf_vm = SingleSegmentVmExecutor::new(e2e_prover.agg_pk().leaf_vm_pk.vm_config.clone()); - let exe_result = leaf_vm.execute( - e2e_prover.leaf_committed_exe().exe.clone(), - verifier_input.write_to_stream(), - )?; - let runtime_pvs: Vec<_> = exe_result - .public_values - .iter() - .map(|v| v.unwrap()) - .collect(); - Ok(runtime_pvs) +fn full_agg_config_for_test() -> FullAggConfig { + FullAggConfig { + agg_config: agg_config_for_test(), + halo2_config: Halo2Config { + verifier_k: 24, + wrapper_k: None, + }, + } } -fn small_test_app_config(log_blowup_factor: usize) -> AppConfig { +fn agg_config_for_test() -> AggConfig { + AggConfig { + max_num_user_public_values: NUM_PUB_VALUES, + leaf_fri_params: standard_fri_params_with_100_bits_conjectured_security(LEAF_LOG_BLOWUP), + internal_fri_params: standard_fri_params_with_100_bits_conjectured_security( + INTERNAL_LOG_BLOWUP, + ), + root_fri_params: standard_fri_params_with_100_bits_conjectured_security(ROOT_LOG_BLOWUP), + compiler_options: CompilerOptions { + enable_cycle_tracker: true, + compile_prints: true, + ..Default::default() + }, + } +} + +fn small_test_app_config(app_log_blowup: usize) -> AppConfig { AppConfig { - app_fri_params: standard_fri_params_with_100_bits_conjectured_security(log_blowup_factor), + app_fri_params: standard_fri_params_with_100_bits_conjectured_security(app_log_blowup), app_vm_config: NativeConfig::new( SystemConfig::default() .with_max_segment_len(200) @@ -121,39 +122,51 @@ fn small_test_app_config(log_blowup_factor: usize) -> AppConfig { .with_public_values(16), Native, ), + leaf_fri_params: standard_fri_params_with_100_bits_conjectured_security(LEAF_LOG_BLOWUP), + compiler_options: CompilerOptions { + enable_cycle_tracker: true, + ..Default::default() + }, } } #[test] fn test_public_values_and_leaf_verification() { - let app_config = small_test_app_config(3); - let (e2e_prover, _) = load_agg_pk_into_e2e_prover(app_config); + let app_log_blowup = 3; + let app_config = small_test_app_config(app_log_blowup); + let app_pk = AppProvingKey::keygen(app_config); + let app_committed_exe = app_committed_exe_for_test(app_log_blowup); + + let agg_config = agg_config_for_test(); + let leaf_vm_config = agg_config.leaf_vm_config(); + let leaf_vm = SingleSegmentVmExecutor::new(leaf_vm_config); + let leaf_committed_exe = app_pk.leaf_committed_exe.clone(); - let app_engine = BabyBearPoseidon2Engine::new(e2e_prover.app_pk.app_vm_pk.fri_params); - let app_vm = VmExecutor::new(e2e_prover.app_pk.app_vm_pk.vm_config.clone()); + let app_engine = BabyBearPoseidon2Engine::new(app_pk.app_vm_pk.fri_params); + let app_vm = VmExecutor::new(app_pk.app_vm_pk.vm_config.clone()); let app_vm_result = app_vm - .execute_and_generate_with_cached_program(e2e_prover.app_committed_exe.clone(), vec![]) + .execute_and_generate_with_cached_program(app_committed_exe.clone(), vec![]) .unwrap(); assert!(app_vm_result.per_segment.len() > 2); let mut app_vm_seg_proofs: Vec<_> = app_vm_result .per_segment .into_iter() - .map(|proof_input| app_engine.prove(&e2e_prover.app_pk.app_vm_pk.vm_pk, proof_input)) + .map(|proof_input| app_engine.prove(&app_pk.app_vm_pk.vm_pk, proof_input)) .collect(); let app_last_proof = app_vm_seg_proofs.pop().unwrap(); - let expected_app_commit: [F; DIGEST_SIZE] = - e2e_prover.app_committed_exe.get_program_commit().into(); + let expected_app_commit: [F; DIGEST_SIZE] = app_committed_exe.get_program_commit().into(); // Verify all segments except the last one. let (first_seg_final_pc, first_seg_final_mem_root) = { let runtime_pvs = run_leaf_verifier( + &leaf_vm, + leaf_committed_exe.clone(), LeafVmVerifierInput { proofs: app_vm_seg_proofs.clone(), public_values_root_proof: None, }, - &e2e_prover, ) .expect("failed to verify the first segment"); @@ -170,7 +183,7 @@ fn test_public_values_and_leaf_verification() { let pv_proof = UserPublicValuesProof::compute( app_vm.config.system.memory_config.memory_dimensions(), - e2e_prover.agg_pk().num_public_values(), + NUM_PUB_VALUES, &vm_poseidon2_hasher(), app_vm_result.final_memory.as_ref().unwrap(), ); @@ -179,11 +192,12 @@ fn test_public_values_and_leaf_verification() { // Verify the last segment with the correct public values root proof. { let runtime_pvs = run_leaf_verifier( + &leaf_vm, + leaf_committed_exe.clone(), LeafVmVerifierInput { proofs: vec![app_last_proof.clone()], public_values_root_proof: Some(pv_root_proof.clone()), }, - &e2e_prover, ) .expect("failed to verify the second segment"); @@ -204,11 +218,12 @@ fn test_public_values_and_leaf_verification() { let mut wrong_pv_root_proof = pv_root_proof.clone(); wrong_pv_root_proof.public_values_commit[0] += F::ONE; let execution_result = run_leaf_verifier( + &leaf_vm, + leaf_committed_exe.clone(), LeafVmVerifierInput { proofs: vec![app_last_proof.clone()], public_values_root_proof: Some(wrong_pv_root_proof), }, - &e2e_prover, ); match execution_result.err().unwrap() { ExecutionError::Fail { .. } => {} @@ -223,11 +238,12 @@ fn test_public_values_and_leaf_verification() { let mut wrong_pv_root_proof = pv_root_proof.clone(); wrong_pv_root_proof.sibling_hashes[0][0] += F::ONE; let execution_result = run_leaf_verifier( + &leaf_vm, + leaf_committed_exe.clone(), LeafVmVerifierInput { proofs: vec![app_last_proof.clone()], public_values_root_proof: Some(wrong_pv_root_proof), }, - &e2e_prover, ); match execution_result.err().unwrap() { ExecutionError::Fail { .. } => {} @@ -236,75 +252,24 @@ fn test_public_values_and_leaf_verification() { } } -#[test] -fn test_e2e_proof_generation() { - let app_config = small_test_app_config(3); - let (e2e_prover, dummy_internal_proof) = load_agg_pk_into_e2e_prover(app_config); - - let air_id_perm = e2e_prover.agg_pk().root_verifier_pk.air_id_permutation(); - let special_air_ids = air_id_perm.get_special_air_ids(); - - let root_proof = e2e_prover.generate_e2e_proof(StdIn::default()); - let root_pvs = RootVmVerifierPvs::from_flatten( - root_proof.per_air[special_air_ids.public_values_air_id] - .public_values - .clone(), - ); - - let app_exe_commit = AppExecutionCommit::compute( - &e2e_prover.app_pk.app_vm_pk.vm_config, - &e2e_prover.app_committed_exe, - e2e_prover.leaf_committed_exe(), - ); - - assert_eq!(root_pvs.exe_commit, app_exe_commit.exe_commit); - assert_eq!( - root_pvs.leaf_verifier_commit, - app_exe_commit.leaf_vm_verifier_commit - ); - - static_verifier::test_static_verifier( - &e2e_prover.agg_pk().root_verifier_pk, - dummy_internal_proof, - &root_proof, - ); -} - #[test] fn test_e2e_app_log_blowup_1() { - let app_config = small_test_app_config(1); - - let (e2e_prover, dummy_internal_proof) = load_agg_pk_into_e2e_prover(app_config); - let root_proof = e2e_prover.generate_e2e_proof(StdIn::default()); - - static_verifier::test_static_verifier( - &e2e_prover.agg_pk().root_verifier_pk, - dummy_internal_proof, - &root_proof, - ); -} - -#[test] -fn test_agg_keygen_store_and_load() { + let app_log_blowup = 1; const AGG_PK_PATH: &str = "temp/agg_pk.out"; + let app_config = small_test_app_config(app_log_blowup); + let app_pk = Sdk.app_keygen(app_config, None::<&str>).unwrap(); + let agg_pk = Sdk + .agg_keygen(full_agg_config_for_test(), Some(AGG_PK_PATH)) + .unwrap(); + let evm_verifier = Sdk + .generate_snark_verifier_contract(&agg_pk, None::<&'static str>) + .unwrap(); - let sdk = Sdk; - let agg_config = AggConfig { - max_num_user_public_values: NUM_PUB_VALUES, - leaf_fri_params: standard_fri_params_with_100_bits_conjectured_security(4), - internal_fri_params: standard_fri_params_with_100_bits_conjectured_security(3), - root_fri_params: standard_fri_params_with_100_bits_conjectured_security(2), - compiler_options: CompilerOptions { - enable_cycle_tracker: true, - compile_prints: true, - ..Default::default() - }, - }; - - let (_, agg_pk) = sdk.agg_keygen(agg_config, Some(AGG_PK_PATH)).unwrap(); - let (file_config, file_pk) = sdk.load_agg_pk_from_file(AGG_PK_PATH).unwrap(); - assert_agg_config_eq(&agg_config, &file_config); - assert_agg_pk_eq(&agg_pk, &file_pk); + let e2e_prover = Sdk + .create_e2e_prover(app_pk, app_committed_exe_for_test(app_log_blowup), agg_pk) + .unwrap(); + let evm_proof = e2e_prover.generate_proof_for_evm(StdIn::default()); + assert!(Sdk.verify_evm_proof(&evm_verifier, &evm_proof)); } #[test] @@ -324,71 +289,3 @@ fn test_sdk_guest_build_and_transpile() { .with_extension(Rv32MTranspilerExtension); let _exe = sdk.transpile(one, transpiler).unwrap(); } - -#[test] -fn test_sdk_vm_config_builder() { - let sdk_vm_config = SdkVmConfig::builder() - .system( - SystemConfig::default() - .with_max_segment_len(200) - .with_continuations() - .with_public_values(16), - ) - .native(Default::default()) - .rv32i(Default::default()) - .build(); - let app_config = AppConfig { - app_fri_params: standard_fri_params_with_100_bits_conjectured_security(1), - app_vm_config: sdk_vm_config, - }; - - let (e2e_prover, dummy_internal_proof) = load_agg_pk_into_e2e_prover(app_config); - let root_proof = e2e_prover.generate_e2e_proof(StdIn::default()); - - static_verifier::test_static_verifier( - &e2e_prover.agg_pk().root_verifier_pk, - dummy_internal_proof, - &root_proof, - ); -} - -mod static_verifier { - use ax_stark_sdk::{ - ax_stark_backend::prover::types::Proof, - config::baby_bear_poseidon2_outer::BabyBearPoseidon2OuterConfig, - }; - use axvm_native_compiler::prelude::Witness; - use axvm_native_recursion::{halo2::wrapper::Halo2WrapperProvingKey, witness::Witnessable}; - use axvm_sdk::keygen::RootVerifierProvingKey; - - use crate::SC; - - pub(crate) fn test_static_verifier( - root_verifier_pk: &RootVerifierProvingKey, - dummy_internal_proof: Proof, - root_proof: &Proof, - ) { - // Here we intend to use a dummy root proof to generate a static verifier circuit in order - // to test if the static verifier circuit can handle a different root proof. - let dummy_root_proof = root_verifier_pk.generate_dummy_root_proof(dummy_internal_proof); - let static_verifier = root_verifier_pk.keygen_static_verifier(24, dummy_root_proof); - let mut witness = Witness::default(); - root_proof.write(&mut witness); - // Here the proof is verified inside. - // FIXME: explicitly verify the proof. - let static_verifier_proof = static_verifier.prove(witness); - let verifier_wrapper = - Halo2WrapperProvingKey::keygen_auto_tune(static_verifier.generate_dummy_snark()); - assert_eq!( - verifier_wrapper - .pinning - .metadata - .config_params - .num_advice_per_phase, - vec![1] - ); - let evm_verifier = verifier_wrapper.generate_evm_verifier(); - let evm_proof = verifier_wrapper.prove_for_evm(static_verifier_proof); - Halo2WrapperProvingKey::evm_verify(evm_verifier, evm_proof); - } -} diff --git a/crates/axvm-sdk/tests/utils.rs b/crates/axvm-sdk/tests/utils.rs deleted file mode 100644 index 4bc271c97d..0000000000 --- a/crates/axvm-sdk/tests/utils.rs +++ /dev/null @@ -1,68 +0,0 @@ -use ax_stark_sdk::{ax_stark_backend::config::StarkGenericConfig, config::FriParameters}; -use axvm_circuit::{prover::types::VmProvingKey, system::program::trace::AxVmCommittedExe}; -use axvm_sdk::{config::AggConfig, keygen::AggProvingKey}; -use serde::Serialize; - -pub fn assert_agg_config_eq(a: &AggConfig, b: &AggConfig) { - assert_eq!(a.max_num_user_public_values, b.max_num_user_public_values); - assert_eq!(a.compiler_options.word_size, b.compiler_options.word_size); - assert_fri_params_eq(&a.leaf_fri_params, &b.leaf_fri_params); - assert_fri_params_eq(&a.internal_fri_params, &b.internal_fri_params); - assert_fri_params_eq(&a.root_fri_params, &b.root_fri_params); - assert_serialize_eq(a, b); -} - -pub fn assert_agg_pk_eq(a: &AggProvingKey, b: &AggProvingKey) { - assert_vm_pk_eq(&a.leaf_vm_pk, &b.leaf_vm_pk); - assert_vm_pk_eq(&a.internal_vm_pk, &b.internal_vm_pk); - assert_vm_pk_eq(&a.root_verifier_pk.vm_pk, &b.root_verifier_pk.vm_pk); - assert_committed_exe_eq(&a.internal_committed_exe, &b.internal_committed_exe); - assert_committed_exe_eq( - &a.root_verifier_pk.root_committed_exe, - &b.root_verifier_pk.root_committed_exe, - ); - assert_eq!( - a.root_verifier_pk.air_heights, - b.root_verifier_pk.air_heights, - ); - assert_serialize_eq(a, b); -} - -fn assert_fri_params_eq(a: &FriParameters, b: &FriParameters) { - assert_eq!(a.log_blowup, b.log_blowup); - assert_eq!(a.num_queries, b.num_queries); - assert_eq!(a.proof_of_work_bits, b.proof_of_work_bits); -} - -fn assert_vm_pk_eq(a: &VmProvingKey, b: &VmProvingKey) { - assert_fri_params_eq(&a.fri_params, &b.fri_params); - assert_eq!(a.vm_pk.max_constraint_degree, b.vm_pk.max_constraint_degree); - assert_eq!(a.vm_pk.per_air.len(), b.vm_pk.per_air.len()); - for (a_pk, b_pk) in a.vm_pk.per_air.iter().zip(b.vm_pk.per_air.iter()) { - assert_eq!(a_pk.air_name, b_pk.air_name); - } -} - -fn assert_committed_exe_eq( - a: &AxVmCommittedExe, - b: &AxVmCommittedExe, -) { - for (a_inst, b_inst) in a - .exe - .program - .instructions() - .iter() - .zip(b.exe.program.instructions().iter()) - { - assert_eq!(a_inst, b_inst); - } - assert_eq!(a.exe.pc_start, b.exe.pc_start); - assert_eq!(a.exe.init_memory, b.exe.init_memory); - assert_eq!(a.committed_program.raw_data, b.committed_program.raw_data); -} - -fn assert_serialize_eq(a: &T, b: &T) { - let a_bytes = bincode::serialize(a).unwrap(); - let b_bytes = bincode::serialize(b).unwrap(); - assert_eq!(a_bytes, b_bytes); -} diff --git a/crates/stark-backend/src/prover/quotient/single.rs b/crates/stark-backend/src/prover/quotient/single.rs index ffb5629d05..07c3a03004 100644 --- a/crates/stark-backend/src/prover/quotient/single.rs +++ b/crates/stark-backend/src/prover/quotient/single.rs @@ -21,7 +21,11 @@ use crate::{ /// /// Designed to be general enough to support RAP with multiple rounds of challenges. #[allow(clippy::too_many_arguments)] -#[instrument(name = "compute single RAP quotient polynomial", skip_all)] +#[instrument( + name = "compute single RAP quotient polynomial", + level = "trace", + skip_all +)] pub fn compute_single_rap_quotient_values<'a, SC, R, Mat>( rap: &'a R, symbolic_constraints: &SymbolicConstraints>, diff --git a/crates/stark-sdk/src/config/fri_params.rs b/crates/stark-sdk/src/config/fri_params.rs index f2bbffbfd9..0585e43ef7 100644 --- a/crates/stark-sdk/src/config/fri_params.rs +++ b/crates/stark-sdk/src/config/fri_params.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)] pub struct FriParameters { pub log_blowup: usize, pub num_queries: usize, diff --git a/crates/toolchain/tests/src/ecc_tests.rs b/crates/toolchain/tests/src/ecc_tests.rs index 1ba09ccc67..89de106227 100644 --- a/crates/toolchain/tests/src/ecc_tests.rs +++ b/crates/toolchain/tests/src/ecc_tests.rs @@ -34,6 +34,7 @@ use axvm_transpiler::{transpiler::Transpiler, FromElf}; use derive_more::derive::From; use eyre::Result; use num_bigint_dig::BigUint; +use serde::{Deserialize, Serialize}; use crate::utils::{build_example_program, build_example_program_with_features}; @@ -157,7 +158,7 @@ fn test_decompress() -> Result<()> { Ok(()) } -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32ModularKeccak256Config { #[system] pub system: SystemConfig, diff --git a/crates/toolchain/tests/tests/transpiler_tests.rs b/crates/toolchain/tests/tests/transpiler_tests.rs index bdfffce4e6..07dbf982bb 100644 --- a/crates/toolchain/tests/tests/transpiler_tests.rs +++ b/crates/toolchain/tests/tests/transpiler_tests.rs @@ -34,6 +34,7 @@ use axvm_transpiler::{elf::Elf, transpiler::Transpiler, FromElf}; use derive_more::derive::From; use eyre::Result; use num_bigint_dig::BigUint; +use serde::{Deserialize, Serialize}; use test_case::test_case; type F = BabyBear; @@ -93,7 +94,7 @@ fn test_rv32im_runtime(elf_path: &str) -> Result<()> { Ok(()) } -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32ModularFp2Int256Config { #[system] pub system: SystemConfig, diff --git a/crates/vm/src/arch/config.rs b/crates/vm/src/arch/config.rs index 39a6bad58e..f1d6973dd3 100644 --- a/crates/vm/src/arch/config.rs +++ b/crates/vm/src/arch/config.rs @@ -3,7 +3,7 @@ use ax_stark_backend::{p3_field::PrimeField32, ChipUsageGetter}; use axvm_circuit::system::memory::MemoryTraceHeights; use axvm_instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES; use derive_new::new; -use serde::{Deserialize, Serialize}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; use super::{ AnyEnum, InstructionExecutor, SystemComplex, SystemExecutor, SystemPeriphery, VmChipComplex, @@ -32,7 +32,7 @@ pub fn vm_poseidon2_config() -> Poseidon2Config::new_p3_baby_bear_16() } -pub trait VmConfig: Clone { +pub trait VmConfig: Clone + Serialize + DeserializeOwned { type Executor: InstructionExecutor + AnyEnum + ChipUsageGetter; type Periphery: AnyEnum + ChipUsageGetter; diff --git a/crates/vm/src/arch/segment.rs b/crates/vm/src/arch/segment.rs index 8a54ea5209..f5f1b0766f 100644 --- a/crates/vm/src/arch/segment.rs +++ b/crates/vm/src/arch/segment.rs @@ -279,10 +279,8 @@ impl> ExecutionSegment { itertools::izip!(self.air_names.clone(), self.current_trace_heights()).collect(); self.collected_metrics.emit(); - if did_terminate { - metrics::counter!("total_cells_used") - .absolute(self.current_trace_cells().into_iter().sum::() as u64); - } + metrics::counter!("total_cells_used") + .absolute(self.current_trace_cells().into_iter().sum::() as u64); } Ok(ExecutionSegmentState { diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs index 234589e48a..b3963e476e 100644 --- a/crates/vm/src/arch/vm.rs +++ b/crates/vm/src/arch/vm.rs @@ -127,8 +127,8 @@ where let mut pc = exe.pc_start; loop { - println!("Executing segment: {}", segments.len()); - let state = segment.execute_from_pc(pc)?; + let state = tracing::info_span!("execute_segment", segment = segments.len()) + .in_scope(|| segment.execute_from_pc(pc))?; pc = state.pc; if state.is_terminated { @@ -247,7 +247,7 @@ where let start = std::time::Instant::now(); let ret = seg.generate_proof_input(committed_program.clone()); #[cfg(feature = "bench-metrics")] - metrics::gauge!("execute_and_trace_gen_time_ms", "segment" => seg_idx.to_string()) + metrics::gauge!("trace_gen_time_ms", "segment" => seg_idx.to_string()) .set(start.elapsed().as_millis() as f64); ret }) diff --git a/crates/vm/tests/integration_test.rs b/crates/vm/tests/integration_test.rs index 653c5922f8..90c73779f5 100644 --- a/crates/vm/tests/integration_test.rs +++ b/crates/vm/tests/integration_test.rs @@ -53,6 +53,7 @@ use axvm_native_compiler::{ use axvm_rv32im_transpiler::BranchEqualOpcode::*; use derive_more::derive::From; use rand::Rng; +use serde::{Deserialize, Serialize}; use test_log::test; const LIMB_BITS: usize = 29; @@ -984,7 +985,7 @@ fn instructions_for_keccak256_test(input: &[u8]) -> Vec> { instructions } -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct NativeKeccakConfig { #[system] pub system: SystemConfig, diff --git a/extensions/algebra/circuit/Cargo.toml b/extensions/algebra/circuit/Cargo.toml index 1e13daacd2..8c442bb59d 100644 --- a/extensions/algebra/circuit/Cargo.toml +++ b/extensions/algebra/circuit/Cargo.toml @@ -18,15 +18,14 @@ axvm-rv32-adapters = { workspace = true } axvm-algebra-transpiler = { workspace = true } itertools = { workspace = true } -hex-literal = { workspace = true } -num-bigint-dig = { workspace = true } +num-bigint-dig = { workspace = true, features = ["serde"] } num-traits = { workspace = true } -once_cell = { workspace = true } tracing = { workspace = true } rand = { workspace = true } derive_more = { workspace = true, features = ["from"] } strum = { workspace = true } derive-new = { workspace = true } +serde.workspace = true [dev-dependencies] halo2curves-axiom = { workspace = true } diff --git a/extensions/algebra/circuit/src/config.rs b/extensions/algebra/circuit/src/config.rs index 80438f36cb..c01aa7086c 100644 --- a/extensions/algebra/circuit/src/config.rs +++ b/extensions/algebra/circuit/src/config.rs @@ -7,10 +7,11 @@ use axvm_circuit_derive::{AnyEnum, InstructionExecutor, VmConfig}; use axvm_rv32im_circuit::*; use derive_more::derive::From; use num_bigint_dig::BigUint; +use serde::{Deserialize, Serialize}; use super::*; -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32ModularConfig { #[system] pub system: SystemConfig, @@ -36,7 +37,7 @@ impl Rv32ModularConfig { } } -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32ModularWithFp2Config { #[system] pub system: SystemConfig, diff --git a/extensions/algebra/circuit/src/fp2_extension.rs b/extensions/algebra/circuit/src/fp2_extension.rs index c1547ded0b..c53fff7eb7 100644 --- a/extensions/algebra/circuit/src/fp2_extension.rs +++ b/extensions/algebra/circuit/src/fp2_extension.rs @@ -16,11 +16,12 @@ use axvm_instructions::{AxVmOpcode, UsizeOpcode}; use axvm_rv32_adapters::Rv32VecHeapAdapterChip; use derive_more::derive::From; use num_bigint_dig::BigUint; +use serde::{Deserialize, Serialize}; use strum::EnumCount; use crate::fp2_chip::{Fp2AddSubChip, Fp2MulDivChip}; -#[derive(Clone, Debug, derive_new::new)] +#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)] pub struct Fp2Extension { pub supported_modulus: Vec, } diff --git a/extensions/algebra/circuit/src/modular_extension.rs b/extensions/algebra/circuit/src/modular_extension.rs index f34ab2a2d2..d3e305f922 100644 --- a/extensions/algebra/circuit/src/modular_extension.rs +++ b/extensions/algebra/circuit/src/modular_extension.rs @@ -17,6 +17,7 @@ use axvm_instructions::{AxVmOpcode, UsizeOpcode}; use axvm_rv32_adapters::{Rv32IsEqualModAdapterChip, Rv32VecHeapAdapterChip}; use derive_more::derive::From; use num_bigint_dig::BigUint; +use serde::{Deserialize, Serialize}; use strum::EnumCount; use crate::modular_chip::{ @@ -24,7 +25,7 @@ use crate::modular_chip::{ ModularMulDivChip, ModularMulDivCoreChip, }; -#[derive(Clone, Debug, derive_new::new)] +#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)] pub struct ModularExtension { pub supported_modulus: Vec, } diff --git a/extensions/bigint/circuit/Cargo.toml b/extensions/bigint/circuit/Cargo.toml index 837306066e..b61181ecf4 100644 --- a/extensions/bigint/circuit/Cargo.toml +++ b/extensions/bigint/circuit/Cargo.toml @@ -18,13 +18,10 @@ axvm-rv32-adapters = { workspace = true } axvm-bigint-transpiler = { workspace = true } axvm-rv32im-transpiler = { workspace = true } -hex-literal.workspace = true -itertools.workspace = true -tracing.workspace = true derive-new.workspace = true derive_more = { workspace = true, features = ["from"] } rand.workspace = true -eyre.workspace = true +serde.workspace = true [dev-dependencies] ax-stark-sdk = { workspace = true } diff --git a/extensions/bigint/circuit/src/extension.rs b/extensions/bigint/circuit/src/extension.rs index e21088c730..2eea7ab1c6 100644 --- a/extensions/bigint/circuit/src/extension.rs +++ b/extensions/bigint/circuit/src/extension.rs @@ -24,10 +24,11 @@ use axvm_rv32im_circuit::{ Rv32MExecutor, Rv32MPeriphery, }; use derive_more::derive::From; +use serde::{Deserialize, Serialize}; use crate::*; -#[derive(Clone, Debug, VmConfig, derive_new::new)] +#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)] pub struct Int256Rv32Config { #[system] pub system: SystemConfig, @@ -53,7 +54,7 @@ impl Default for Int256Rv32Config { } } -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] pub struct Int256 { pub range_tuple_checker_sizes: [u32; 2], } diff --git a/extensions/ecc/circuit/Cargo.toml b/extensions/ecc/circuit/Cargo.toml index e54bc52ab6..c51d68e355 100644 --- a/extensions/ecc/circuit/Cargo.toml +++ b/extensions/ecc/circuit/Cargo.toml @@ -25,6 +25,7 @@ derive-new = { workspace = true } once_cell = { workspace = true } eyre = { workspace = true } num-integer = { workspace = true } +serde = { workspace = true } [dev-dependencies] ax-stark-sdk = { workspace = true } diff --git a/extensions/ecc/circuit/src/config.rs b/extensions/ecc/circuit/src/config.rs index ae8d03f38b..2672b88428 100644 --- a/extensions/ecc/circuit/src/config.rs +++ b/extensions/ecc/circuit/src/config.rs @@ -7,10 +7,11 @@ use axvm_circuit::arch::{ use axvm_circuit_derive::{AnyEnum, InstructionExecutor, VmConfig}; use axvm_rv32im_circuit::*; use derive_more::derive::From; +use serde::{Deserialize, Serialize}; use super::*; -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32WeierstrassConfig { #[system] pub system: SystemConfig, diff --git a/extensions/ecc/circuit/src/weierstrass_extension.rs b/extensions/ecc/circuit/src/weierstrass_extension.rs index ad845ca845..40f8f9739b 100644 --- a/extensions/ecc/circuit/src/weierstrass_extension.rs +++ b/extensions/ecc/circuit/src/weierstrass_extension.rs @@ -19,11 +19,12 @@ use derive_more::derive::From; use num_bigint_dig::BigUint; use num_traits::{FromPrimitive, Zero}; use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; use strum::EnumCount; use super::{EcAddNeChip, EcDoubleChip}; -#[derive(Clone, Debug, derive_new::new)] +#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)] pub struct CurveConfig { /// The coordinate modulus of the curve. pub modulus: BigUint, @@ -42,7 +43,7 @@ pub static SECP256K1_CONFIG: Lazy = Lazy::new(|| CurveConfig { b: BigUint::from_u8(7u8).unwrap(), }); -#[derive(Clone, Debug, derive_new::new)] +#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)] pub struct WeierstrassExtension { pub supported_curves: Vec, } diff --git a/extensions/keccak256/circuit/Cargo.toml b/extensions/keccak256/circuit/Cargo.toml index c8863ea69c..f044ff88d5 100644 --- a/extensions/keccak256/circuit/Cargo.toml +++ b/extensions/keccak256/circuit/Cargo.toml @@ -27,6 +27,7 @@ derive-new.workspace = true derive_more = { workspace = true, features = ["from"] } rand.workspace = true eyre.workspace = true +serde.workspace = true [dev-dependencies] ax-stark-sdk = { workspace = true } diff --git a/extensions/keccak256/circuit/src/extension.rs b/extensions/keccak256/circuit/src/extension.rs index a0505d2714..e09afd6b98 100644 --- a/extensions/keccak256/circuit/src/extension.rs +++ b/extensions/keccak256/circuit/src/extension.rs @@ -15,11 +15,12 @@ use axvm_rv32im_circuit::{ Rv32MExecutor, Rv32MPeriphery, }; use derive_more::derive::From; +use serde::{Deserialize, Serialize}; use strum::IntoEnumIterator; use crate::*; -#[derive(Clone, Debug, VmConfig, derive_new::new)] +#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)] pub struct Keccak256Rv32Config { #[system] pub system: SystemConfig, @@ -45,7 +46,7 @@ impl Default for Keccak256Rv32Config { } } -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)] pub struct Keccak256; #[derive(ChipUsageGetter, Chip, InstructionExecutor, From, AnyEnum)] diff --git a/extensions/native/recursion/Cargo.toml b/extensions/native/recursion/Cargo.toml index 764b618a9c..bf86545391 100644 --- a/extensions/native/recursion/Cargo.toml +++ b/extensions/native/recursion/Cargo.toml @@ -32,6 +32,7 @@ cfg-if = { workspace = true } [dev-dependencies] axvm-native-recursion = { workspace = true, features = ["test-utils"] } tempfile = "3.14.0" +bson.workspace = true [features] default = ["parallel", "mimalloc"] diff --git a/extensions/native/recursion/src/halo2/mod.rs b/extensions/native/recursion/src/halo2/mod.rs index 366baba9a6..e309a3953f 100644 --- a/extensions/native/recursion/src/halo2/mod.rs +++ b/extensions/native/recursion/src/halo2/mod.rs @@ -18,7 +18,8 @@ use itertools::Itertools; use serde::{ de, de::{MapAccess, Visitor}, - Deserialize, Deserializer, Serialize, + ser::SerializeStruct, + Deserialize, Deserializer, Serialize, Serializer, }; use snark_verifier_sdk::{ halo2::{gen_dummy_snark_from_vk, gen_snark_shplonk}, @@ -60,9 +61,8 @@ pub struct DslOperations { /// Necessary metadata to prove a Halo2 circuit /// Attention: Deserializer of this struct is not generic. It only works for verifier/wrapper circuit. -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone)] pub struct Halo2ProvingPinning { - #[serde(serialize_with = "pk_serializer")] pub pk: ProvingKey, pub metadata: Halo2ProvingMetadata, } @@ -268,11 +268,19 @@ impl Halo2Prover { } } -fn pk_serializer(value: &ProvingKey, serializer: S) -> Result -where - S: serde::Serializer, -{ - serializer.serialize_bytes(&value.to_bytes(SerdeFormat::RawBytes)) +impl Serialize for Halo2ProvingPinning { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let pk_bytes = self.pk.to_bytes(SerdeFormat::RawBytes); + + // Start serializing as a struct with 2 fields: "pk" and "metadata" + let mut state = serializer.serialize_struct("Halo2ProvingPinning", 2)?; + state.serialize_field("pk", &pk_bytes)?; + state.serialize_field("metadata", &self.metadata)?; + state.end() + } } impl<'de> Deserialize<'de> for Halo2ProvingPinning { diff --git a/extensions/native/recursion/src/halo2/tests/mod.rs b/extensions/native/recursion/src/halo2/tests/mod.rs index 5ebb333860..9469822b99 100644 --- a/extensions/native/recursion/src/halo2/tests/mod.rs +++ b/extensions/native/recursion/src/halo2/tests/mod.rs @@ -1,3 +1,5 @@ +use std::io::Write; + use ax_stark_backend::p3_field::{ reduce_32 as reduce_32_gt, split_32 as split_32_gt, AbstractField, }; @@ -171,9 +173,16 @@ fn test_wrapper_select_k() { #[test] fn test_pinning_serde() { let (_, _, pinning) = snarks_dummy_circuit(); + // Something went wrong when Halo2ProvingPinning is a field. So we explicitly test with a struct + // contains Haplo2ProvingPinning. + let wrapper = Halo2WrapperProvingKey { + pinning: pinning.clone(), + }; + let mut f = tempfile::NamedTempFile::new().unwrap(); - serde_json::to_writer_pretty(&mut f, &pinning).unwrap(); - let new_pinning: Halo2ProvingPinning = serde_json::from_reader(f.reopen().unwrap()).unwrap(); + f.write_all(&bson::to_vec(&wrapper).unwrap()).unwrap(); + let new_wrapper: Halo2WrapperProvingKey = bson::from_reader(f.reopen().unwrap()).unwrap(); + let new_pinning = new_wrapper.pinning; let params = gen_kzg_params(DUMMY_K as u32); let mut builder = BaseCircuitBuilder::from_stage(Prover) .use_params(new_pinning.metadata.config_params.clone()) diff --git a/extensions/native/recursion/src/halo2/wrapper.rs b/extensions/native/recursion/src/halo2/wrapper.rs index 2ac41e4c05..b1110492aa 100644 --- a/extensions/native/recursion/src/halo2/wrapper.rs +++ b/extensions/native/recursion/src/halo2/wrapper.rs @@ -22,6 +22,9 @@ use crate::halo2::{ EvmProof, Halo2Params, Halo2ProvingMetadata, Halo2ProvingPinning, }; +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EvmVerifier(pub Vec); + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Halo2WrapperProvingKey { pub pinning: Halo2ProvingPinning, @@ -66,22 +69,22 @@ impl Halo2WrapperProvingKey { } /// A helper function for testing to verify the proof of this circuit with evm verifier. // FIXME: the signature is not ideal. It should return an Error instead of panicking when the verification fails. - pub fn evm_verify(evm_verifier_deployment_codes: Vec, evm_proof: EvmProof) { + pub fn evm_verify(evm_verifier: &EvmVerifier, evm_proof: &EvmProof) { evm_verify( - evm_verifier_deployment_codes, - evm_proof.instances, - evm_proof.proof, + evm_verifier.0.clone(), + evm_proof.instances.clone(), + evm_proof.proof.clone(), ); } /// Return deployment code for EVM verifier which can verify the snark of this circuit. - pub fn generate_evm_verifier(&self) -> Vec { + pub fn generate_evm_verifier(&self) -> EvmVerifier { let params = read_params(self.pinning.metadata.config_params.k as u32); - gen_evm_verifier_shplonk::( + EvmVerifier(gen_evm_verifier_shplonk::( ¶ms, self.pinning.pk.get_vk(), self.pinning.metadata.num_pvs.clone(), None, - ) + )) } pub fn prove_for_evm(&self, snark_to_verify: Snark) -> EvmProof { let k = self.pinning.metadata.config_params.k; diff --git a/extensions/pairing/circuit/Cargo.toml b/extensions/pairing/circuit/Cargo.toml index 21ec3b5f33..d01c077523 100644 --- a/extensions/pairing/circuit/Cargo.toml +++ b/extensions/pairing/circuit/Cargo.toml @@ -33,6 +33,7 @@ once_cell = { workspace = true } rand = { workspace = true } itertools = { workspace = true } eyre = { workspace = true } +serde = { workspace = true, features = ["derive", "std"] } [target.'cfg(not(target_os = "zkvm"))'.dependencies] axvm-pairing-guest = { workspace = true } diff --git a/extensions/pairing/circuit/src/config.rs b/extensions/pairing/circuit/src/config.rs index e3e3ca4bdd..7a4e5c2b55 100644 --- a/extensions/pairing/circuit/src/config.rs +++ b/extensions/pairing/circuit/src/config.rs @@ -8,10 +8,11 @@ use axvm_circuit_derive::{AnyEnum, InstructionExecutor, VmConfig}; use axvm_ecc_circuit::*; use axvm_rv32im_circuit::*; use derive_more::derive::From; +use serde::{Deserialize, Serialize}; use super::*; -#[derive(Clone, Debug, VmConfig)] +#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)] pub struct Rv32PairingConfig { #[system] pub system: SystemConfig, diff --git a/extensions/pairing/circuit/src/pairing_extension.rs b/extensions/pairing/circuit/src/pairing_extension.rs index 595705a3b5..20c31635c6 100644 --- a/extensions/pairing/circuit/src/pairing_extension.rs +++ b/extensions/pairing/circuit/src/pairing_extension.rs @@ -22,12 +22,13 @@ use axvm_rv32_adapters::{Rv32VecHeapAdapterChip, Rv32VecHeapTwoReadsAdapterChip} use derive_more::derive::From; use num_bigint_dig::BigUint; use num_traits::{FromPrimitive, Zero}; +use serde::{Deserialize, Serialize}; use strum::{EnumCount, FromRepr}; use super::*; // All the supported pairing curves. -#[derive(Clone, Copy, Debug, FromRepr)] +#[derive(Clone, Copy, Debug, FromRepr, Serialize, Deserialize)] #[repr(usize)] pub enum PairingCurve { Bn254, @@ -60,7 +61,7 @@ impl PairingCurve { } } -#[derive(Clone, Debug, derive_new::new)] +#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)] pub struct PairingExtension { pub supported_curves: Vec, } diff --git a/extensions/rv32im/circuit/Cargo.toml b/extensions/rv32im/circuit/Cargo.toml index c080fd5f1b..702e51d05a 100644 --- a/extensions/rv32im/circuit/Cargo.toml +++ b/extensions/rv32im/circuit/Cargo.toml @@ -26,6 +26,7 @@ eyre.workspace = true # for div_rem: num-bigint.workspace = true num-integer.workspace = true +serde = { workspace = true, features = ["derive", "std"] } [dev-dependencies] ax-stark-sdk = { workspace = true } diff --git a/extensions/rv32im/circuit/src/extension.rs b/extensions/rv32im/circuit/src/extension.rs index 10686145b7..c0589eca68 100644 --- a/extensions/rv32im/circuit/src/extension.rs +++ b/extensions/rv32im/circuit/src/extension.rs @@ -21,12 +21,13 @@ use axvm_rv32im_transpiler::{ Rv32LoadStoreOpcode, Rv32Phantom, ShiftOpcode, }; use derive_more::derive::From; +use serde::{Deserialize, Serialize}; use strum::IntoEnumIterator; use crate::{adapters::*, *}; /// Config for a VM with base extension and IO extension -#[derive(Clone, Debug, VmConfig, derive_new::new)] +#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)] pub struct Rv32IConfig { #[system] pub system: SystemConfig, @@ -37,7 +38,7 @@ pub struct Rv32IConfig { } /// Config for a VM with base extension, IO extension, and multiplication extension -#[derive(Clone, Debug, VmConfig, derive_new::new)] +#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)] pub struct Rv32ImConfig { #[system] pub system: SystemConfig, @@ -122,15 +123,15 @@ impl Rv32ImConfig { // ============ Extension Implementations ============ /// RISC-V 32-bit Base (RV32I) Extension -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)] pub struct Rv32I; /// RISC-V Extension for handling IO (not to be confused with I base extension) -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)] pub struct Rv32Io; /// RISC-V 32-bit Multiplication Extension (RV32M) Extension -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] pub struct Rv32M { pub range_tuple_checker_sizes: [u32; 2], }