Compare air_public_inputs against python vm + Fix how public memory is built

.github/workflows/rust.yml

@@ -315,6 +315,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-reference-trace-cache-${{ hashFiles('cairo_programs/**/*.cairo', 'examples/wasm-demo/src/array_sum.cairo') }}
         restore-keys: ${{ matrix.program-target }}-reference-trace-cache-
 
@@ -354,7 +355,7 @@ jobs:
         include:
         - program-target: cairo_proof_programs
           programs-dir: cairo_programs/proof_programs
-          extra-args: '--proof_mode'
+          extra-args: '--proof_mode --air_public_input {program}.rs.air_public_input '
         - program-target: cairo_test_programs
           programs-dir: cairo_programs
           extra-args: ''
@@ -392,6 +393,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-release-trace-cache-${{ github.sha }}
 
 
@@ -466,6 +468,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-reference-trace-cache-${{ hashFiles('cairo_programs/**/*.cairo', 'examples/wasm-demo/src/array_sum.cairo') }}
         fail-on-cache-miss: true
 
@@ -475,14 +478,17 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-release-trace-cache-${{ github.sha }}
         fail-on-cache-miss: true
 
     - name: Run comparison script
       run: |
         if [ ${{ matrix.program-target }} = cairo_proof_programs ]; then
-          PROOF=proof_mode
+          PROOF=proof_mode 
+          AIR_PUBLIC_INPUT=air_public_input
         fi
+        ./vm/src/tests/compare_vm_state.sh trace memory $PROOF $AIR_PUBLIC_INPUT
         ./vm/src/tests/compare_vm_state.sh trace memory $PROOF
 
   wasm-demo:

.gitignore

@@ -8,6 +8,7 @@
 **/*.sierra
 **/*.trace
 **/*.memory
+**/*.air_public_input
 **/*.swp
 bench/results
 .python-version

CHANGELOG.md

@@ -24,6 +24,21 @@
 
 * feat: Implement `CairoRunner.get_cairo_pie`[#1375](https://github.com/lambdaclass/cairo-vm/pull/1375)
 
+* fix: Compare air_public_inputs against python vm + Fix how public memory is built [#391](https://github.com/lambdaclass/cairo-vm/pull/1391)
+
+    BugFixes:
+
+    *  `CairoRunner.finalize_segments` now builds the output builtin's public memory (if applicable).
+    * `MemorySegmentManager.get_public_memory_addresses` logic fixed.
+    * `MemorySegmentManager.finalize` no longer skips segments when their public memory is None
+
+    Minor changes:
+
+    * `VirtualMachine.get_public_memory_addresses` now strips the "_builtin" suffix from builtin names
+    * `MemorySegmentAddresses.stop_address` renamed to `stop_ptr`
+
+    Overall these changes make the the air public input file (obtained through the --air_public_input flag) equivalent to the ones outputted by the cairo-lang version
+
 * fix: Fix `SPLIT_FELT` hint [#1387](https://github.com/lambdaclass/cairo-vm/pull/1387)
 
 * refactor: combine `Program.hints` and `Program.hints_ranges` into custom collection [#1366](https://github.com/lambdaclass/cairo-vm/pull/1366)

Makefile

@@ -11,7 +11,7 @@ STARKNET_SIERRA_COMPILE_CAIRO_2:=cairo2/bin/starknet-sierra-compile
 	deps deps-macos cargo-deps build run check test clippy coverage benchmark flamegraph \
 	compare_benchmarks_deps compare_benchmarks docs clean \
 	compare_vm_output compare_trace_memory compare_trace compare_memory \
-	compare_trace_memory_proof compare_trace_proof compare_memory_proof \
+	compare_trace_memory_proof  compare_all_proof compare_trace_proof compare_memory_proof compare_air_public_input \
 	cairo_bench_programs cairo_proof_programs cairo_test_programs cairo_1_test_contracts cairo_2_test_contracts \
 	cairo_trace cairo-vm_trace cairo_proof_trace cairo-vm_proof_trace \
 	$(RELBIN) $(DBGBIN)
@@ -27,10 +27,14 @@ STARKNET_SIERRA_COMPILE_CAIRO_2:=cairo2/bin/starknet-sierra-compile
 TEST_PROOF_DIR=cairo_programs/proof_programs
 TEST_PROOF_FILES:=$(wildcard $(TEST_PROOF_DIR)/*.cairo)
 COMPILED_PROOF_TESTS:=$(patsubst $(TEST_PROOF_DIR)/%.cairo, $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_FILES))
+
 CAIRO_MEM_PROOF:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.memory, $(COMPILED_PROOF_TESTS))
 CAIRO_TRACE_PROOF:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.trace, $(COMPILED_PROOF_TESTS))
+CAIRO_AIR_PUBLIC_INPUT:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.air_public_input, $(COMPILED_PROOF_TESTS))
+
 CAIRO_RS_MEM_PROOF:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.rs.memory, $(COMPILED_PROOF_TESTS))
 CAIRO_RS_TRACE_PROOF:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.rs.trace, $(COMPILED_PROOF_TESTS))
+CAIRO_RS_AIR_PUBLIC_INPUT:=$(patsubst $(TEST_PROOF_DIR)/%.json, $(TEST_PROOF_DIR)/%.rs.air_public_input, $(COMPILED_PROOF_TESTS))
 
 PROOF_BENCH_DIR=cairo_programs/benchmarks
 PROOF_BENCH_FILES:=$(wildcard $(PROOF_BENCH_DIR)/*.cairo)
@@ -39,11 +43,11 @@ PROOF_COMPILED_BENCHES:=$(patsubst $(PROOF_BENCH_DIR)/%.cairo, $(PROOF_BENCH_DIR
 $(TEST_PROOF_DIR)/%.json: $(TEST_PROOF_DIR)/%.cairo
 	cairo-compile --cairo_path="$(TEST_PROOF_DIR):$(PROOF_BENCH_DIR)" $< --output $@ --proof_mode
 
-$(TEST_PROOF_DIR)/%.rs.trace $(TEST_PROOF_DIR)/%.rs.memory: $(TEST_PROOF_DIR)/%.json $(RELBIN)
-	cargo llvm-cov run -p cairo-vm-cli --release --no-report -- --layout starknet_with_keccak --proof_mode $< --trace_file $@ --memory_file $(@D)/$(*F).rs.memory
+$(TEST_PROOF_DIR)/%.rs.trace $(TEST_PROOF_DIR)/%.rs.memory $(TEST_PROOF_DIR)/%.rs.air_public_input: $(TEST_PROOF_DIR)/%.json $(RELBIN)
+	cargo llvm-cov run -p cairo-vm-cli --release --no-report -- --layout starknet_with_keccak --proof_mode $< --trace_file $@ --memory_file $(@D)/$(*F).rs.memory --air_public_input $(@D)/$(*F).rs.air_public_input
 
-$(TEST_PROOF_DIR)/%.trace $(TEST_PROOF_DIR)/%.memory: $(TEST_PROOF_DIR)/%.json
-	cairo-run --layout starknet_with_keccak --proof_mode --program $< --trace_file $@ --memory_file $(@D)/$(*F).memory
+$(TEST_PROOF_DIR)/%.trace $(TEST_PROOF_DIR)/%.memory $(TEST_PROOF_DIR)/%.air_public_input: $(TEST_PROOF_DIR)/%.json
+	cairo-run --layout starknet_with_keccak --proof_mode --program $< --trace_file $(@D)/$(*F).trace  --air_public_input $(@D)/$(*F).air_public_input --memory_file $(@D)/$(*F).memory
 
 $(PROOF_BENCH_DIR)/%.json: $(PROOF_BENCH_DIR)/%.cairo
 	cairo-compile --cairo_path="$(TEST_PROOF_DIR):$(PROOF_BENCH_DIR)" $< --output $@ --proof_mode
@@ -215,8 +219,8 @@ cairo_bench_programs: $(COMPILED_BENCHES)
 cairo_1_test_contracts: $(CAIRO_1_COMPILED_CASM_CONTRACTS)
 cairo_2_test_contracts: $(CAIRO_2_COMPILED_CASM_CONTRACTS)
 
-cairo_proof_trace: $(CAIRO_TRACE_PROOF) $(CAIRO_MEM_PROOF)
-cairo-vm_proof_trace: $(CAIRO_RS_TRACE_PROOF) $(CAIRO_RS_MEM_PROOF)
+cairo_proof_trace: $(CAIRO_TRACE_PROOF) $(CAIRO_MEM_PROOF) $(CAIRO_AIR_PUBLIC_INPUT)
+cairo-vm_proof_trace: $(CAIRO_RS_TRACE_PROOF) $(CAIRO_RS_MEM_PROOF) $(CAIRO_RS_AIR_PUBLIC_INPUT)
 
 cairo_trace: $(CAIRO_TRACE) $(CAIRO_MEM)
 cairo-vm_trace: $(CAIRO_RS_TRACE) $(CAIRO_RS_MEM)
@@ -271,12 +275,18 @@ compare_memory: $(CAIRO_RS_MEM) $(CAIRO_MEM)
 compare_trace_memory_proof: $(COMPILED_PROOF_TESTS) $(CAIRO_RS_TRACE_PROOF) $(CAIRO_TRACE_PROOF) $(CAIRO_RS_MEM_PROOF) $(CAIRO_MEM_PROOF)
 	cd vm/src/tests; ./compare_vm_state.sh trace memory proof_mode
 
+compare_all_proof: $(COMPILED_PROOF_TESTS) $(CAIRO_RS_TRACE_PROOF) $(CAIRO_TRACE_PROOF) $(CAIRO_RS_MEM_PROOF) $(CAIRO_MEM_PROOF) $(CAIRO_RS_AIR_PUBLIC_INPUT) $(CAIRO_AIR_PUBLIC_INPUT)
+	cd vm/src/tests; ./compare_vm_state.sh trace memory proof_mode air_public_input
+
 compare_trace_proof: $(CAIRO_RS_TRACE_PROOF) $(CAIRO_TRACE_PROOF)
 	cd vm/src/tests; ./compare_vm_state.sh trace proof_mode
 
 compare_memory_proof: $(CAIRO_RS_MEM_PROOF) $(CAIRO_MEM_PROOF)
 	cd vm/src/tests; ./compare_vm_state.sh memory proof_mode
 
+compare_air_public_input: $(CAIRO_RS_AIR_PUBLIC_INPUT) $(CAIRO_AIR_PUBLIC_INPUT)
+	cd vm/src/tests; ./compare_vm_state.sh memory proof_mode air_public_input
+
 # Run with nightly enable the `doc_cfg` feature wich let us provide clear explaination about which parts of the code are behind a feature flag
 docs:
 	RUSTDOCFLAGS="--cfg docsrs" cargo +nightly doc --verbose --release --locked --no-deps --all-features --open
@@ -292,6 +302,7 @@ clean:
 	rm -f $(TEST_PROOF_DIR)/*.json
 	rm -f $(TEST_PROOF_DIR)/*.memory
 	rm -f $(TEST_PROOF_DIR)/*.trace
+	rm -f $(TEST_PROOF_DIR)/*.air_public_input
 	rm -rf cairo-vm-env
 	rm -rf cairo-vm-pypy-env
 	rm -rf cairo

vm/src/air_public_input.rs

@@ -17,9 +17,9 @@ use crate::{
 #[derive(Serialize, Debug)]
 pub struct PublicMemoryEntry {
     address: usize,
-    page: usize,
     #[serde(serialize_with = "mem_value_serde::serialize")]
     value: Option<Felt252>,
+    page: usize,
 }
 
 mod mem_value_serde {
@@ -41,15 +41,15 @@ mod mem_value_serde {
 #[derive(Serialize, Debug)]
 struct MemorySegmentAddresses {
     begin_addr: usize,
-    stop_addr: usize,
+    stop_ptr: usize,
 }
 
 impl From<(usize, usize)> for MemorySegmentAddresses {
     fn from(addresses: (usize, usize)) -> Self {
-        let (begin_addr, stop_addr) = addresses;
+        let (begin_addr, stop_ptr) = addresses;
         MemorySegmentAddresses {
             begin_addr,
-            stop_addr,
+            stop_ptr,
         }
     }
 }

vm/src/tests/air_public_input_comparator.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import sys
+import json
+
+
+filename1 = sys.argv[1]
+filename2 = sys.argv[2]
+
+with open(filename1, 'r') as cairo_lang_input_file, open(filename2, 'r') as cairo_vm_input_file:
+    cairo_lang_input = json.load(cairo_lang_input_file)
+    cairo_vm_input = json.load(cairo_vm_input_file)
+
+    if cairo_lang_input == cairo_vm_input:
+
+        print(f"Comparison succesful for {filename1} vs {filename2}")
+    else:
+        print(f"Comparison unsuccesful for {filename1} vs {filename2}")
+        exit(1)
+

vm/src/tests/compare_vm_state.sh

@@ -8,6 +8,7 @@ proof_tests_path="${tests_path}/proof_programs"
 exit_code=0
 trace=false
 memory=false
+air_public_input=false
 passed_tests=0
 failed_tests=0
 
@@ -22,11 +23,21 @@ for i in $@; do
         "proof_mode") tests_path=$proof_tests_path
         echo "Requested proof mode usage"
         ;;
+        "air_public_input") air_public_input=true
+        echo "Requested air_public_input comparison"
+        ;;
         *)
         ;;
     esac
 done
 
+if $air_public_input; then
+    if [ $tests_path != $proof_tests_path ]; then
+        echo "Can't compare air_public_input without proof_mode"
+        exit 1
+    fi
+fi
+
 files=$(ls $tests_path)
 EXIT_CODE=$?
 if [ ${EXIT_CODE} != 0 ]; then
@@ -55,6 +66,16 @@ for file in $(ls $tests_path | grep .cairo$ | sed -E 's/\.cairo$//'); do
             passed_tests=$((passed_tests + 1))
         fi
     fi
+
+    if $air_public_input; then
+        if ! ./air_public_input_comparator.py $path_file.air_public_input $path_file.rs.air_public_input; then
+            echo "Air Public Input differs for $file"
+            exit_code=1
+            failed_tests=$((failed_tests + 1))
+        else
+            passed_tests=$((passed_tests + 1))
+        fi
+    fi
 done
 
 if test $failed_tests != 0; then

vm/src/vm/errors/memory_errors.rs

@@ -98,6 +98,8 @@ pub enum MemoryError {
     VecCapacityExceeded,
     #[error("Memory wasn't relocated")]
     UnrelocatedMemory,
+    #[error("Malformed public memory")]
+    MalformedPublicMemory,
 }
 
 #[derive(Debug, PartialEq, Eq, Error)]

vm/src/vm/runners/cairo_runner.rs

@@ -51,7 +51,7 @@ use num_traits::{ToPrimitive, Zero};
 use serde::Deserialize;
 
 use super::{
-    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner},
+    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner, OUTPUT_BUILTIN_NAME},
     cairo_pie::{self, CairoPie, CairoPieMetadata},
 };
 
@@ -938,8 +938,14 @@ impl CairoRunner {
             let (_, size) = builtin_runner
                 .get_used_cells_and_allocated_size(vm)
                 .map_err(RunnerError::FinalizeSegements)?;
-            vm.segments
-                .finalize(Some(size), builtin_runner.base(), None)
+            if builtin_runner.name() == OUTPUT_BUILTIN_NAME {
+                let public_memory = (0..size).map(|i| (i, 0)).collect();
+                vm.segments
+                    .finalize(Some(size), builtin_runner.base(), Some(&public_memory))
+            } else {
+                vm.segments
+                    .finalize(Some(size), builtin_runner.base(), None)
+            }
         }
         self.segments_finalized = true;
         Ok(())

vm/src/vm/vm_core.rs

@@ -996,7 +996,9 @@ impl VirtualMachine {
     /// Returns a list of addresses of memory cells that constitute the public memory.
     pub fn get_public_memory_addresses(&self) -> Result<Vec<(usize, usize)>, VirtualMachineError> {
         if let Some(relocation_table) = &self.relocation_table {
-            Ok(self.segments.get_public_memory_addresses(relocation_table))
+            self.segments
+                .get_public_memory_addresses(relocation_table)
+                .map_err(VirtualMachineError::Memory)
         } else {
             Err(MemoryError::UnrelocatedMemory.into())
         }
@@ -1028,7 +1030,10 @@ impl VirtualMachine {
                         return Err(RunnerError::NoStopPointer(Box::new(builtin.name())).into());
                     };
 
-                Ok((builtin.name(), relocate(addresses)?))
+                Ok((
+                    builtin.name().strip_suffix("_builtin").unwrap_or_default(),
+                    relocate(addresses)?,
+                ))
             })
             .collect()
     }

vm/src/vm/vm_memory/memory_segments.rs

@@ -215,21 +215,25 @@ impl MemorySegmentManager {
 
     /// Returns a list of addresses of memory cells that constitute the public memory.
     /// segment_offsets is the result of self.relocate_segments()
-    pub fn get_public_memory_addresses(&self, segment_offsets: &[usize]) -> Vec<(usize, usize)> {
-        let len = self.num_segments().min(self.public_memory_offsets.len());
-        let mut addresses = Vec::with_capacity(len);
-
-        for (offsets, segment_start) in self
-            .public_memory_offsets
-            .values()
-            .zip(segment_offsets.iter())
-            .take(len)
-        {
+    pub fn get_public_memory_addresses(
+        &self,
+        segment_offsets: &[usize],
+    ) -> Result<Vec<(usize, usize)>, MemoryError> {
+        let mut addresses = Vec::with_capacity(self.num_segments());
+        let empty_vec = vec![];
+        for segment_index in 0..self.num_segments() {
+            let offsets = &self
+                .public_memory_offsets
+                .get(&segment_index)
+                .unwrap_or(&empty_vec);
+            let segment_start = segment_offsets
+                .get(segment_index)
+                .ok_or(MemoryError::MalformedPublicMemory)?;
             for (offset, page_id) in offsets.iter() {
                 addresses.push((segment_start + offset, *page_id));
             }
         }
-        addresses
+        Ok(addresses)
     }
 
     // Writes the following information for the given segment:
@@ -245,10 +249,8 @@ impl MemorySegmentManager {
         if let Some(size) = size {
             self.segment_sizes.insert(segment_index, size);
         }
-        if let Some(public_memory) = public_memory {
-            self.public_memory_offsets
-                .insert(segment_index, public_memory.clone());
-        }
+        self.public_memory_offsets
+            .insert(segment_index, public_memory.cloned().unwrap_or_default());
     }
 }
 
@@ -828,12 +830,12 @@ mod tests {
 
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
-    fn finalize_no_size_nor_memory_no_change() {
+    fn finalize_no_size_nor_memory() {
         let mut segments = MemorySegmentManager::new();
         segments.finalize(None, 0, None);
         assert!(segments.memory.data.is_empty());
         assert!(segments.memory.temp_data.is_empty());
-        assert!(segments.public_memory_offsets.is_empty());
+        assert_eq!(segments.public_memory_offsets, HashMap::from([(0, vec![])]));
         assert_eq!(segments.num_segments(), 0);
         assert_eq!(segments.num_temp_segments(), 0);
     }
@@ -843,7 +845,7 @@ mod tests {
     fn finalize_no_memory() {
         let mut segments = MemorySegmentManager::new();
         segments.finalize(Some(42), 0, None);
-        assert!(segments.public_memory_offsets.is_empty());
+        assert_eq!(segments.public_memory_offsets, HashMap::from([(0, vec![])]));
         assert_eq!(segments.segment_sizes, HashMap::from([(0, 42)]));
     }