Add --air_public_input to runner (#1268)

* Add air_public_input arg

* Added get_public_memory_addresses()

* Added get_memory_segment_addresses

* Added write_air_public_input skeleton

* Added air_public_input.rs

* Added fixme for unwrap

* Fix get_memory_segment_addresses()

* Finished implementation

* Fixed relocated memory issue

* Fixed trace issue

* Removed rc limits unwrap

* Refactor

* Added changelog entry

* Remove memory addresses structs

* Fix changelog

* Refactor get_public_memory_addresses

* Fixed relocation issue on builtin addresses

* Replaced get_trace() with get_relocated_trace()

* Use references on new()

* Removed trace indexing

* Remove allow unused from air public input

* Fix all builtins stop_ptr

* Handle write result

* Removed indexing

* Remove indexing

* Partial fix cli error

* Fixed error formatting

* Added get_air_public_input() to cairo runner

* Fix clippy

* Change to clearer code

* Update vm/src/air_public_input.rs

Co-authored-by: fmoletta <[email protected]>

* Update changelog

* Update vm/src/air_public_input.rs

Co-authored-by: fmoletta <[email protected]>

* Reorder to avoid cloning

* Fix main test

* Custom ser for mem value

* Another fix for main test

* Return NoStopPtr err instead of using Option

* Make get_layout private

* Clone page_id

* Return err if rc values are missing

* Added new MemorySegmentAddresses

* Fix changelog

* Fix test main

* Fix wasm build

* Fix clippy

* Fix nostd again

* Dummy commit

* Revert "Dummy commit"

This reverts commit 6b01c66.

* New main test

* Added air_public_input test, changed main

* Allow too many args in test

* Try deny deny

* Fix clippy

* Fix clap

* Fix argument

* Changed run rstest

---------

Co-authored-by: fmoletta <[email protected]>
Co-authored-by: Pedro Fontana <[email protected]>
Co-authored-by: Mario Rugiero <[email protected]>

CHANGELOG.md

@@ -1,6 +1,9 @@
 ## Cairo-VM Changelog
 
 #### Upcoming Changes
+* feat: implement the `--air_public_input` flag to the runner for outputting public inputs into a file [#1268](https://github.com/lambdaclass/cairo-rs/pull/1268)
+
+* fix: CLI errors bad formatting and handling
 
 * fix: return error when a parsed hint's PC is invalid [#1340](https://github.com/lambdaclass/cairo-vm/pull/1340)
 

cairo-vm-cli/src/main.rs

@@ -1,12 +1,13 @@
 #![deny(warnings)]
 #![forbid(unsafe_code)]
 use bincode::enc::write::Writer;
+use cairo_vm::air_public_input::PublicInputError;
 use cairo_vm::cairo_run::{self, EncodeTraceError};
 use cairo_vm::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor;
 use cairo_vm::vm::errors::cairo_run_errors::CairoRunError;
 use cairo_vm::vm::errors::trace_errors::TraceError;
 use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
-use clap::{Parser, ValueHint};
+use clap::{CommandFactory, Parser, ValueHint};
 use std::io::{self, Write};
 use std::path::PathBuf;
 use thiserror::Error;
@@ -37,6 +38,8 @@ struct Args {
     proof_mode: bool,
     #[structopt(long = "secure_run")]
     secure_run: Option<bool>,
+    #[clap(long = "air_public_input")]
+    air_public_input: Option<String>,
 }
 
 fn validate_layout(value: &str) -> Result<String, String> {
@@ -68,6 +71,8 @@ enum Error {
     VirtualMachine(#[from] VirtualMachineError),
     #[error(transparent)]
     Trace(#[from] TraceError),
+    #[error(transparent)]
+    PublicInput(#[from] PublicInputError),
 }
 
 struct FileWriter {
@@ -104,20 +109,22 @@ impl FileWriter {
 }
 
 fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
-    let args = Args::try_parse_from(args);
-    let args = match args {
-        Ok(args) => args,
-        Err(error) => {
-            eprintln!("{error}");
-            return Err(Error::Cli(error));
-        }
-    };
-    let trace_enabled = args.trace_file.is_some();
+    let args = Args::try_parse_from(args)?;
+
+    if args.air_public_input.is_some() && !args.proof_mode {
+        let error = Args::command().error(
+            clap::error::ErrorKind::ArgumentConflict,
+            "--air_public_input can only be used in proof_mode.",
+        );
+        return Err(Error::Cli(error));
+    }
+
+    let trace_enabled = args.trace_file.is_some() || args.air_public_input.is_some();
     let mut hint_executor = BuiltinHintProcessor::new_empty();
     let cairo_run_config = cairo_run::CairoRunConfig {
         entrypoint: &args.entrypoint,
         trace_enabled,
-        relocate_mem: args.memory_file.is_some(),
+        relocate_mem: args.memory_file.is_some() || args.air_public_input.is_some(),
         layout: &args.layout,
         proof_mode: args.proof_mode,
         secure_run: args.secure_run,
@@ -160,21 +167,28 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
         memory_writer.flush()?;
     }
 
+    if let Some(file_path) = args.air_public_input {
+        let json = cairo_runner.get_air_public_input(&vm)?.serialize_json()?;
+        std::fs::write(file_path, json)?;
+    }
+
     Ok(())
 }
 
 fn main() -> Result<(), Error> {
+    #[cfg(test)]
+    return Ok(());
+
+    #[cfg(not(test))]
     match run(std::env::args()) {
-        Ok(()) => Ok(()),
-        Err(Error::Cli(_)) => {
-            Ok(()) // Exit with code 0 to avoid printing CLI error message
-        }
-        Err(error) => Err(error),
+        Err(Error::Cli(err)) => err.exit(),
+        other => other,
     }
 }
 
 #[cfg(test)]
 mod tests {
+    #![allow(clippy::too_many_arguments)]
     use super::*;
     use assert_matches::assert_matches;
     use rstest::rstest;
@@ -215,11 +229,18 @@ mod tests {
         #[values(false, true)] secure_run: bool,
         #[values(false, true)] print_output: bool,
         #[values(false, true)] entrypoint: bool,
+        #[values(false, true)] air_public_input: bool,
     ) {
         let mut args = vec!["cairo-vm-cli".to_string()];
         if let Some(layout) = layout {
             args.extend_from_slice(&["--layout".to_string(), layout.to_string()]);
         }
+        if air_public_input {
+            args.extend_from_slice(&[
+                "--air_public_input".to_string(),
+                "air_input.pub".to_string(),
+            ]);
+        }
         if proof_mode {
             trace_file = true;
             args.extend_from_slice(&["--proof_mode".to_string()]);
@@ -239,8 +260,13 @@ mod tests {
         if print_output {
             args.extend_from_slice(&["--print_output".to_string()]);
         }
+
         args.push("../cairo_programs/proof_programs/fibonacci.json".to_string());
-        assert_matches!(run(args.into_iter()), Ok(_));
+        if air_public_input && !proof_mode {
+            assert_matches!(run(args.into_iter()), Err(_));
+        } else {
+            assert_matches!(run(args.into_iter()), Ok(_));
+        }
     }
 
     #[test]
@@ -266,7 +292,7 @@ mod tests {
     //to fool Codecov.
     #[test]
     fn test_main() {
-        assert!(main().is_ok());
+        main().unwrap();
     }
 
     #[test]

vm/src/air_public_input.rs

@@ -0,0 +1,140 @@
+use felt::Felt252;
+use serde::Serialize;
+use thiserror_no_std::Error;
+
+use crate::{
+    stdlib::{
+        collections::HashMap,
+        prelude::{String, Vec},
+    },
+    types::layout::CairoLayout,
+    vm::{
+        errors::{trace_errors::TraceError, vm_errors::VirtualMachineError},
+        trace::trace_entry::TraceEntry,
+    },
+};
+
+#[derive(Serialize, Debug)]
+pub struct PublicMemoryEntry {
+    address: usize,
+    page: usize,
+    #[serde(serialize_with = "mem_value_serde::serialize")]
+    value: Option<Felt252>,
+}
+
+mod mem_value_serde {
+    use super::*;
+    use serde::Serializer;
+
+    pub(crate) fn serialize<S: Serializer>(
+        value: &Option<Felt252>,
+        serializer: S,
+    ) -> Result<S::Ok, S::Error> {
+        if let Some(value) = value {
+            serializer.serialize_str(&format!("0x{}", value.to_str_radix(16)))
+        } else {
+            serializer.serialize_none()
+        }
+    }
+}
+
+#[derive(Serialize, Debug)]
+struct MemorySegmentAddresses {
+    begin_addr: usize,
+    stop_addr: usize,
+}
+
+impl From<(usize, usize)> for MemorySegmentAddresses {
+    fn from(addresses: (usize, usize)) -> Self {
+        let (begin_addr, stop_addr) = addresses;
+        MemorySegmentAddresses {
+            begin_addr,
+            stop_addr,
+        }
+    }
+}
+
+#[derive(Serialize, Debug)]
+pub struct PublicInput<'a> {
+    layout: &'a str,
+    layout_params: Option<&'a CairoLayout>,
+    rc_min: isize,
+    rc_max: isize,
+    n_steps: usize,
+    memory_segments: HashMap<&'a str, MemorySegmentAddresses>,
+    public_memory: Vec<PublicMemoryEntry>,
+}
+
+impl<'a> PublicInput<'a> {
+    pub fn new(
+        memory: &[Option<Felt252>],
+        layout: &'a str,
+        dyn_layout_params: Option<&'a CairoLayout>,
+        public_memory_addresses: &[(usize, usize)],
+        memory_segment_addresses: HashMap<&'static str, (usize, usize)>,
+        trace: &[TraceEntry],
+        rc_limits: (isize, isize),
+    ) -> Result<Self, PublicInputError> {
+        let memory_entry =
+            |addresses: &(usize, usize)| -> Result<PublicMemoryEntry, PublicInputError> {
+                let (address, page) = addresses;
+                Ok(PublicMemoryEntry {
+                    address: *address,
+                    page: *page,
+                    value: memory
+                        .get(*address)
+                        .ok_or(PublicInputError::MemoryNotFound(*address))?
+                        .clone(),
+                })
+            };
+        let public_memory = public_memory_addresses
+            .iter()
+            .map(memory_entry)
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let (rc_min, rc_max) = rc_limits;
+
+        let trace_first = trace.first().ok_or(PublicInputError::EmptyTrace)?;
+        let trace_last = trace.last().ok_or(PublicInputError::EmptyTrace)?;
+
+        Ok(PublicInput {
+            layout,
+            layout_params: dyn_layout_params,
+            rc_min,
+            rc_max,
+            n_steps: trace.len(),
+            memory_segments: {
+                let mut memory_segment_addresses = memory_segment_addresses
+                    .into_iter()
+                    .map(|(n, s)| (n, s.into()))
+                    .collect::<HashMap<_, MemorySegmentAddresses>>();
+
+                memory_segment_addresses.insert("program", (trace_first.pc, trace_last.pc).into());
+                memory_segment_addresses
+                    .insert("execution", (trace_first.ap, trace_last.ap).into());
+                memory_segment_addresses
+            },
+            public_memory,
+        })
+    }
+
+    pub fn serialize_json(&self) -> Result<String, PublicInputError> {
+        serde_json::to_string_pretty(&self).map_err(PublicInputError::from)
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum PublicInputError {
+    #[error("The trace slice provided is empty")]
+    EmptyTrace,
+    #[error("The provided memory doesn't contain public address {0}")]
+    MemoryNotFound(usize),
+    #[error("Range check values are missing")]
+    NoRangeCheckLimits,
+    #[error("Failed to (de)serialize data")]
+    Serde(#[from] serde_json::Error),
+    #[error(transparent)]
+    VirtualMachine(#[from] VirtualMachineError),
+    #[error(transparent)]
+    Trace(#[from] TraceError),
+}

vm/src/lib.rs

@@ -54,6 +54,7 @@ mod stdlib {
 }
 
 pub extern crate felt;
+pub mod air_public_input;
 pub mod cairo_run;
 pub mod hint_processor;
 pub mod math_utils;

vm/src/types/instance_definitions/bitwise_instance_def.rs

@@ -1,7 +1,9 @@
+use serde::Serialize;
+
 pub(crate) const CELLS_PER_BITWISE: u32 = 5;
 pub(crate) const INPUT_CELLS_PER_BITWISE: u32 = 2;
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Serialize, Clone, Debug, PartialEq)]
 pub(crate) struct BitwiseInstanceDef {
     pub(crate) ratio: Option<u32>,
     pub(crate) total_n_bits: u32,

vm/src/types/instance_definitions/builtins_instance_def.rs

@@ -4,8 +4,9 @@ use super::{
     pedersen_instance_def::PedersenInstanceDef, poseidon_instance_def::PoseidonInstanceDef,
     range_check_instance_def::RangeCheckInstanceDef,
 };
+use serde::Serialize;
 
-#[derive(Debug, PartialEq)]
+#[derive(Serialize, Debug, PartialEq)]
 pub(crate) struct BuiltinsInstanceDef {
     pub(crate) output: bool,
     pub(crate) pedersen: Option<PedersenInstanceDef>,

vm/src/types/instance_definitions/cpu_instance_def.rs

@@ -1,4 +1,6 @@
-#[derive(Debug, PartialEq)]
+use serde::Serialize;
+
+#[derive(Serialize, Debug, PartialEq)]
 pub(crate) struct CpuInstanceDef {
     pub(crate) _safe_call: bool,
 }

vm/src/types/instance_definitions/diluted_pool_instance_def.rs

@@ -1,4 +1,6 @@
-#[derive(Debug, PartialEq)]
+use serde::Serialize;
+
+#[derive(Serialize, Debug, PartialEq)]
 pub(crate) struct DilutedPoolInstanceDef {
     pub(crate) units_per_step: u32,
     pub(crate) spacing: u32,

vm/src/types/instance_definitions/ec_op_instance_def.rs

@@ -1,7 +1,9 @@
+use serde::Serialize;
+
 pub(crate) const CELLS_PER_EC_OP: u32 = 7;
 pub(crate) const INPUT_CELLS_PER_EC_OP: u32 = 5;
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Serialize, Clone, Debug, PartialEq)]
 pub(crate) struct EcOpInstanceDef {
     pub(crate) ratio: Option<u32>,
     pub(crate) scalar_height: u32,

vm/src/types/instance_definitions/ecdsa_instance_def.rs

@@ -1,7 +1,9 @@
+use serde::Serialize;
+
 pub(crate) const CELLS_PER_SIGNATURE: u32 = 2;
 pub(crate) const _INPUTCELLS_PER_SIGNATURE: u32 = 2;
 
-#[derive(Debug, PartialEq)]
+#[derive(Serialize, Debug, PartialEq)]
 pub(crate) struct EcdsaInstanceDef {
     pub(crate) ratio: Option<u32>,
     pub(crate) _repetitions: u32,

vm/src/types/instance_definitions/keccak_instance_def.rs

@@ -1,6 +1,7 @@
 use crate::stdlib::prelude::*;
+use serde::Serialize;
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Serialize, Clone, Debug, PartialEq)]
 pub(crate) struct KeccakInstanceDef {
     pub(crate) ratio: Option<u32>,
     pub(crate) _state_rep: Vec<u32>,

vm/src/types/instance_definitions/pedersen_instance_def.rs

@@ -1,9 +1,10 @@
 use num_bigint::{BigInt, Sign};
+use serde::Serialize;
 
 pub(crate) const CELLS_PER_HASH: u32 = 3;
 pub(crate) const INPUT_CELLS_PER_HASH: u32 = 2;
 
-#[derive(Debug, PartialEq)]
+#[derive(Serialize, Debug, PartialEq)]
 pub(crate) struct PedersenInstanceDef {
     pub(crate) ratio: Option<u32>,
     pub(crate) _repetitions: u32,

vm/src/types/instance_definitions/poseidon_instance_def.rs

@@ -1,7 +1,9 @@
+use serde::Serialize;
+
 pub(crate) const CELLS_PER_POSEIDON: u32 = 6;
 pub(crate) const INPUT_CELLS_PER_POSEIDON: u32 = 3;
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Serialize, Clone, Debug, PartialEq)]
 pub(crate) struct PoseidonInstanceDef {
     pub(crate) ratio: Option<u32>,
 }