feat: Refactor run_benchmark_ere

crates/benchmark-runner/src/lib.rs

@@ -47,64 +47,76 @@ pub enum Action {
 
 pub fn run_benchmark_ere<V>(host_name: &str, zkvm_instance: V, action: Action) -> Result<()>
 where
-    // C: Compiler + Send + Sync,
-    // C::Error: std::error::Error + Send + Sync + 'static,
     V: zkVM + Sync,
-    // V::Error: std::error::Error + Send + Sync + 'static,
 {
     println!("Benchmarking `{}`…", host_name);
-
     let zkvm_ref = Arc::new(&zkvm_instance);
-
     let corpuses = generate_stateless_witness::generate();
 
-    // Iterate through test corpus and generate reports
-    // TODO: note that when proving, processing these in parallel will likely skew the results
-    corpuses.into_par_iter().try_for_each(|bw| -> Result<()> {
-        println!(" {} ({} blocks)", bw.name, bw.blocks_and_witnesses.len());
-
-        let mut reports = Vec::new();
-        for ci in bw.blocks_and_witnesses {
-            let block_number = ci.block.number;
+    match action {
+        Action::Execute => {
+            // Use parallel iteration for execution
+            corpuses.into_par_iter().try_for_each(|bw| -> Result<()> {
+                process_corpus(bw, zkvm_ref.clone(), &action, host_name)
+            })?;
+        }
+        Action::Prove => {
+            // Use sequential iteration for proving
+            corpuses.into_iter().try_for_each(|bw| -> Result<()> {
+                process_corpus(bw, Arc::new(&*zkvm_ref), &action, host_name)
+            })?;
+        }
+    }
 
-            let mut stdin = Input::new();
-            stdin.write(ci);
-            stdin.write(bw.network);
+    Ok(())
+}
 
-            let workload_metrics = match action {
-                Action::Execute => {
-                    let report = zkvm_ref.execute(&stdin)?;
-                    let region_cycles: HashMap<_, _> = report.region_cycles.into_iter().collect();
-                    WorkloadMetrics::Execution {
-                        name: format!("{}-{}", bw.name, block_number),
-                        total_num_cycles: report.total_num_cycles,
-                        region_cycles,
-                    }
+fn process_corpus<V>(
+    bw: BlocksAndWitnesses,
+    zkvm_ref: Arc<&V>,
+    action: &Action,
+    host_name: &str,
+) -> Result<()>
+where
+    V: zkVM + Sync,
+{
+    println!(" {} ({} blocks)", bw.name, bw.blocks_and_witnesses.len());
+    let mut reports = Vec::new();
+
+    for ci in bw.blocks_and_witnesses {
+        let block_number = ci.block.number;
+        let mut stdin = Input::new();
+        stdin.write(ci);
+        stdin.write(bw.network);
+
+        let workload_metrics = match action {
+            Action::Execute => {
+                let report = zkvm_ref.execute(&stdin)?;
+                let region_cycles: HashMap<_, _> = report.region_cycles.into_iter().collect();
+                WorkloadMetrics::Execution {
+                    name: format!("{}-{}", bw.name, block_number),
+                    total_num_cycles: report.total_num_cycles,
+                    region_cycles,
                 }
-                Action::Prove => {
-                    let (proof, report) = zkvm_ref.prove(&stdin)?;
-
-                    WorkloadMetrics::Proving {
-                        name: format!("{}-{}", bw.name, block_number),
-                        proving_time_ms: report.proving_time.as_millis(),
-                    }
+            }
+            Action::Prove => {
+                let (_, report) = zkvm_ref.prove(&stdin)?;
+                WorkloadMetrics::Proving {
+                    name: format!("{}-{}", bw.name, block_number),
+                    proving_time_ms: report.proving_time.as_millis(),
                 }
-            };
-
-            reports.push(workload_metrics);
-        }
-
-        let out_path = format!(
-            "{}/zkevm-metrics/{}/{}.json",
-            env!("CARGO_WORKSPACE_DIR"),
-            host_name,
-            bw.name
-        );
-        WorkloadMetrics::to_path(out_path, &reports)?;
-        println!("wrote {} reports", reports.len());
-
-        Ok(())
-    })?;
-
+            }
+        };
+        reports.push(workload_metrics);
+    }
+
+    let out_path = format!(
+        "{}/zkevm-metrics/{}/{}.json",
+        env!("CARGO_WORKSPACE_DIR"),
+        host_name,
+        bw.name
+    );
+    WorkloadMetrics::to_path(out_path, &reports)?;
+    println!("wrote {} reports", reports.len());
     Ok(())
 }

crates/metrics/src/lib.rs

@@ -9,6 +9,7 @@ use thiserror::Error;
 /// Stores the total cycle count and a breakdown of cycle count per named region.
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub enum WorkloadMetrics {
+    /// Metrics produced when benchmarking in execution mode
     Execution {
         /// Name of the workload (e.g., "fft", "aes").
         name: String,
@@ -17,6 +18,7 @@ pub enum WorkloadMetrics {
         /// Region-specific cycles, mapping region names (e.g., "setup", "compute") to their cycle counts.
         region_cycles: HashMap<String, u64>,
     },
+    /// Metrics produced when benchmarking in proving mode
     Proving {
         /// Name of the workload (e.g., "fft", "aes").
         name: String,