Add air_public_input flag to cairo1-run

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 #### Upcoming Changes
 
+* feat: Add `air_public_input` flag to `cairo1-run` [#1539] (https://github.com/lambdaclass/cairo-vm/pull/1539)
+
+* feat: Add `proof_mode` flag to `cairo1-run` [#1537] (https://github.com/lambdaclass/cairo-vm/pull/1537)
+  * The cairo1-run crate no longer compiles and executes in proof_mode by default
+  * Add flag `proof_mode` to cairo1-run crate. Activating this flag will enable proof_mode compilation and execution
+
 * feat(BREAKING): Replace `cairo-felt` crate with `starknet-types-core` (0.0.5) [#1408](https://github.com/lambdaclass/cairo-vm/pull/1408)
 
 * feat(BREAKING): Add Cairo 1 proof mode compilation and execution [#1517] (https://github.com/lambdaclass/cairo-vm/pull/1517)

cairo1-run/src/main.rs

@@ -47,6 +47,10 @@
 use cairo_vm::vm::errors::runner_errors::RunnerError;
 use cairo_vm::vm::errors::trace_errors::TraceError;
 use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
+use cairo_vm::vm::runners::builtin_runner::{
+    BITWISE_BUILTIN_NAME, EC_OP_BUILTIN_NAME, HASH_BUILTIN_NAME, OUTPUT_BUILTIN_NAME,
+    POSEIDON_BUILTIN_NAME, RANGE_CHECK_BUILTIN_NAME, SIGNATURE_BUILTIN_NAME,
+};
 use cairo_vm::vm::runners::cairo_runner::RunnerMode;
 use cairo_vm::{
     serde::deserialize_program::ReferenceManager,
@@ -77,6 +81,10 @@
     memory_file: Option<PathBuf>,
     #[clap(long = "layout", default_value = "plain", value_parser=validate_layout)]
     layout: String,
+    #[clap(long = "proof_mode", value_parser)]
+    proof_mode: bool,
+    #[clap(long = "air_public_input", value_parser)]
+    air_public_input: Option<PathBuf>,
 }
 
 fn validate_layout(value: &str) -> Result<String, String> {
@@ -171,6 +179,13 @@
 
 fn run(args: impl Iterator<Item = String>) -> Result<Vec<MaybeRelocatable>, Error> {
     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 compiler_config = CompilerConfig {
         replace_ids: true,
@@ -202,29 +217,34 @@
         &type_sizes,
         main_func,
         initial_gas,
+        args.proof_mode,
     )?;
 
-    println!("Compiling with proof mode and running ...");
-
-    // This information can be useful for the users using the prover.
-    println!("Builtins used: {:?}", builtins);
-
-    // Prepare "canonical" proof mode instructions. These are usually added by the compiler in cairo 0
-    let mut ctx = casm! {};
-    casm_extend! {ctx,
-        call rel 4;
-        jmp rel 0;
-    };
-    let proof_mode_header = ctx.instructions;
-
     // Get the user program instructions
     let program_instructions = casm_program.instructions.iter();
 
     // This footer is used by lib funcs
     let libfunc_footer = create_code_footer();
 
-    // This is the program we are actually proving
-    // With embedded proof mode, cairo1 header and the libfunc footer
+    let proof_mode_header = if args.proof_mode {
+        println!("Compiling with proof mode and running ...");
+
+        // This information can be useful for the users using the prover.
+        println!("Builtins used: {:?}", builtins);
+
+        // Prepare "canonical" proof mode instructions. These are usually added by the compiler in cairo 0
+        let mut ctx = casm! {};
+        casm_extend! {ctx,
+            call rel 4;
+            jmp rel 0;
+        };
+        ctx.instructions
+    } else {
+        casm! {}.instructions
+    };
+
+    // This is the program we are actually running/proving
+    // With (embedded proof mode), cairo1 header and the libfunc footer
     let instructions = chain!(
         proof_mode_header.iter(),
         entry_code.iter(),
@@ -244,36 +264,53 @@
 
     let data_len = data.len();
 
-    let starting_pc = 0;
-
-    let program = Program::new_for_proof(
-        builtins,
-        data,
-        starting_pc,
-        // Proof mode is on top
-        // jmp rel 0 is on PC == 2
-        2,
-        program_hints,
-        ReferenceManager {
-            references: Vec::new(),
-        },
-        HashMap::new(),
-        vec![],
-        None,
-    )?;
+    let program = if args.proof_mode {
+        Program::new_for_proof(
+            builtins,
+            data,
+            0,
+            // Proof mode is on top
+            // jmp rel 0 is on PC == 2
+            2,
+            program_hints,
+            ReferenceManager {
+                references: Vec::new(),
+            },
+            HashMap::new(),
+            vec![],
+            None,
+        )?
+    } else {
+        Program::new(
+            builtins,
+            data,
+            Some(0),
+            program_hints,
+            ReferenceManager {
+                references: Vec::new(),
+            },
+            HashMap::new(),
+            vec![],
+            None,
+        )?
+    };
 
-    let mut runner = CairoRunner::new_v2(&program, &args.layout, RunnerMode::ProofModeCairo1)?;
+    let runner_mode = if args.proof_mode {
+        RunnerMode::ProofModeCairo1
+    } else {
+        RunnerMode::ExecutionMode
+    };
 
-    let mut vm = VirtualMachine::new(args.trace_file.is_some());
+    let mut runner = CairoRunner::new_v2(&program, &args.layout, runner_mode)?;
+
+    let mut vm = VirtualMachine::new(args.trace_file.is_some() || args.air_public_input.is_some());
     let end = runner.initialize(&mut vm)?;
 
     additional_initialization(&mut vm, data_len)?;
 
-    // Run it until the infinite loop
+    // Run it until the end/ infinite loop in proof_mode
     runner.run_until_pc(end, &mut vm, &mut hint_processor)?;
-
-    // Then pad it to the power of 2
-    runner.run_until_next_power_of_2(&mut vm, &mut hint_processor)?;
+    runner.end_run(false, false, &mut vm, &mut hint_processor)?;
 
     // Fetch return type data
     let return_type_id = main_func
@@ -322,8 +359,59 @@
         }
     }
 
+    // Set stop pointers for builtins so we can obtain the air public input
+    if args.air_public_input.is_some() {
+        // Cairo 1 programs have other return values aside from the used builtin's final pointers, so we need to hand-pick them
+        let ret_types_sizes = main_func
+            .signature
+            .ret_types
+            .iter()
+            .map(|id| type_sizes.get(id).cloned().unwrap_or_default());
+        let ret_types_and_sizes = main_func
+            .signature
+            .ret_types
+            .iter()
+            .zip(ret_types_sizes.clone());
+
+        let full_ret_types_size: i16 = ret_types_sizes.sum();
+        let mut stack_pointer = (vm.get_ap() - (full_ret_types_size as usize).saturating_sub(1))
+            .map_err(VirtualMachineError::Math)?;
+
+        // Calculate the stack_ptr for each return builtin in the return values
+        let mut builtin_name_to_stack_pointer = HashMap::new();
+        for (id, size) in ret_types_and_sizes {
+            if let Some(ref name) = id.debug_name {
+                let builtin_name = match &*name.to_string() {
+                    "RangeCheck" => RANGE_CHECK_BUILTIN_NAME,
+                    "Poseidon" => POSEIDON_BUILTIN_NAME,
+                    "EcOp" => EC_OP_BUILTIN_NAME,
+                    "Bitwise" => BITWISE_BUILTIN_NAME,
+                    "Pedersen" => HASH_BUILTIN_NAME,
+                    "Output" => OUTPUT_BUILTIN_NAME,
+                    "Ecdsa" => SIGNATURE_BUILTIN_NAME,
+                    _ => {
+                        stack_pointer.offset += size as usize;
+                        continue;
+                    }
+                };
+                builtin_name_to_stack_pointer.insert(builtin_name, stack_pointer);
+            }
+            stack_pointer.offset += size as usize;
+        }
+        // Set stop pointer for each builtin
+        vm.builtins_final_stack_from_stack_pointer_dict(&builtin_name_to_stack_pointer)?;
+
+        // Build execution public memory
+        runner.finalize_segments(&mut vm)?;
+    }
+
     runner.relocate(&mut vm, true)?;
 
+    if let Some(file_path) = args.air_public_input {
+        let json = runner.get_air_public_input(&vm)?.serialize_json()?;
+        std::fs::write(file_path, json)?;
+    }
+
     if let Some(trace_path) = args.trace_file {
         let relocated_trace = runner
             .relocated_trace
@@ -465,6 +553,7 @@
     type_sizes: &UnorderedHashMap<ConcreteTypeId, i16>,
     func: &Function,
     initial_gas: usize,
+    proof_mode: bool,
 ) -> Result<(Vec<Instruction>, Vec<BuiltinName>), Error> {
     let mut ctx = casm! {};
     // The builtins in the formatting expected by the runner.
@@ -497,8 +586,11 @@
         let ty_size = type_sizes[ty];
         let generic_ty = &info.long_id.generic_id;
         if let Some(offset) = builtin_offset.get(generic_ty) {
-            // Everything is off by 2 due to the proof mode header
-            let offset = offset + 2;
+            let mut offset = *offset;
+            if proof_mode {
+                // Everything is off by 2 due to the proof mode header
+                offset += 2;
+            }
             casm_extend! {ctx,
                 [ap + 0] = [fp - offset], ap++;
             }
@@ -661,111 +753,127 @@
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/fibonacci.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/fibonacci.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_fibonacci_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(89)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/factorial.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/factorial.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_factorial_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(3628800)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/array_get.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/array_get.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_array_get_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(3)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/enum_flow.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/enum_flow.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_enum_flow_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(300)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/enum_match.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/enum_match.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_enum_match_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(10), MaybeRelocatable::from(felt_str("3618502788666131213697322783095070105623107215331596699973092056135872020471"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/hello.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/hello.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_hello_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(1), MaybeRelocatable::from(1234)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/ops.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/ops.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_ops_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(6)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/print.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/print.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_print_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/recursion.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/recursion.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_recursion_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(felt_str("1154076154663935037074198317650845438095734251249125412074882362667803016453"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/sample.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/sample.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_sample_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(felt_str("5050"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/poseidon.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/poseidon.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_poseidon_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(felt_str("1099385018355113290651252669115094675591288647745213771718157553170111442461"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/poseidon_pedersen.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/poseidon_pedersen.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_poseidon_pedersen_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(felt_str("1036257840396636296853154602823055519264738423488122322497453114874087006398"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/pedersen_example.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/pedersen_example.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_pedersen_example_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(felt_str("1089549915800264549621536909767699778745926517555586332772759280702396009108"))]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/simple.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/simple.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_simple_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(1)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/simple_struct.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/simple_struct.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_simple_struct_ok(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(100)]);
     }
 
     #[rstest]
     #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/dictionaries.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo"].as_slice())]
+    #[case(["cairo1-run", "../cairo_programs/cairo-1-programs/dictionaries.cairo", "--trace_file", "/dev/null", "--memory_file", "/dev/null", "--layout", "all_cairo", "--proof_mode", "--air_public_input", "/dev/null"].as_slice())]
     fn test_run_dictionaries(#[case] args: &[&str]) {
         let args = args.iter().cloned().map(String::from);
         assert_matches!(run(args), Ok(res) if res == vec![MaybeRelocatable::from(1024)]);