Convert relocation table to HashMap<usize, MaybeRelocatable>

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 #### Upcoming Changes
 
+* fix: [#1862](https://github.com/lambdaclass/cairo-vm/pull/1862):
+  * Use MaybeRelocatable for relocation table
+
 #### [2.0.0-rc1] - 2024-11-20
 
 * feat: add `EvalCircuit` and `TestLessThanOrEqualAddress` hints [#1843](https://github.com/lambdaclass/cairo-vm/pull/1843)

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)?;

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;
             }

vm/src/vm/vm_core.rs

@@ -1000,14 +1000,19 @@ 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.
     ///   - There shouldn't already be relocation at the source segment.
     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)
     }

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)];