Restrict addresses to Relocatable + fix some error variants used in signature.rs

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 #### Upcoming Changes
 
+* Restrict addresses to Relocatable + fix some error variants used in signature.rs [#792](https://github.com/lambdaclass/cairo-rs/pull/792)
+    * Public Api Changes:
+        * Change `ValidationRule` inner type to `Box<dyn Fn(&Memory, &Relocatable) -> Result<Vec<Relocatable>, MemoryError>>`.
+        * Change `validated_addresses` field of `Memory` to `HashSet<Relocatable>`.
+        * Change `validate_memory_cell(&mut self, address: &MaybeRelocatable) -> Result<(), MemoryError>` to `validate_memory_cell(&mut self, addr: &Relocatable) -> Result<(), MemoryError>`.
+
 * Add `VmException` to `CairoRunner::run_from_entrypoint`[#775](https://github.com/lambdaclass/cairo-rs/pull/775)
     * Public Api Changes:
         * Change error return type of `CairoRunner::run_from_entrypoint` to `CairoRunError`.

src/vm/errors/memory_errors.rs

@@ -62,4 +62,6 @@ pub enum MemoryError {
     ErrorVerifyingSignature,
     #[error("Couldn't obtain a mutable accessed offset")]
     CantGetMutAccessedOffset,
+    #[error("Failed to convert String: {0} to FieldElement")]
+    FailedStringToFieldElementConversion(String),
 }

src/vm/runners/builtin_runner/mod.rs

@@ -447,6 +447,7 @@ impl From<SignatureBuiltinRunner> for BuiltinRunner {
 mod tests {
     use super::*;
     use crate::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor;
+    use crate::relocatable;
     use crate::types::instance_definitions::ecdsa_instance_def::EcdsaInstanceDef;
     use crate::types::instance_definitions::keccak_instance_def::KeccakInstanceDef;
     use crate::types::program::Program;
@@ -1194,9 +1195,7 @@ mod tests {
             BitwiseBuiltinRunner::new(&BitwiseInstanceDef::default(), true).into();
 
         let mut vm = vm!();
-        vm.memory
-            .validated_addresses
-            .insert(mayberelocatable!(0, 2));
+        vm.memory.validated_addresses.insert(relocatable!(0, 2));
 
         vm.memory.data = vec![vec![
             mayberelocatable!(0, 0).into(),

src/vm/runners/builtin_runner/range_check.rs

@@ -86,9 +86,7 @@ impl RangeCheckBuiltinRunner {
 
     pub fn add_validation_rule(&self, memory: &mut Memory) -> Result<(), RunnerError> {
         let rule: ValidationRule = ValidationRule(Box::new(
-            |memory: &Memory,
-             address: &MaybeRelocatable|
-             -> Result<Vec<MaybeRelocatable>, MemoryError> {
+            |memory: &Memory, address: &Relocatable| -> Result<Vec<Relocatable>, MemoryError> {
                 if let MaybeRelocatable::Int(ref num) = memory
                     .get(address)?
                     .ok_or(MemoryError::FoundNonInt)?

src/vm/runners/builtin_runner/signature.rs

@@ -56,9 +56,9 @@ impl SignatureBuiltinRunner {
         let s_string = s.to_str_radix(10);
         let (r_felt, s_felt) = (
             FieldElement::from_dec_str(&r_string)
-                .map_err(|_| MemoryError::AddressNotRelocatable)?,
+                .map_err(|_| MemoryError::FailedStringToFieldElementConversion(r_string))?,
             FieldElement::from_dec_str(&s_string)
-                .map_err(|_| MemoryError::AddressNotRelocatable)?,
+                .map_err(|_| MemoryError::FailedStringToFieldElementConversion(s_string))?,
         );
 
         let signature = Signature {
@@ -100,13 +100,8 @@ impl SignatureBuiltinRunner {
         let signatures = Rc::clone(&self.signatures);
         let rule: ValidationRule = ValidationRule(Box::new(
             move |memory: &Memory,
-                  address: &MaybeRelocatable|
-                  -> Result<Vec<MaybeRelocatable>, MemoryError> {
-                let address = match address {
-                    MaybeRelocatable::RelocatableValue(address) => *address,
-                    _ => return Err(MemoryError::MissingAccessedAddresses),
-                };
-
+                  address: &Relocatable|
+                  -> Result<Vec<Relocatable>, MemoryError> {
                 let address_offset = address.offset.mod_floor(&(cells_per_instance as usize));
                 let mem_addr_sum = memory.get(&(address + 1_i32));
                 let mem_addr_less = if address.offset > 0 {
@@ -122,14 +117,14 @@ impl SignatureBuiltinRunner {
                     (0, Ok(Some(_element)), _) => {
                         let pubkey_addr = address;
                         let msg_addr = address + 1_i32;
-                        (pubkey_addr, msg_addr)
+                        (*pubkey_addr, msg_addr)
                     }
                     (1, _, Ok(Some(_element))) if address.offset > 0 => {
                         let pubkey_addr = address
                             .sub_usize(1)
                             .map_err(|_| MemoryError::EffectiveSizesNotCalled)?;
                         let msg_addr = address;
-                        (pubkey_addr, msg_addr)
+                        (pubkey_addr, *msg_addr)
                     }
                     _ => return Ok(Vec::new()),
                 };

src/vm/runners/cairo_runner.rs

@@ -1516,11 +1516,11 @@ mod tests {
         assert!(vm
             .memory
             .validated_addresses
-            .contains(&MaybeRelocatable::from((2, 0))));
+            .contains(&Relocatable::from((2, 0))));
         assert!(vm
             .memory
             .validated_addresses
-            .contains(&MaybeRelocatable::from((2, 1))));
+            .contains(&Relocatable::from((2, 1))));
         assert_eq!(vm.memory.validated_addresses.len(), 2);
     }
 

src/vm/vm_memory/memory.rs

@@ -12,7 +12,7 @@ use std::{
 
 pub struct ValidationRule(
     #[allow(clippy::type_complexity)]
-    pub  Box<dyn Fn(&Memory, &MaybeRelocatable) -> Result<Vec<MaybeRelocatable>, MemoryError>>,
+    pub  Box<dyn Fn(&Memory, &Relocatable) -> Result<Vec<Relocatable>, MemoryError>>,
 );
 
 pub struct Memory {
@@ -21,7 +21,7 @@ pub struct Memory {
     // 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...
     pub(crate) relocation_rules: HashMap<usize, Relocatable>,
-    pub validated_addresses: HashSet<MaybeRelocatable>,
+    pub validated_addresses: HashSet<Relocatable>,
     validation_rules: HashMap<usize, ValidationRule>,
 }
 
@@ -31,7 +31,7 @@ impl Memory {
             data: Vec::<Vec<Option<MaybeRelocatable>>>::new(),
             temp_data: Vec::<Vec<Option<MaybeRelocatable>>>::new(),
             relocation_rules: HashMap::new(),
-            validated_addresses: HashSet::<MaybeRelocatable>::new(),
+            validated_addresses: HashSet::<Relocatable>::new(),
             validation_rules: HashMap::new(),
         }
     }
@@ -81,7 +81,7 @@ impl Memory {
                 }
             }
         };
-        self.validate_memory_cell(&MaybeRelocatable::from(key))
+        self.validate_memory_cell(&relocatable)
     }
 
     /// Retrieve a value from memory (either normal or temporary) and apply relocation rules
@@ -249,27 +249,23 @@ impl Memory {
         self.validation_rules.insert(segment_index, rule);
     }
 
-    fn validate_memory_cell(&mut self, address: &MaybeRelocatable) -> Result<(), MemoryError> {
-        if let &MaybeRelocatable::RelocatableValue(ref rel_addr) = address {
-            if !self.validated_addresses.contains(address) {
-                for (index, validation_rule) in self.validation_rules.iter() {
-                    if rel_addr.segment_index == *index as isize {
-                        self.validated_addresses
-                            .extend(validation_rule.0(self, address)?);
-                    }
+    fn validate_memory_cell(&mut self, addr: &Relocatable) -> Result<(), MemoryError> {
+        if !self.validated_addresses.contains(addr) {
+            for (index, validation_rule) in self.validation_rules.iter() {
+                if addr.segment_index == *index as isize {
+                    self.validated_addresses
+                        .extend(validation_rule.0(self, addr)?);
                 }
             }
-            Ok(())
-        } else {
-            Err(MemoryError::AddressNotRelocatable)
         }
+        Ok(())
     }
     ///Applies validation_rules to the current memory
     //Should be called during initialization, as None values will raise a FoundNonInt error
     pub fn validate_existing_memory(&mut self) -> Result<(), MemoryError> {
         for i in 0..self.data.len() {
             for j in 0..self.data[i].len() {
-                self.validate_memory_cell(&MaybeRelocatable::from((i as isize, j)))?;
+                self.validate_memory_cell(&Relocatable::from((i as isize, j)))?;
             }
         }
         Ok(())
@@ -593,7 +589,7 @@ mod memory_tests {
         memory.validate_existing_memory().unwrap();
         assert!(memory
             .validated_addresses
-            .contains(&MaybeRelocatable::from((0, 0))));
+            .contains(&Relocatable::from((0, 0))));
     }
 
     #[test]