Remove unsafe impl of Add<usize> for &'a Relocatable

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* fix(BREAKING): Remove unsafe impl of `Add<usize> for &'a Relocatable`[#1718](https://github.com/lambdaclass/cairo-vm/pull/1718)
+
 * Bump `starknet-types-core` version + Use the lib's pedersen hash [#1692](https://github.com/lambdaclass/cairo-vm/pull/1692)
 
 * refactor: Remove unused code & use constants whenever possible for builtin instance definitions[#1707](https://github.com/lambdaclass/cairo-vm/pull/1707)

vm/src/types/relocatable.rs

@@ -356,17 +356,6 @@ impl MaybeRelocatable {
     }
 }
 
-impl<'a> Add<usize> for &'a Relocatable {
-    type Output = Relocatable;
-
-    fn add(self, other: usize) -> Self::Output {
-        Relocatable {
-            segment_index: self.segment_index,
-            offset: self.offset + other,
-        }
-    }
-}
-
 /// Turns a MaybeRelocatable into a Felt252 value.
 /// If the value is an Int, it will extract the Felt252 value from it.
 /// If the value is RelocatableValue, it will relocate it according to the relocation_table

vm/src/vm/runners/builtin_runner/signature.rs

@@ -183,8 +183,12 @@ impl SignatureBuiltinRunner {
     pub fn air_private_input(&self, memory: &Memory) -> Vec<PrivateInput> {
         let mut private_inputs = vec![];
         for (addr, signature) in self.signatures.borrow().iter() {
-            if let (Ok(pubkey), Ok(msg)) = (memory.get_integer(*addr), memory.get_integer(addr + 1))
-            {
+            if let (Ok(pubkey), Some(msg)) = (
+                memory.get_integer(*addr),
+                (*addr + 1_usize)
+                    .ok()
+                    .and_then(|addr| memory.get_integer(addr).ok()),
+            ) {
                 private_inputs.push(PrivateInput::Signature(PrivateInputSignature {
                     index: addr
                         .offset

vm/src/vm/runners/cairo_runner.rs

@@ -568,10 +568,7 @@ impl CairoRunner {
             } else {
                 let mut stack_prefix = vec![
                     Into::<MaybeRelocatable>::into(
-                        self.execution_base
-                            .as_ref()
-                            .ok_or(RunnerError::NoExecBase)?
-                            + target_offset,
+                        (self.execution_base.ok_or(RunnerError::NoExecBase)? + target_offset)?,
                     ),
                     MaybeRelocatable::from(Felt252::zero()),
                 ];
@@ -589,20 +586,16 @@ impl CairoRunner {
                 )?;
             }
 
-            self.initial_fp = Some(
-                self.execution_base
-                    .as_ref()
-                    .ok_or(RunnerError::NoExecBase)?
-                    + target_offset,
-            );
+            self.initial_fp =
+                Some((self.execution_base.ok_or(RunnerError::NoExecBase)? + target_offset)?);
 
             self.initial_ap = self.initial_fp;
-            return Ok(self.program_base.as_ref().ok_or(RunnerError::NoProgBase)?
+            return Ok((self.program_base.ok_or(RunnerError::NoProgBase)?
                 + self
                     .program
                     .shared_program_data
                     .end
-                    .ok_or(RunnerError::NoProgramEnd)?);
+                    .ok_or(RunnerError::NoProgramEnd)?)?);
         }
 
         let return_fp = vm.segments.add();

vm/src/vm/vm_memory/memory.rs

@@ -178,24 +178,23 @@
             &self.data
         };
         let (i, j) = from_relocatable_to_indexes(relocatable);
-        Some(self.relocate_value(data.get(i)?.get(j)?.as_ref()?.get_value()))
+        self.relocate_value(data.get(i)?.get(j)?.as_ref()?.get_value())
+            .ok()
     }
 
     // Version of Memory.relocate_value() that doesn't require a self reference
     fn relocate_address(
         addr: Relocatable,
         relocation_rules: &HashMap<usize, Relocatable>,
-    ) -> MaybeRelocatable {
-        let segment_idx = addr.segment_index;
-        if segment_idx >= 0 {
-            return addr.into();
-        }
-
-        // Adjust the segment index to begin at zero, as per the struct field's
-        match relocation_rules.get(&(-(segment_idx + 1) as usize)) {
-            Some(x) => (x + addr.offset).into(),
-            None => addr.into(),
+    ) -> 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((*x + addr.offset)?.into());
+            }
         }
+        Ok(addr.into())
     }
 
     /// Relocates the memory according to the relocation rules and clears `self.relocaction_rules`.
@@ -209,7 +208,7 @@
                 let value = cell.get_value_mut();
                 match value {
                     MaybeRelocatable::RelocatableValue(addr) if addr.segment_index < 0 => {
-                        *value = Memory::relocate_address(*addr, &self.relocation_rules);
+                        *value = Memory::relocate_address(*addr, &self.relocation_rules)?;
                     }
                     _ => {}
                 }
@@ -581,39 +580,42 @@
 
 /// Applies `relocation_rules` to a value
 pub(crate) trait RelocateValue<'a, Input: 'a, Output: 'a> {
-    fn relocate_value(&self, value: Input) -> Output;
+    fn relocate_value(&self, value: Input) -> Result<Output, MemoryError>;
 }
 
 impl RelocateValue<'_, Relocatable, Relocatable> for Memory {
-    fn relocate_value(&self, addr: Relocatable) -> Relocatable {
-        let segment_idx = addr.segment_index;
-        if segment_idx >= 0 {
-            return addr;
-        }
-
-        // Adjust the segment index to begin at zero, as per the struct field's
-        // comment.
-        match self.relocation_rules.get(&(-(segment_idx + 1) as usize)) {
-            Some(x) => x + addr.offset,
-            None => addr,
+    fn relocate_value(&self, addr: Relocatable) -> Result<Relocatable, 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 (*x + addr.offset).map_err(MemoryError::Math);
+            }
         }
+        Ok(addr)
     }
 }
 
 impl<'a> RelocateValue<'a, &'a Felt252, &'a Felt252> for Memory {
-    fn relocate_value(&self, value: &'a Felt252) -> &'a Felt252 {
-        value
+    fn relocate_value(&self, value: &'a Felt252) -> Result<&'a Felt252, MemoryError> {
+        Ok(value)
     }
 }
 
 impl<'a> RelocateValue<'a, &'a MaybeRelocatable, Cow<'a, MaybeRelocatable>> for Memory {
-    fn relocate_value(&self, value: &'a MaybeRelocatable) -> Cow<'a, MaybeRelocatable> {
-        match value {
+    fn relocate_value(
+        &self,
+        value: &'a MaybeRelocatable,
+    ) -> Result<Cow<'a, MaybeRelocatable>, MemoryError> {
+        Ok(match value {
             MaybeRelocatable::Int(_) => Cow::Borrowed(value),
             MaybeRelocatable::RelocatableValue(addr) => {
-                Cow::Owned(self.relocate_value(*addr).into())
+                Cow::Owned(self.relocate_value(*addr)?.into())
             }
-        }
+        })
     }
 }
 
@@ -1043,7 +1045,9 @@
 
         // Test when value is Some(BigInt):
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::Int(Felt252::from(0))),
+            memory
+                .relocate_value(&MaybeRelocatable::Int(Felt252::from(0)))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::Int(Felt252::from(0))),
         );
     }
@@ -1061,11 +1065,15 @@
 
         // Test when value is Some(MaybeRelocatable) with segment_index >= 0:
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((0, 0).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((0, 0).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((0, 0).into())),
         );
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((5, 0).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((5, 0).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((5, 0).into())),
         );
     }
@@ -1084,7 +1092,9 @@
         // Test when value is Some(MaybeRelocatable) with segment_index < 0 and
         // there are no applicable relocation rules:
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((-5, 0).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((-5, 0).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((-5, 0).into())),
         );
     }
@@ -1103,19 +1113,27 @@
         // Test when value is Some(MaybeRelocatable) with segment_index < 0 and
         // there are applicable relocation rules:
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((-1, 0).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((-1, 0).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((2, 0).into())),
         );
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((-2, 0).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((-2, 0).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((2, 2).into())),
         );
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((-1, 5).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((-1, 5).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((2, 5).into())),
         );
         assert_eq!(
-            memory.relocate_value(&MaybeRelocatable::RelocatableValue((-2, 5).into())),
+            memory
+                .relocate_value(&MaybeRelocatable::RelocatableValue((-2, 5).into()))
+                .unwrap(),
             Cow::Owned(MaybeRelocatable::RelocatableValue((2, 7).into())),
         );
     }
@@ -1498,11 +1516,11 @@
             .unwrap();
 
         assert_eq!(
-            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules),
+            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((2, 0).into()),
         );
         assert_eq!(
-            Memory::relocate_address((-2, 1).into(), &memory.relocation_rules),
+            Memory::relocate_address((-2, 1).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((2, 3).into()),
         );
     }
@@ -1512,11 +1530,11 @@
     fn relocate_address_no_rules() {
         let memory = Memory::new();
         assert_eq!(
-            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules),
+            Memory::relocate_address((-1, 0).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((-1, 0).into()),
         );
         assert_eq!(
-            Memory::relocate_address((-2, 1).into(), &memory.relocation_rules),
+            Memory::relocate_address((-2, 1).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((-2, 1).into()),
         );
     }
@@ -1526,11 +1544,11 @@
     fn relocate_address_real_addr() {
         let memory = Memory::new();
         assert_eq!(
-            Memory::relocate_address((1, 0).into(), &memory.relocation_rules),
+            Memory::relocate_address((1, 0).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((1, 0).into()),
         );
         assert_eq!(
-            Memory::relocate_address((1, 1).into(), &memory.relocation_rules),
+            Memory::relocate_address((1, 1).into(), &memory.relocation_rules).unwrap(),
             MaybeRelocatable::RelocatableValue((1, 1).into()),
         );
     }