Add prove_stdio.rs

scripts/Cargo.toml

@@ -16,6 +16,7 @@ clap = { workspace = true, features = ["derive"] }
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
 sysinfo = "0.32.0"
+num_cpus = "1.0"
 
 [lints]
 workspace = true

scripts/prove_stdio.rs

@@ -1,6 +1,5 @@
-use std::{env::set_var, fmt::Display, fs::create_dir_all, path::PathBuf, process::Command};
+use std::{env::set_var, fs::File, path::PathBuf, process::Command};
 
-use alloy::{eips::BlockId, transports::http::reqwest::Url};
 use anyhow::{ensure, Ok};
 use clap::{arg, Args, ValueEnum, ValueHint};
 
@@ -23,32 +22,121 @@ pub struct ProveStdioArgs {
     /// The end of the block range to prove. If None, start_block-1 is used.
     #[arg(long, default_value_t = false)]
     use_test_config: bool,
+    /// The batch size for block fetching.
+    #[arg(long, default_value_t = 8)]
+    block_batch_size: u32,
+    /// The directory to output the proof files. If it does not exist, it will
+    /// recursively be created.
+    #[arg(short = 'o', long, value_hint = ValueHint::DirPath, default_value = ".")]
+    output_dir: PathBuf,
 }
 
 pub fn prove_via_stdio(args: ProveStdioArgs) -> anyhow::Result<()> {
-    // Set rustc environment variables.
-    set_var("RUST_MIN_STACK", "33554432");
-    set_var("RUST_BACKTRACE", "full");
-    set_var("RUST_LOG", "info");
-    // Script users are running locally, and might benefit from extra perf.
-    // See also .cargo/config.toml.
-    set_var("RUSTFLAGS", "-C target-cpu=native -Zlinker-features=-lld");
+    // Get number of cores of the system.
+    let num_cpus = num_cpus::get().to_string();
+    let mut envs = vec![
+        ("RUST_MIN_STACK", "33554432"),
+        ("RUST_BACKTRACE", "full"),
+        ("RUST_LOG", "info"),
+        ("RUSTFLAGS", "-C target-cpu=native -Zlinker-features=-lld"),
+        ("RAYON_NUM_THREADS", num_cpus.as_str()),
+        ("TOKIO_WORKER_THREADS", num_cpus.as_str()),
+    ];
 
     match args.mode {
         RunMode::Test => {
+            envs.extend([
+                ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+                ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+                ("CPU_CIRCUIT_SIZE", "9..20"),
+                ("KECCAK_CIRCUIT_SIZE", "7..18"),
+                ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
+                ("LOGIC_CIRCUIT_SIZE", "5..17"),
+                ("MEMORY_CIRCUIT_SIZE", "17..22"),
+                ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
+                ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
+                // TODO(Robin): update Poseidon ranges here and below once Kernel ASM supports
+                ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+            ]);
             let witness_filename = args
                 .input_witness_file
                 .to_str()
                 .ok_or(anyhow::anyhow!("Invalid witness file path"))?;
             if witness_filename.contains("witness_b19807080") {
+                envs.extend([
+                    ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+                    ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+                    ("CPU_CIRCUIT_SIZE", "9..20"),
+                    ("KECCAK_CIRCUIT_SIZE", "7..18"),
+                    ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..14"),
+                    ("LOGIC_CIRCUIT_SIZE", "5..17"),
+                    ("MEMORY_CIRCUIT_SIZE", "17..22"),
+                    ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..20"),
+                    ("MEMORY_AFTER_CIRCUIT_SIZE", "7..20"),
+                    ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+                ]);
             } else if witness_filename.contains("witness_b3_b6") {
-            } else {
+                envs.extend([
+                    ("ARITHMETIC_CIRCUIT_SIZE", "16..18"),
+                    ("BYTE_PACKING_CIRCUIT_SIZE", "8..15"),
+                    ("CPU_CIRCUIT_SIZE", "10..20"),
+                    ("KECCAK_CIRCUIT_SIZE", "4..13"),
+                    ("KECCAK_SPONGE_CIRCUIT_SIZE", "8..9"),
+                    ("LOGIC_CIRCUIT_SIZE", "4..14"),
+                    ("MEMORY_CIRCUIT_SIZE", "17..22"),
+                    ("MEMORY_BEFORE_CIRCUIT_SIZE", "16..18"),
+                    ("MEMORY_AFTER_CIRCUIT_SIZE", "7..8"),
+                    ("POSEIDON_CIRCUIT_SIZE", "4..8"),
+                ]);
             }
-            todo!("Test mode");
+
+            run_proof(args, envs)
         }
         RunMode::Verify => {
-            todo!("Verify mode");
+            // Build the targets before timing.
+            let status = Command::new("cargo")
+                .envs(envs.clone())
+                .args(["build", "--release", "--jobs", num_cpus.as_str()])
+                .spawn()?
+                .wait()?;
+            ensure!(status.success(), "command failed with {}", status);
+
+            // Time the proof.
+            let start = std::time::Instant::now();
+            run_proof(args, envs)?;
+            let elapsed = start.elapsed();
+            println!("Elapsed: {:?}", elapsed);
+            Ok(())
         }
     }
+}
+
+fn run_proof(args: ProveStdioArgs, envs: Vec<(&str, &str)>) -> anyhow::Result<()> {
+    let witness_file = File::open(&args.input_witness_file)?;
+    let mut cmd = Command::new("cargo");
+    cmd.envs(envs).stdin(witness_file);
+    cmd.args([
+        "run",
+        "--release",
+        "--package",
+        "zero",
+        "--bin",
+        "leader",
+        "--",
+        "--runtime",
+        "in-memory",
+        "--load-strategy",
+        "on-demand",
+        "--block-batch-size",
+        args.block_batch_size.to_string().as_str(),
+        "--proof-output-dir",
+        args.output_dir.to_str().unwrap(),
+        "stdio",
+    ]);
+    if args.use_test_config {
+        cmd.arg("--use-test-config");
+    }
+    let status = cmd.spawn()?.wait()?;
+    ensure!(status.success(), "command failed with {}", status);
     Ok(())
 }