feat: use criterion for benchmarks

provekit/Cargo.toml

@@ -6,8 +6,9 @@ edition = "2024"
 [dependencies]
 rand = { workspace = true }
 utils = { workspace = true }
+criterion = { workspace = true }
 noir-r1cs = { git = "https://github.com/worldfnd/ProveKit", branch = "main" }
 
 [[bench]]
-name = "provekit"
+name = "prove_verify"
 harness = false

provekit/README.md

@@ -0,0 +1,16 @@
+# ProveKit SHA256 Benchmark
+
+This benchmark code is using the ProveKit of World Foundation(https://github.com/worldfnd/ProveKit).
+
+
+## Benchmarking
+
+```bash
+cargo bench
+```
+
+Measure RAM footprint:
+
+```bash
+cargo run -r --bin provekit-measure
+```

provekit/benches/prove_verify.rs

@@ -0,0 +1,68 @@
+use criterion::{BatchSize, Criterion, black_box, criterion_group, criterion_main};
+use provekit::{ProvekitSha256Benchmark, WORKSPACE_ROOT};
+use std::path::PathBuf;
+use utils::bench::{SubMetrics, SubMetricsTable, measure_peak_memory};
+
+const INPUT_EXPONENTS: [u32; 1] = [11];
+
+fn sha256_benchmarks(c: &mut Criterion) {
+    let (bench_harness, preprocessing_peak_memory) =
+        measure_peak_memory(|| ProvekitSha256Benchmark::new(&INPUT_EXPONENTS));
+
+    let mut all_metrics = Vec::new();
+
+    for &exp in INPUT_EXPONENTS.iter() {
+        let mut metrics = SubMetrics::new(1 << exp);
+        metrics.preprocessing_peak_memory = preprocessing_peak_memory;
+
+        let package_name = format!("sha256_bench_2e{exp}");
+        let circuit_path = PathBuf::from(WORKSPACE_ROOT)
+            .join("target")
+            .join(format!("{package_name}.json"));
+        let toml_path = PathBuf::from(WORKSPACE_ROOT)
+            .join("circuits/hash/sha256-provekit")
+            .join(format!("sha256-bench-2e{exp}"))
+            .join("Prover.toml");
+
+        metrics.preprocessing_size = std::fs::metadata(circuit_path)
+            .map(|m| m.len())
+            .unwrap_or(0) as usize
+            + std::fs::metadata(toml_path).map(|m| m.len()).unwrap_or(0) as usize;
+
+        let (proof, proving_peak_memory) = measure_peak_memory(|| bench_harness.run_prove(exp));
+        metrics.proving_peak_memory = proving_peak_memory;
+        metrics.proof_size = proof.whir_r1cs_proof.transcript.len();
+
+        all_metrics.push(metrics);
+    }
+
+    println!("{}", SubMetricsTable(all_metrics));
+
+    let mut group = c.benchmark_group("SHA256 Prove & Verify");
+    group.sample_size(10);
+
+    for &exp in INPUT_EXPONENTS.iter() {
+        let input_size = 1 << exp;
+        let prove_id = format!("Prove ({} bytes)", input_size);
+        group.bench_function(prove_id, |bench| {
+            bench.iter(|| {
+                let proof = bench_harness.run_prove(exp);
+                black_box(proof);
+            });
+        });
+
+        let verify_id = format!("Verify ({} bytes)", input_size);
+        group.bench_function(verify_id, |bench| {
+            bench.iter_batched(
+                || bench_harness.prepare_verify(exp),
+                |(proof, proof_scheme)| bench_harness.run_verify(&proof, proof_scheme).unwrap(),
+                BatchSize::SmallInput,
+            );
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, sha256_benchmarks);
+criterion_main!(benches);

provekit/src/lib.rs

@@ -1,138 +1,107 @@
-use noir_r1cs::NoirProofScheme;
+use noir_r1cs::{NoirProof, NoirProofScheme};
 use rand::RngCore;
-use std::fs::{self, File};
-use std::io::{Read, Write};
+use std::collections::HashMap;
+use std::fs;
+use std::path::PathBuf;
 use std::process::Command;
-use std::time::Instant;
-use utils::bench::{Metrics, benchmark};
 
-pub const INPUT_EXP: [u32; 5] = [8, 9, 10, 11, 12];
-pub const TMP_DIR: &str = "tmp";
-pub const CIRCUIT_ROOT: &str = "circuits/hash/sha256-provekit";
-pub const CSV_OUTPUT: &str = "tmp/provekit_sha256.csv";
+pub const WORKSPACE_ROOT: &str = "circuits";
+pub const CIRCUIT_SUB_PATH: &str = "hash/sha256-provekit";
 
-/// Generates random input files for hashing benchmarks.
-pub fn generate_hash_inputs() -> Result<(), &'static str> {
-    let input_dir = format!("{}/hash-input", TMP_DIR);
-    fs::create_dir_all(&input_dir).map_err(|_| "Failed to create input directory")?;
-
-    let mut rng = rand::thread_rng();
-
-    for exp in INPUT_EXP {
-        let size = 1usize << exp;
-        let bin_path = format!("{}/input_2e{}.bin", input_dir, exp);
-
-        let mut data = vec![0u8; size];
-        rng.fill_bytes(&mut data);
-
-        let mut file = File::create(&bin_path).map_err(|_| "Failed to create input file")?;
-        file.write_all(&data)
-            .map_err(|_| "Failed to write input file")?;
-    }
-
-    Ok(())
+/// Provekit benchmark harness for SHA256.
+pub struct ProvekitSha256Benchmark {
+    proof_schemes: HashMap<u32, NoirProofScheme>,
+    toml_paths: HashMap<u32, PathBuf>,
 }
 
-/// Compiles all Noir circuits.
-pub fn compile_all_circuits() -> Result<(), &'static str> {
-    let output = Command::new("nargo")
-        .args([
-            "compile",
-            "--silence-warnings",
-            "--skip-brillig-constraints-check",
-        ])
-        .current_dir("circuits")
-        .output()
-        .map_err(|_| "Failed to execute nargo")?;
-
-    if !output.status.success() {
-        eprintln!("{}", String::from_utf8_lossy(&output.stderr));
-        return Err("Compilation failed");
+impl ProvekitSha256Benchmark {
+    /// Compiles the circuits and creates a new benchmark harness.
+    pub fn new(exponents: &[u32]) -> Self {
+        let output = Command::new("nargo")
+            .args([
+                "compile",
+                "--workspace",
+                "--silence-warnings",
+                "--skip-brillig-constraints-check",
+            ])
+            .current_dir(WORKSPACE_ROOT)
+            .output()
+            .expect("Failed to run nargo compile");
+
+        if !output.status.success() {
+            panic!(
+                "Workspace compilation failed: {}",
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+
+        let mut rng = rand::thread_rng();
+        let workspace_path = PathBuf::from(WORKSPACE_ROOT);
+        let mut proof_schemes = HashMap::new();
+        let mut toml_paths = HashMap::new();
+
+        for &exp in exponents {
+            let size = 1usize << exp;
+            let package_name = format!("sha256_bench_2e{exp}");
+            let circuit_path = workspace_path
+                .join("target")
+                .join(format!("{package_name}.json"));
+
+            let proof_scheme = NoirProofScheme::from_file(circuit_path.to_str().unwrap())
+                .unwrap_or_else(|e| panic!("Failed to load proof scheme for exp {exp}: {e}"));
+            proof_schemes.insert(exp, proof_scheme);
+
+            let dir_name = format!("sha256-bench-2e{exp}");
+            let circuit_member_dir = workspace_path.join(CIRCUIT_SUB_PATH).join(dir_name);
+            fs::create_dir_all(&circuit_member_dir).expect("Failed to create circuit dir");
+
+            let mut data = vec![0u8; size];
+            rng.fill_bytes(&mut data);
+            let toml_content = format!(
+                "input = [{}]\ninput_len = {size}",
+                data.iter()
+                    .map(u8::to_string)
+                    .collect::<Vec<_>>()
+                    .join(", "),
+            );
+
+            let toml_path = circuit_member_dir.join("Prover.toml");
+            fs::write(&toml_path, toml_content).expect("Failed to write Prover.toml");
+            toml_paths.insert(exp, toml_path);
+        }
+
+        Self {
+            proof_schemes,
+            toml_paths,
+        }
     }
 
-    Ok(())
-}
-
-/// Generates a Prover.toml file from input data.
-pub fn generate_prover_toml(
-    toml_path: &str,
-    input_file: &str,
-    size: usize,
-) -> Result<(), &'static str> {
-    let mut file = File::open(input_file).map_err(|_| "Failed to open input file")?;
-    let mut buffer = Vec::new();
-    file.read_to_end(&mut buffer)
-        .map_err(|_| "Failed to read input file")?;
-
-    let content = format!(
-        "input = [\n{}\n]\ninput_len = {}\n",
-        buffer
-            .iter()
-            .map(|b| format!("    {}", b))
-            .collect::<Vec<_>>()
-            .join(",\n"),
-        size
-    );
-
-    fs::write(toml_path, content).map_err(|_| "Failed to write Prover.toml")?;
-
-    Ok(())
-}
-
-/// Sets up Prover.toml files for all circuits.
-pub fn setup_circuits() -> Result<(), &'static str> {
-    for exp in INPUT_EXP {
-        let size = 1usize << exp;
-        let input_file = format!("{}/hash-input/input_2e{}.bin", TMP_DIR, exp);
-        let circuit_dir = format!("{}/sha256-bench-2e{}", CIRCUIT_ROOT, exp);
-        let toml_path = format!("{}/Prover.toml", circuit_dir);
-
-        generate_prover_toml(&toml_path, &input_file, size)?;
+    /// Runs the proving algorithm.
+    pub fn run_prove(&self, input_exp: u32) -> NoirProof {
+        let proof_scheme = self.proof_schemes.get(&input_exp).unwrap();
+        let toml_path = self.toml_paths.get(&input_exp).unwrap();
+        let witness_map = proof_scheme
+            .read_witness(toml_path.to_str().unwrap())
+            .expect("Failed to read witness");
+
+        proof_scheme
+            .prove(&witness_map)
+            .expect("Proof generation failed")
     }
 
-    Ok(())
-}
-
-/// Sets up the benchmark.
-pub fn setup() -> Result<(), &'static str> {
-    fs::create_dir_all(TMP_DIR).map_err(|_| "Failed to create tmp directory")?;
-
-    generate_hash_inputs()?;
-    compile_all_circuits()?;
-    setup_circuits()?;
-
-    Ok(())
-}
-
-/// Benchmarks provekit with sha256 for all input exponents.
-pub fn bench_sha256() {
-    let inputs: Vec<u32> = INPUT_EXP.to_vec();
-
-    benchmark(
-        |input_exp| {
-            let size = 1usize << input_exp;
-            let mut metrics = Metrics::new(size);
-
-            let circuit_dir = format!("{}/sha256-bench-2e{}", CIRCUIT_ROOT, input_exp);
-            let circuit_path = format!("circuits/target/sha256_bench_2e{}.json", input_exp);
-            let prover_toml_path = format!("{}/Prover.toml", circuit_dir);
-
-            let proof_scheme = NoirProofScheme::from_file(&circuit_path).unwrap();
-            let input_map = proof_scheme.read_witness(&prover_toml_path).unwrap();
-
-            let prove_start = Instant::now();
-            let proof = proof_scheme.prove(&input_map).unwrap();
-            metrics.proof_duration = prove_start.elapsed();
-
-            let verify_start = Instant::now();
-            proof_scheme.verify(&proof).unwrap();
-            metrics.verify_duration = verify_start.elapsed();
-
-            metrics.proof_size = proof.whir_r1cs_proof.transcript.len();
+    /// Prepares inputs for verification.
+    pub fn prepare_verify(&self, input_exp: u32) -> (NoirProof, &NoirProofScheme) {
+        let proof_scheme = self.proof_schemes.get(&input_exp).unwrap();
+        let proof = self.run_prove(input_exp);
+        (proof, proof_scheme)
+    }
 
-            metrics
-        },
-        &inputs,
-        CSV_OUTPUT,
-    );
+    /// Runs the verification algorithm.
+    pub fn run_verify(
+        &self,
+        proof: &NoirProof,
+        proof_scheme: &NoirProofScheme,
+    ) -> Result<(), &'static str> {
+        proof_scheme.verify(proof).map_err(|_| "Proof is not valid")
+    }
 }

utils/src/bench.rs

@@ -157,6 +157,18 @@ impl SubMetrics {
     }
 }
 
+pub struct SubMetricsTable(pub Vec<SubMetrics>);
+
+impl Display for SubMetricsTable {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.0.is_empty() {
+            return Ok(());
+        }
+        let mut table = Table::new(&self.0);
+        table.with(Style::modern());
+        write!(f, "{table}")
+    }
+}
 pub fn write_json_submetrics(output_path: &str, metrics: &SubMetrics) {
     let json = serde_json::to_string_pretty(metrics).unwrap();
     std::fs::write(output_path, json).unwrap();