[DNM] Segment arena relocation fix upstream

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* fix: Always use a normal segment in first SegmentArena segment [#1845](https://github.com/lambdaclass/cairo-vm/pull/1845)
+
 * chore: update cairo-lang dependencies to 2.12.0-dev.0 #[2040](https://github.com/lambdaclass/cairo-vm/pull/2040)
 
 * feat: add get_current_step getter [#2034](https://github.com/lambdaclass/cairo-vm/pull/2034)

Makefile

@@ -105,6 +105,10 @@ KZG_DA_CAIRO0_HINTS_DIR=cairo_programs/cairo-0-kzg-da-hints
 KZG_DA_CAIRO0_HINTS_FILES:=$(wildcard $(KZG_DA_CAIRO0_HINTS_DIR)/*.cairo)
 COMPILED_KZG_DA_CAIRO0_HINTS:=$(patsubst $(KZG_DA_CAIRO0_HINTS_DIR)/%.cairo, $(KZG_DA_CAIRO0_HINTS_DIR)/%.json, $(KZG_DA_CAIRO0_HINTS_FILES))
 
+SEGMENT_ARENA_CAIRO0_HINTS_DIR=cairo_programs/segment_arena
+SEGMENT_ARENA_CAIRO0_HINTS_FILES:=$(wildcard $(SEGMENT_ARENA_CAIRO0_HINTS_DIR)/*.cairo)
+COMPILED_SEGMENT_ARENA_CAIRO0_HINTS:=$(patsubst $(SEGMENT_ARENA_CAIRO0_HINTS_DIR)/%.cairo, $(SEGMENT_ARENA_CAIRO0_HINTS_DIR)/%.json, $(SEGMENT_ARENA_CAIRO0_HINTS_FILES))
+
 PRINT_TEST_DIR=cairo_programs/print_feature
 PRINT_TEST_FILES:=$(wildcard $(PRINT_TEST_DIR)/*.cairo)
 COMPILED_PRINT_TESTS:=$(patsubst $(PRINT_TEST_DIR)/%.cairo, $(PRINT_TEST_DIR)/%.json, $(PRINT_TEST_FILES))
@@ -254,7 +258,7 @@ run:
 check:
 	cargo check
 
-cairo_test_programs: $(COMPILED_TESTS) $(COMPILED_BAD_TESTS) $(COMPILED_NORETROCOMPAT_TESTS) $(COMPILED_PRINT_TESTS) $(COMPILED_MOD_BUILTIN_TESTS) $(COMPILED_SECP_CAIRO0_HINTS) $(COMPILED_KZG_DA_CAIRO0_HINTS)
+cairo_test_programs: $(COMPILED_TESTS) $(COMPILED_BAD_TESTS) $(COMPILED_NORETROCOMPAT_TESTS) $(COMPILED_PRINT_TESTS) $(COMPILED_MOD_BUILTIN_TESTS) $(COMPILED_SECP_CAIRO0_HINTS) $(COMPILED_KZG_DA_CAIRO0_HINTS) $(COMPILED_SEGMENT_ARENA_CAIRO0_HINTS)
 cairo_proof_programs: $(COMPILED_PROOF_TESTS) $(COMPILED_MOD_BUILTIN_PROOF_TESTS) $(COMPILED_STWO_EXCLUSIVE_TESTS)
 cairo_bench_programs: $(COMPILED_BENCHES)
 cairo_1_test_contracts: $(CAIRO_1_COMPILED_CASM_CONTRACTS)

cairo_programs/segment_arena/test_segment_arena.cairo

@@ -0,0 +1,172 @@
+// This test contains parts imported from cairo-lang
+// https://github.com/starkware-libs/cairo-lang/blob/master/src/starkware/starknet/builtins/segment_arena/segment_arena_test.cairo
+
+from starkware.cairo.common.alloc import alloc
+
+// Represents the information about a single segment allocated by the arena.
+struct SegmentInfo {
+    // A pointer to the first element of this segment.
+    start: felt*,
+    // A pointer to the end of this segment (the first unused element).
+    end: felt*,
+    // A sequential id, assigned to the segment when it is finalized.
+    // This value is used to guarantee that 'end' is not assigned twice.
+    finalization_index: felt,
+}
+
+// Represents the status of the segment arena.
+struct SegmentArenaBuiltin {
+    // A pointer to a list of SegmentInfo. infos[i] contains information about the i-th segment
+    // (ordered by construction).
+    // The value is fixed during the execution of an entry point.
+    infos: SegmentInfo*,
+    // The number of segments that were created so far.
+    n_segments: felt,
+    // The number of segments that were finalized so far.
+    n_finalized: felt,
+}
+
+// Constructs a new segment for the segment arena builtin and initializes it with an empty instance
+// of `SegmentArenaBuiltin`.
+func new_arena() -> SegmentArenaBuiltin* {
+    let (segment_arena: SegmentArenaBuiltin*) = alloc();
+    assert segment_arena[0] = SegmentArenaBuiltin(
+        infos=cast(nondet %{ segments.add() %}, SegmentInfo*), n_segments=0, n_finalized=0
+    );
+    return &segment_arena[1];
+}
+
+// Validates the segment arena builtin.
+//
+// In particular, relocates the temporary segments such that the start of segment i is strictly
+// larger than the end of segment i+1.
+func validate_segment_arena(segment_arena: SegmentArenaBuiltin*) {
+    tempvar n_segments = segment_arena.n_segments;
+    tempvar n_finalized = segment_arena.n_finalized;
+    // The following line should follow from the fact that every allocated segment
+    // must be finalized exactly once.
+    // We keep it both as a sanity check and since Sierra compilation is not proven yet.
+    assert n_segments = n_finalized;
+
+    if (n_segments == 0) {
+        return ();
+    }
+
+    // The following call also implies that n_segments > 0.
+    _verify_continuity(infos=segment_arena.infos, n_segments_minus_one=n_segments - 1);
+    return ();
+}
+
+// Helper function for validate_segment_arena.
+func _verify_continuity(infos: SegmentInfo*, n_segments_minus_one: felt) {
+    if (n_segments_minus_one == 0) {
+        // If there is only one segment left, there is no need to check anything.
+        return ();
+    }
+
+    // Enforce an empty cell between two consecutive segments so that the start of a segment
+    // is strictly bigger than the end of the previous segment.
+    // This is required for proving the soundness of this construction, in the case where a segment
+    // has length zero.
+
+    // Note: the following code was copied from relocate_segment() for efficiency reasons.
+    let src_ptr = infos[1].start;
+    let dest_ptr = infos[0].end + 1;
+    %{ memory.add_relocation_rule(src_ptr=ids.src_ptr, dest_ptr=ids.dest_ptr) %}
+    assert src_ptr = dest_ptr;
+
+    return _verify_continuity(infos=&infos[1], n_segments_minus_one=n_segments_minus_one - 1);
+}
+
+
+// Creates a new segment using the segment arena.
+func new_segment{segment_arena: SegmentArenaBuiltin*}() -> felt* {
+    let prev_segment_arena = &segment_arena[-1];
+    tempvar n_segments = prev_segment_arena.n_segments;
+    tempvar infos = prev_segment_arena.infos;
+
+    %{
+        if 'segment_index_to_arena_index' not in globals():
+            # A map from the relocatable value segment index to the index in the arena.
+            segment_index_to_arena_index = {}
+
+        # The segment is placed at the end of the arena.
+        index = ids.n_segments
+
+        # Create a segment or a temporary segment.
+        start = segments.add_temp_segment() if index > 0 else segments.add()
+
+        # Update 'SegmentInfo::start' and 'segment_index_to_arena_index'.
+        ids.prev_segment_arena.infos[index].start = start
+        segment_index_to_arena_index[start.segment_index] = index
+    %}
+    assert segment_arena[0] = SegmentArenaBuiltin(
+        infos=infos, n_segments=n_segments + 1, n_finalized=prev_segment_arena.n_finalized
+    );
+    let segment_arena = &segment_arena[1];
+    return infos[n_segments].start;
+}
+
+// Finalizes a given segment and returns the corresponding start.
+func finalize_segment{segment_arena: SegmentArenaBuiltin*}(segment_end: felt*) -> felt* {
+    let prev_segment_arena = &segment_arena[-1];
+    tempvar n_segments = prev_segment_arena.n_segments;
+    tempvar n_finalized = prev_segment_arena.n_finalized;
+
+    // Guess the index of the segment.
+    tempvar index = nondet %{ segment_index_to_arena_index[ids.segment_end.segment_index] %};
+
+    // Write segment_end in the manager.
+    tempvar infos: SegmentInfo* = prev_segment_arena.infos;
+    tempvar segment_info: SegmentInfo* = &infos[index];
+    // Writing n_finalized to 'finalization_index' guarantees 'segment_info.end' was not assigned
+    // a value before.
+    assert segment_info.finalization_index = n_finalized;
+    assert segment_info.end = segment_end;
+
+    assert segment_arena[0] = SegmentArenaBuiltin(
+        infos=infos, n_segments=n_segments, n_finalized=n_finalized + 1
+    );
+
+    let segment_arena = &segment_arena[1];
+    return segment_info.start;
+}
+
+func test_segment_arena() -> (felt*, SegmentInfo*) {
+    alloc_locals;
+
+    local segment_arena_start: SegmentArenaBuiltin* = new_arena();
+    let segment_arena = segment_arena_start;
+
+    with segment_arena {
+        let segment0 = new_segment();
+        let segment1 = new_segment();
+        let segment2 = new_segment();
+
+        assert segment0[0] = 1;
+        assert segment0[1] = 2;
+
+        assert segment1[0] = 3;
+        assert segment1[1] = 4;
+
+        assert segment2[0] = 5;
+
+        assert finalize_segment(segment0 + 2) = segment0;
+        assert finalize_segment(segment1 + 2) = segment1;
+
+        let segment3 = new_segment();
+
+        assert segment3[0] = 6;
+        assert segment3[1] = 7;
+
+        assert finalize_segment(segment3 + 2) = segment3;
+        assert finalize_segment(segment2 + 1) = segment2;
+    }
+    validate_segment_arena(segment_arena=&segment_arena[-1]);
+    return (segment0, segment_arena_start[-1].infos);
+}
+
+func main{}() {
+    let (_segment0, _infos) = test_segment_arena();
+    ret;
+}

vm/src/hint_processor/builtin_hint_processor/segments.rs

@@ -19,11 +19,35 @@ pub fn relocate_segment(
     ap_tracking: &ApTracking,
 ) -> Result<(), HintError> {
     let src_ptr = get_ptr_from_var_name("src_ptr", vm, ids_data, ap_tracking)?;
-    #[cfg(not(feature = "extensive_hints"))]
-    let dest_ptr = get_ptr_from_var_name("dest_ptr", vm, ids_data, ap_tracking)?;
-    #[cfg(feature = "extensive_hints")]
-    let dest_ptr = crate::hint_processor::builtin_hint_processor::hint_utils::get_maybe_relocatable_from_var_name("dest_ptr", vm, ids_data, ap_tracking)?;
-
+    // Bugfix: this is a workaround for an issue in the vm related to the way it computes
+    // `dest_ptr` in `segment_arena.cairo`. For some reason it sets `dest_ptr` to 0 (felt) instead
+    // of a valid relocatable value.
+    // As this hint is used in other places, we need to determine if we are in `segments_arena`
+    // first. We check this by determining if the "infos" variable is declared.
+    // If it is the case, we simply recompute `dest_ptr` manually.
+    let dest_ptr = if let Ok(infos) = get_ptr_from_var_name("infos", vm, ids_data, ap_tracking) {
+        let infos_0_end = vm.get_relocatable((infos + 1)?)?;
+
+        #[cfg(not(feature = "extensive_hints"))]
+        {
+            (infos_0_end + 1u32)?
+        }
+
+        #[cfg(feature = "extensive_hints")]
+        {
+            crate::types::relocatable::MaybeRelocatable::RelocatableValue((infos_0_end + 1u32)?)
+        }
+    } else {
+        #[cfg(not(feature = "extensive_hints"))]
+        {
+            get_ptr_from_var_name("dest_ptr", vm, ids_data, ap_tracking)?
+        }
+
+        #[cfg(feature = "extensive_hints")]
+        {
+            crate::hint_processor::builtin_hint_processor::hint_utils::get_maybe_relocatable_from_var_name("dest_ptr", vm, ids_data, ap_tracking)?
+        }
+    };
     vm.add_relocation_rule(src_ptr, dest_ptr)
         .map_err(HintError::Memory)?;
     Ok(())

vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs

@@ -57,7 +57,7 @@ impl DictManagerExecScope {
 
     /// Allocates a new segment for a new dictionary and return the start of the segment.
     pub fn new_default_dict(&mut self, vm: &mut VirtualMachine) -> Result<Relocatable, HintError> {
-        let dict_segment = if self.use_temporary_segments {
+        let dict_segment = if self.use_temporary_segments && !self.trackers.is_empty() {
             vm.add_temporary_segment()
         } else {
             vm.add_memory_segment()
@@ -123,9 +123,28 @@ impl DictManagerExecScope {
     /// Relocates all dictionaries into a single segment
     /// Does nothing if use_temporary_segments is set to false
     pub fn relocate_all_dictionaries(&mut self, vm: &mut VirtualMachine) -> Result<(), HintError> {
-        if self.use_temporary_segments {
-            let mut prev_end = vm.add_memory_segment();
-            for tracker in &self.trackers {
+        // We expect the first segment to be a normal one, which doesn't require relocation. So
+        // there is nothing to do unless there are at least two segments.
+        if self.use_temporary_segments && !self.trackers.is_empty() {
+            let first_segment = self.trackers.first().ok_or(HintError::CustomHint(
+                "Trackers must have a first element".into(),
+            ))?;
+            if first_segment.start.segment_index < 0 {
+                return Err(HintError::CustomHint(
+                    "First dict segment should not be temporary"
+                        .to_string()
+                        .into_boxed_str(),
+                ));
+            }
+            let mut prev_end = first_segment.end.unwrap_or_default();
+            for tracker in &self.trackers[1..] {
+                if tracker.start.segment_index >= 0 {
+                    return Err(HintError::CustomHint(
+                        "Dict segment should be temporary"
+                            .to_string()
+                            .into_boxed_str(),
+                    ));
+                }
                 #[cfg(feature = "extensive_hints")]
                 {
                     vm.add_relocation_rule(
@@ -220,3 +239,77 @@ impl DictSquashExecScope {
         self.current_access_indices()?.pop()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::stdlib::collections::HashMap;
+    use crate::types::relocatable::Relocatable;
+    use crate::vm::vm_core::VirtualMachine;
+
+    /// Test for relocate_all_dictionaries error cases
+    #[test]
+    fn test_relocate_all_dictionaries_errors() {
+        let mut vm = VirtualMachine::new(false, false);
+
+        // Test 1: First segment is a temporary segment (should error)
+        {
+            let mut dict_manager = DictManagerExecScope::new(true);
+            let first_dict_start = Relocatable::from((-1, 0)); // Temporary segment
+
+            dict_manager.trackers.push(DictTrackerExecScope {
+                data: HashMap::default(),
+                start: first_dict_start,
+                end: Some(Relocatable::from((-1, 10))),
+            });
+
+            let result = dict_manager.relocate_all_dictionaries(&mut vm);
+            assert!(matches!(
+                result,
+                Err(HintError::CustomHint(_)) if result.unwrap_err().to_string().contains("First dict segment should not be temporary")
+            ));
+        }
+
+        // Test 2: Non-temporary dictionary segment
+        {
+            let mut dict_manager = DictManagerExecScope::new(true);
+            let first_dict_start = Relocatable::from((0, 0)); // Non-temporary segment
+            let second_dict_start = Relocatable::from((1, 0)); // Non-temporary segment
+
+            dict_manager.trackers.push(DictTrackerExecScope {
+                data: HashMap::default(),
+                start: first_dict_start,
+                end: Some(Relocatable::from((0, 10))),
+            });
+            dict_manager.trackers.push(DictTrackerExecScope {
+                data: HashMap::default(),
+                start: second_dict_start,
+                end: Some(Relocatable::from((1, 10))),
+            });
+
+            let result = dict_manager.relocate_all_dictionaries(&mut vm);
+            assert!(matches!(
+                result,
+                Err(HintError::CustomHint(_)) if result.unwrap_err().to_string().contains("Dict segment should be temporary")
+            ));
+        }
+    }
+
+    /// Test for relocate_all_dictionaries when no temporary segments
+    #[test]
+    fn test_relocate_all_dictionaries_no_temporary_segments() {
+        let mut vm = VirtualMachine::new(false, false);
+        let mut dict_manager = DictManagerExecScope::new(false);
+
+        // Adding some trackers should not cause any errors
+        dict_manager.trackers.push(DictTrackerExecScope {
+            data: HashMap::default(),
+            start: Relocatable::from((0, 0)),
+            end: Some(Relocatable::from((0, 10))),
+        });
+
+        // Should not error and essentially do nothing
+        let result = dict_manager.relocate_all_dictionaries(&mut vm);
+        assert!(result.is_ok());
+    }
+}

vm/src/tests/cairo_run_test.rs

@@ -1,4 +1,5 @@
 use crate::{tests::*, types::layout_name::LayoutName};
+
 #[cfg(feature = "mod_builtin")]
 use crate::{
     utils::test_utils::Program,

vm/src/tests/mod.rs

@@ -36,6 +36,7 @@ mod run_deprecated_contract_class_simplified;
 mod cairo_1_run_from_entrypoint_tests;
 mod cairo_run_test;
 mod pedersen_test;
+mod segment_arena_test;
 mod struct_test;
 
 mod cairo_pie_test;