Repository: lambdaclass/cairo-vm (language: rust)

Issue:
Memory operations should sometimes accept an Integer
**Describe the bug**

The pythonic implementation of some memory operations accepts either an integer or a relocatable in some isolated cases. These cases are used by `cairo-lang` and break the OS when not supported.

_This issue replaces #1856 which assumed that the downstream problem was caused by the way `cast(0, felt*)` was used. However, the behavior described there is consistent with the way the pythonic implementation works (see code below), so I will close it as "works as intended."_

There are two specific functions used in the OS that depend on support for integers, discussed below.

#### add_relocation_rule()

[Source](https://github.com/starkware-libs/cairo-lang/blob/a86e92bfde9c171c0856d7b46580c66e004922f3/src/starkware/cairo/lang/vm/memory_dict.py#L111), called from [here](https://github.com/starkware-libs/cairo-lang/blob/a86e92bfde9c171c0856d7b46580c66e004922f3/src/starkware/cairo/common/segments.cairo#L6).

While the python code clearly declares `dest_ptr: RelocatableValue`, neither function does any checks against it (except `assert src_ptr = dest_ptr;` in [the latter](https://github.com/starkware-libs/cairo-lang/blob/a86e92bfde9c171c0856d7b46580c66e004922f3/src/starkware/cairo/common/segments.cairo#L9C5-L9C31); note that this assert makes it impossible to work around the issue).

Consider the following code snippet which demonstrates that the relocation system works transparently when `dest_ptr` is a `felt` and not a `felt*`:

```cairo
from starkware.cairo.common.segments import relocate_segment

func main() {
    alloc_locals;

    local temp_segment: felt*;
    local temp_segment_2: felt*;

    let ptr_from_cast_0: felt* = cast(0, felt*);
    let ptr_from_cast_1: felt* = cast(1, felt*);

    %{
        ids.temp_segment = segments.add_temp_segment()
        ids.temp_segment_2 = segments.add_temp_segment()

        # the following prints "-1:0"
        print("temp_segment before relocation:", ids.temp_segment)

        # the following prints "-2:0"
        print("temp_segment_2 before relocation:", ids.temp_segment_2)

    %}

    relocate_segment(src_ptr=temp_segment, dest_ptr=ptr_from_cast_0);
    relocate_segment(src_ptr=temp_segment_2, dest_ptr=ptr_from_cast_1);

    %{
        # the following prints "0"
        print("temp_segment after relocation:", ids.temp_segment)

        # the following prints "1"
        print("temp_segment_2 after relocation:", ids.temp_segment_2)
    %}

    ret;
}
```

#### get_range()

[Source](https://github.com/starkware-libs/cairo-lang/blob/a86e92bfde9c171c0856d7b46580c66e004922f3/src/starkware/cairo/lang/vm/memory_dict.py#L237)

Similarly, this function is happy to operate on `addr` of type `felt` so long as `size == 0`. Unlike the `add_relocation_rule()` issue, this one may be easy to workaround.

Consider the following code snippet which clearly shows that `get_range` can be called as described above:

```cairo
func main() {
    alloc_locals;

    let ptr_from_cast_0: felt* = cast(0, felt*);
    let ptr_from_cast_1: felt* = cast(1, felt*);

    %{
        # note that ptr_from_cast_0 and ptr_from_cast_1 are both internally treated as felt
        # and not felt*:
        # the following prints "0"
        print("ptr_from_cast_0:", ids.ptr_from_cast_0)
        # the following prints "1"
        print("ptr_from_cast_1:", ids.ptr_from_cast_1)

        # get_range(0,0) works
        range = memory.get_range(ids.ptr_from_cast_0, 0)
        # the following prints "[]"
        print("get_range(0,0)", range)

        # get_range(1,0) works
        range = memory.get_range(ids.ptr_from_cast_1, 0)
        # the following prints "[]"
        print("get_range(1,0)", range)

        # get_range(0,1) fails
        # range = memory.get_range(ids.ptr_from_cast_0, 1)

        # get_range(1,1) fails
        # range = memory.get_range(ids.ptr_from_cast_1, 1)
    %}
}
```

**To Reproduce**
Compile (with `cairo-compile`) the above code snippets and run with `cairo-run`, e.g.:

```bash
cairo-compile test_cast.cairo --output test_cast_compiled.json
cairo-run --program=test_cast_compiled.json --print_output --print_info --relocate_prints
```

**Expected behavior**

Changing the underlying types used by `cairo-vm` to be `MaybeRelocatable` could fix this. It seems like a step backwards in sanity, but it does support the behavior that is required by the OS.

I'll propose a concrete suggestion (and I'm happy to implement it, although I'd love feedback as well):

* Modify the `relocation_rules` table [here](https://github.com/lambdaclass/cairo-vm/blob/159f67da19964cc54a95423a69470a26e534a13d/vm/src/vm/vm_memory/memory.rs#L164) to be a `HashMap<usize, MaybeRelocatable>`
* Keep existing accessors for adding a relocation rule (that is, the functions which accept `Relocatable`)
* Add new functions which accept a `MaybeRelocatable` and give them an obnoxious name like `fn add_relocation_rule_maybe_relocatable()`

**What version/commit are you on?**

git commit `58363ad8065ed891e3b14a8191b707677c7c7cb5b9d10030822506786d8d8108` (roughly `v1.0.1`)

**Additional context**

These requirements come from the OS code in [cairo-lang](https://github.com/starkware-libs/cairo-lang). They are very much corner cases as well (triggered in fewer than 1% of blocks on Sepolia testnet), which suggests that the behavior may not be intentional. It could be argued that these are bugs in the OS code... :shrug:



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e04a5fcb3..70c53d3317 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 #### Upcoming Changes
 
+* fix: [#1862](https://github.com/lambdaclass/cairo-vm/pull/1862):
+  * Use MaybeRelocatable for relocation table
+
 * chore: bump pip `cairo-lang` 0.13.3 [#1884](https://github.com/lambdaclass/cairo-vm/pull/1884)
 
 #### [2.0.0-rc1] - 2024-11-20
diff --git a/vm/src/hint_processor/builtin_hint_processor/segments.rs b/vm/src/hint_processor/builtin_hint_processor/segments.rs
index c996c2f9b1..bc95c6c1ba 100644
--- a/vm/src/hint_processor/builtin_hint_processor/segments.rs
+++ b/vm/src/hint_processor/builtin_hint_processor/segments.rs
@@ -19,7 +19,10 @@ 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)?;
 
     vm.add_relocation_rule(src_ptr, dest_ptr)
         .map_err(HintError::Memory)?;
diff --git a/vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs b/vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs
index a82a4fa2db..f3fe2fb9a0 100644
--- a/vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs
+++ b/vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs
@@ -126,7 +126,7 @@ impl DictManagerExecScope {
         if self.use_temporary_segments {
             let mut prev_end = vm.add_memory_segment();
             for tracker in &self.trackers {
-                vm.add_relocation_rule(tracker.start, prev_end)?;
+                vm.add_relocation_rule(tracker.start, prev_end.into())?;
                 prev_end += (tracker.end.unwrap_or_default() - tracker.start)?;
                 prev_end += 1;
             }
diff --git a/vm/src/vm/vm_core.rs b/vm/src/vm/vm_core.rs
index a65014a7b3..5862aee093 100644
--- a/vm/src/vm/vm_core.rs
+++ b/vm/src/vm/vm_core.rs
@@ -1000,6 +1000,10 @@ impl VirtualMachine {
 
     /// Add a new relocation rule.
     ///
+    /// When using feature "extensive_hints" the destination is allowed to be an Integer (via
+    /// MaybeRelocatable). Relocating memory to anything other than a `Relocatable` is generally
+    /// not useful, but it does make the implementation consistent with the pythonic version.
+    ///
     /// Will return an error if any of the following conditions are not met:
     ///   - Source address's segment must be negative (temporary).
     ///   - Source address's offset must be zero.
@@ -1007,7 +1011,8 @@ impl VirtualMachine {
     pub fn add_relocation_rule(
         &mut self,
         src_ptr: Relocatable,
-        dst_ptr: Relocatable,
+        #[cfg(not(feature = "extensive_hints"))] dst_ptr: Relocatable,
+        #[cfg(feature = "extensive_hints")] dst_ptr: MaybeRelocatable,
     ) -> Result<(), MemoryError> {
         self.segments.memory.add_relocation_rule(src_ptr, dst_ptr)
     }
diff --git a/vm/src/vm/vm_memory/memory.rs b/vm/src/vm/vm_memory/memory.rs
index 9dbfc9a790..affd0a06d4 100644
--- a/vm/src/vm/vm_memory/memory.rs
+++ b/vm/src/vm/vm_memory/memory.rs
@@ -161,7 +161,10 @@ pub struct Memory {
     pub(crate) temp_data: Vec<Vec<MemoryCell>>,
     // relocation_rules's keys map to temp_data's indices and therefore begin at
     // zero; that is, segment_index = -1 maps to key 0, -2 to key 1...
+    #[cfg(not(feature = "extensive_hints"))]
     pub(crate) relocation_rules: HashMap<usize, Relocatable>,
+    #[cfg(feature = "extensive_hints")]
+    pub(crate) relocation_rules: HashMap<usize, MaybeRelocatable>,
     pub validated_addresses: AddressSet,
     validation_rules: Vec<Option<ValidationRule>>,
 }
@@ -247,6 +250,7 @@ impl Memory {
     }
 
     // Version of Memory.relocate_value() that doesn't require a self reference
+    #[cfg(not(feature = "extensive_hints"))]
     fn relocate_address(
         addr: Relocatable,
         relocation_rules: &HashMap<usize, Relocatable>,
@@ -260,6 +264,23 @@ impl Memory {
         }
         Ok(addr.into())
     }
+    #[cfg(feature = "extensive_hints")]
+    fn relocate_address(
+        addr: Relocatable,
+        relocation_rules: &HashMap<usize, MaybeRelocatable>,
+    ) -> Result<MaybeRelocatable, MemoryError> {
+        if addr.segment_index < 0 {
+            // Adjust the segment index to begin at zero, as per the struct field's
+            // comment.
+            if let Some(x) = relocation_rules.get(&(-(addr.segment_index + 1) as usize)) {
+                return Ok(match x {
+                    MaybeRelocatable::RelocatableValue(r) => (*r + addr.offset)?.into(),
+                    MaybeRelocatable::Int(i) => i.into(),
+                });
+            }
+        }
+        Ok(addr.into())
+    }
 
     /// Relocates the memory according to the relocation rules and clears `self.relocaction_rules`.
     pub fn relocate_memory(&mut self) -> Result<(), MemoryError> {
@@ -289,6 +310,15 @@ impl Memory {
         for index in (0..self.temp_data.len()).rev() {
             if let Some(base_addr) = self.relocation_rules.get(&index) {
                 let data_segment = self.temp_data.remove(index);
+
+                #[cfg(feature = "extensive_hints")]
+                let base_addr = match base_addr {
+                    MaybeRelocatable::RelocatableValue(addr) => addr,
+                    MaybeRelocatable::Int(_) => {
+                        continue;
+                    }
+                };
+
                 // Insert the to-be relocated segment into the real memory
                 let mut addr = *base_addr;
                 if let Some(s) = self.data.get_mut(addr.segment_index as usize) {
@@ -310,13 +340,17 @@ impl Memory {
         self.relocation_rules.clear();
         Ok(())
     }
-
     /// Add a new relocation rule.
     ///
+    /// When using feature "extensive_hints" the destination is allowed to be an Integer (via
+    /// MaybeRelocatable). Relocating memory to anything other than a `Relocatable` is generally
+    /// not useful, but it does make the implementation consistent with the pythonic version.
+    ///
     /// Will return an error if any of the following conditions are not met:
     ///   - Source address's segment must be negative (temporary).
     ///   - Source address's offset must be zero.
     ///   - There shouldn't already be relocation at the source segment.
+    #[cfg(not(feature = "extensive_hints"))]
     pub(crate) fn add_relocation_rule(
         &mut self,
         src_ptr: Relocatable,
@@ -341,6 +375,31 @@ impl Memory {
         self.relocation_rules.insert(segment_index, dst_ptr);
         Ok(())
     }
+    #[cfg(feature = "extensive_hints")]
+    pub(crate) fn add_relocation_rule(
+        &mut self,
+        src_ptr: Relocatable,
+        dst: MaybeRelocatable,
+    ) -> Result<(), MemoryError> {
+        if src_ptr.segment_index >= 0 {
+            return Err(MemoryError::AddressNotInTemporarySegment(
+                src_ptr.segment_index,
+            ));
+        }
+        if src_ptr.offset != 0 {
+            return Err(MemoryError::NonZeroOffset(src_ptr.offset));
+        }
+
+        // Adjust the segment index to begin at zero, as per the struct field's
+        // comment.
+        let segment_index = -(src_ptr.segment_index + 1) as usize;
+        if self.relocation_rules.contains_key(&segment_index) {
+            return Err(MemoryError::DuplicatedRelocation(src_ptr.segment_index));
+        }
+
+        self.relocation_rules.insert(segment_index, dst);
+        Ok(())
+    }
 
     /// Gets the value from memory address as a Felt252 value.
     /// Returns an Error if the value at the memory address is missing or not a Felt252.
@@ -646,6 +705,7 @@ pub(crate) trait RelocateValue<'a, Input: 'a, Output: 'a> {
     fn relocate_value(&self, value: Input) -> Result<Output, MemoryError>;
 }
 
+#[cfg(not(feature = "extensive_hints"))]
 impl RelocateValue<'_, Relocatable, Relocatable> for Memory {
     fn relocate_value(&self, addr: Relocatable) -> Result<Relocatable, MemoryError> {
         if addr.segment_index < 0 {
@@ -661,6 +721,27 @@ impl RelocateValue<'_, Relocatable, Relocatable> for Memory {
         Ok(addr)
     }
 }
+#[cfg(feature = "extensive_hints")]
+impl RelocateValue<'_, Relocatable, MaybeRelocatable> for Memory {
+    fn relocate_value(&self, addr: Relocatable) -> Result<MaybeRelocatable, MemoryError> {
+        if addr.segment_index < 0 {
+            // Adjust the segment index to begin at zero, as per the struct field's
+            // comment.
+            if let Some(x) = self
+                .relocation_rules
+                .get(&(-(addr.segment_index + 1) as usize))
+            {
+                return Ok(match x {
+                    MaybeRelocatable::RelocatableValue(r) => {
+                        (*r + addr.offset).map_err(MemoryError::Math)?.into()
+                    }
+                    MaybeRelocatable::Int(i) => i.into(),
+                });
+            }
+        }
+        Ok(addr.into())
+    }
+}
 
 impl<'a> RelocateValue<'a, &'a Felt252, &'a Felt252> for Memory {
     fn relocate_value(&self, value: &'a Felt252) -> Result<&'a Felt252, MemoryError> {
@@ -676,7 +757,12 @@ impl<'a> RelocateValue<'a, &'a MaybeRelocatable, Cow<'a, MaybeRelocatable>> for
         Ok(match value {
             MaybeRelocatable::Int(_) => Cow::Borrowed(value),
             MaybeRelocatable::RelocatableValue(addr) => {
-                Cow::Owned(self.relocate_value(*addr)?.into())
+                #[cfg(not(feature = "extensive_hints"))]
+                let v = self.relocate_value(*addr)?.into();
+                #[cfg(feature = "extensive_hints")]
+                let v = self.relocate_value(*addr)?;
+
+                Cow::Owned(v)
             }
         })
     }
@@ -1617,6 +1703,64 @@ mod memory_tests {
         );
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    #[cfg(feature = "extensive_hints")]
+    fn relocate_address_to_integer() {
+        let mut memory = Memory::new();
+        memory
+            .add_relocation_rule((-1, 0).into(), 0.into())
+            .unwrap();
+        memory
+            .add_relocation_rule((-2, 0).into(), 42.into())
+            .unwrap();
+
+        assert_eq!(
+            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules).unwrap(),
+            MaybeRelocatable::Int(0.into()),
+        );
+        assert_eq!(
+            Memory::relocate_address((-2, 0).into(), &memory.relocation_rules).unwrap(),
+            MaybeRelocatable::Int(42.into()),
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    #[cfg(feature = "extensive_hints")]
+    fn relocate_address_integer_no_duplicates() {
+        let mut memory = Memory::new();
+        memory
+            .add_relocation_rule((-1, 0).into(), 1.into())
+            .unwrap();
+        assert_eq!(
+            memory.add_relocation_rule((-1, 0).into(), 42.into()),
+            Err(MemoryError::DuplicatedRelocation(-1))
+        );
+        assert_eq!(
+            memory.add_relocation_rule((-1, 0).into(), (2, 0).into()),
+            Err(MemoryError::DuplicatedRelocation(-1))
+        );
+
+        assert_eq!(
+            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules).unwrap(),
+            MaybeRelocatable::Int(1.into()),
+        );
+
+        memory
+            .add_relocation_rule((-2, 0).into(), (3, 0).into())
+            .unwrap();
+        assert_eq!(
+            memory.add_relocation_rule((-2, 0).into(), 1.into()),
+            Err(MemoryError::DuplicatedRelocation(-2))
+        );
+
+        assert_eq!(
+            Memory::relocate_address((-2, 0).into(), &memory.relocation_rules).unwrap(),
+            MaybeRelocatable::RelocatableValue((3, 0).into()),
+        );
+    }
+
     #[test]
     fn mark_address_as_accessed() {
         let mut memory = memory![((0, 0), 0)];


Test patch: