heaptrack stone prover

Author: katat

src/args.rs

@@ -60,6 +60,9 @@ pub struct ProveArgs {
 
     #[clap(long = "stone_version", default_value = "v6", value_enum)]
     pub stone_version: StoneVersion,
+
+    #[clap(long = "bench-memory")]
+    pub bench_memory: Option<bool>,
 }
 
 #[derive(Args, Debug)]

src/main.rs

@@ -29,6 +29,7 @@ fn main() -> anyhow::Result<()> {
                         &run_cairo_result.air_public_input,
                         &run_cairo_result.air_private_input,
                         &tmp_dir,
+                        args.bench_memory.unwrap_or(false),
                     )
                     .map_err(|e| anyhow::anyhow!("Failed to run stone prover: {}", e))
                 });

src/prover.rs

@@ -52,6 +52,7 @@ pub fn run_stone_prover(
     air_public_input: &PathBuf,
     air_private_input: &PathBuf,
     tmp_dir: &tempfile::TempDir,
+    bench_memory: bool,
 ) -> Result<(), ProverError> {
     println!("Running prover...");
 
@@ -65,6 +66,7 @@ pub fn run_stone_prover(
         air_public_input,
         air_private_input,
         tmp_dir,
+        bench_memory,
     )?;
 
     println!("Prover finished successfully");
@@ -101,6 +103,7 @@ pub fn run_stone_prover_bootloader(
         air_public_input,
         air_private_input,
         tmp_dir,
+        false,
     )?;
 
     println!("Prover finished successfully");
@@ -118,6 +121,7 @@ fn run_stone_prover_internal(
     air_public_input: &PathBuf,
     air_private_input: &PathBuf,
     tmp_dir: &tempfile::TempDir,
+    bench_memory: bool,
 ) -> Result<(), ProverError> {
     let tmp_prover_parameters_path = tmp_dir.path().join("prover_parameters.json");
 
@@ -150,6 +154,7 @@ fn run_stone_prover_internal(
         output_file,
         true,
         stone_version,
+        bench_memory,
     )?;
 
     Ok(())
@@ -163,26 +168,47 @@ fn run_prover_from_command_line_with_annotations(
     output_file: &PathBuf,
     generate_annotations: bool,
     stone_version: &StoneVersion,
+    bench_memory: bool,
 ) -> Result<(), ProverError> {
     // TODO: Add better error handling
     let prover_run_path = match stone_version {
         StoneVersion::V5 => std::env::var("CPU_AIR_PROVER_V5").unwrap(),
         StoneVersion::V6 => std::env::var("CPU_AIR_PROVER_V6").unwrap(),
     };
 
-    let mut command = Command::new("heaptrack");
-    command
-        .arg(prover_run_path)
-        .arg("--out-file")
-        .arg(output_file)
-        .arg("--public-input-file")
-        .arg(public_input_file)
-        .arg("--private-input-file")
-        .arg(private_input_file)
-        .arg("--prover-config-file")
-        .arg(prover_config_file)
-        .arg("--parameter-file")
-        .arg(prover_parameter_file);
+    let mut command;
+    if bench_memory {
+        command = Command::new("heaptrack");
+        command
+            .arg("-o")
+            .arg("heaptrack-prover")
+            .arg(prover_run_path)
+            .arg("--out-file")
+            .arg(output_file)
+            .arg("--public-input-file")
+            .arg(public_input_file)
+            .arg("--private-input-file")
+            .arg(private_input_file)
+            .arg("--prover-config-file")
+            .arg(prover_config_file)
+            .arg("--parameter-file")
+            .arg(prover_parameter_file);
+    }
+    else {
+        command = Command::new(prover_run_path);
+        command
+            .arg("--out-file")
+            .arg(output_file)
+            .arg("--public-input-file")
+            .arg(public_input_file)
+            .arg("--private-input-file")
+            .arg(private_input_file)
+            .arg("--prover-config-file")
+            .arg(prover_config_file)
+            .arg("--parameter-file")
+            .arg(prover_parameter_file);
+    }
+
     if generate_annotations {
         command.arg("--generate-annotations");
     }