Stav/add prover input info flag (#2047)

Author: Stavbe

cairo-vm-cli/src/main.rs

@@ -14,7 +14,7 @@ use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
 use cairo_vm::vm::runners::cairo_pie::CairoPie;
 #[cfg(feature = "with_tracer")]
 use cairo_vm::vm::runners::cairo_runner::CairoRunner;
-use cairo_vm::vm::runners::cairo_runner::RunResources;
+use cairo_vm::vm::runners::cairo_runner::{ProverInputInfoError, RunResources};
 #[cfg(feature = "with_tracer")]
 use cairo_vm_tracer::error::trace_data_errors::TraceDataError;
 #[cfg(feature = "with_tracer")]
@@ -69,6 +69,8 @@ struct Args {
         conflicts_with_all = ["proof_mode", "air_private_input", "air_public_input"]
     )]
     cairo_pie_output: Option<String>,
+    #[arg(long = "prover_input_info")]
+    prover_input_info: Option<String>,
     #[arg(long = "merge_extra_segments")]
     merge_extra_segments: bool,
     #[arg(long = "allow_missing_builtins")]
@@ -102,6 +104,8 @@ enum Error {
     #[error(transparent)]
     PublicInput(#[from] PublicInputError),
     #[error(transparent)]
+    ProveInputInfo(#[from] ProverInputInfoError),
+    #[error(transparent)]
     #[cfg(feature = "with_tracer")]
     TraceData(#[from] TraceDataError),
 }
@@ -168,7 +172,9 @@ fn start_tracer(cairo_runner: &CairoRunner) -> Result<(), TraceDataError> {
 fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
     let args = Args::try_parse_from(args)?;
 
-    let trace_enabled = args.trace_file.is_some() || args.air_public_input.is_some();
+    let trace_enabled = args.trace_file.is_some()
+        || args.air_public_input.is_some()
+        || args.prover_input_info.is_some();
 
     let cairo_layout_params = match args.cairo_layout_params_file {
         Some(file) => Some(CairoLayoutParams::from_file(&file)?),
@@ -235,6 +241,15 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
         memory_writer.flush()?;
     }
 
+    if let Some(prover_input_info_path) = args.prover_input_info {
+        let prover_input_info = cairo_runner.get_prover_input_info().map_err(|error| {
+            eprintln!("{error}");
+            CairoRunError::Runner(error)
+        })?;
+        let json = prover_input_info.serialize_json()?;
+        std::fs::write(prover_input_info_path, json)?;
+    }
+
     if let Some(file_path) = args.air_public_input {
         let json = cairo_runner.get_air_public_input()?.serialize_json()?;
         std::fs::write(file_path, json)?;
@@ -358,6 +373,7 @@ mod tests {
         #[values(false, true)] air_public_input: bool,
         #[values(false, true)] air_private_input: bool,
         #[values(false, true)] cairo_pie_output: bool,
+        #[values(false, true)] prover_input_info: bool,
     ) {
         let mut args = vec!["cairo-vm-cli".to_string()];
         if let Some(layout) = layout {
@@ -388,6 +404,9 @@ mod tests {
         if print_output {
             args.extend_from_slice(&["--print_output".to_string()]);
         }
+        if prover_input_info {
+            args.extend_from_slice(&["--prover_input_info".to_string(), "/dev/null".to_string()]);
+        }
 
         args.push("../cairo_programs/proof_programs/fibonacci.json".to_string());
         if air_public_input && !proof_mode

vm/src/vm/runners/cairo_runner.rs

@@ -48,6 +48,7 @@ use crate::{
 use num_integer::div_rem;
 use num_traits::{ToPrimitive, Zero};
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
 
 use super::{builtin_runner::ModBuiltinRunner, cairo_pie::CairoPieAdditionalData};
 use super::{
@@ -1542,6 +1543,7 @@ impl CairoRunner {
     }
 }
 
+// TODO(Stav): move to specified file.
 //* ----------------------
 //*   ProverInputInfo
 //* ----------------------
@@ -1559,6 +1561,19 @@ pub struct ProverInputInfo {
     pub builtins_segments: BTreeMap<usize, BuiltinName>,
 }
 
+impl ProverInputInfo {
+    pub fn serialize_json(&self) -> Result<String, ProverInputInfoError> {
+        serde_json::to_string_pretty(&self).map_err(ProverInputInfoError::from)
+    }
+}
+
+// TODO(Stav): add TraceNotEnabled error.
+#[derive(Debug, Error)]
+pub enum ProverInputInfoError {
+    #[error("Failed to (de)serialize data")]
+    Serde(#[from] serde_json::Error),
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct SegmentInfo {
     pub index: isize,
@@ -1666,6 +1681,7 @@ mod tests {
     use crate::types::instance_definitions::bitwise_instance_def::CELLS_PER_BITWISE;
     use crate::types::instance_definitions::keccak_instance_def::CELLS_PER_KECCAK;
     use crate::vm::vm_memory::memory::MemoryCell;
+    use rstest::rstest;
 
     use crate::felt_hex;
     use crate::{
@@ -5671,4 +5687,44 @@ mod tests {
         assert!(prover_input.relocatable_memory[6].len() as u32 % CELLS_PER_BITWISE == 0);
         assert!(cairo_runner.vm.segments.memory.data[8].len() as u32 % CELLS_PER_KECCAK == 0);
     }
+
+    #[rstest]
+    #[case(include_bytes!("../../../../cairo_programs/proof_programs/fibonacci.json"))]
+    #[case(include_bytes!("../../../../cairo_programs/proof_programs/bitwise_output.json"))]
+    #[case(include_bytes!("../../../../cairo_programs/proof_programs/poseidon_builtin.json"))]
+    #[case(include_bytes!("../../../../cairo_programs/proof_programs/relocate_temporary_segment_append.json"))]
+    #[case(include_bytes!("../../../../cairo_programs/proof_programs/pedersen_test.json"))]
+    fn serialize_and_deserialize_prover_input_info(#[case] program_content: &[u8]) {
+        use crate::types::layout_name::LayoutName;
+
+        let config = crate::cairo_run::CairoRunConfig {
+            proof_mode: false,
+            relocate_mem: false,
+            trace_enabled: true,
+            layout: LayoutName::all_cairo_stwo,
+            ..Default::default()
+        };
+        let runner = crate::cairo_run::cairo_run(program_content, &config, &mut crate::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor::new_empty()).unwrap();
+        let prover_input_info = runner.get_prover_input_info().unwrap();
+        let serialized_prover_input_info = prover_input_info.serialize_json().unwrap();
+        let deserialized_prover_input_info: ProverInputInfo =
+            serde_json::from_str(&serialized_prover_input_info).unwrap();
+        // Check that the deserialized prover input info is equal to the original one.
+        assert_eq!(
+            prover_input_info.relocatable_memory,
+            deserialized_prover_input_info.relocatable_memory
+        );
+        assert_eq!(
+            prover_input_info.relocatable_trace,
+            deserialized_prover_input_info.relocatable_trace
+        );
+        assert_eq!(
+            prover_input_info.builtins_segments,
+            deserialized_prover_input_info.builtins_segments
+        );
+        assert_eq!(
+            prover_input_info.public_memory_offsets,
+            deserialized_prover_input_info.public_memory_offsets
+        );
+    }
 }