feat: implemented prevent_segment_execution function

Signed-off-by: Dori Medini <[email protected]>

Author: dorimedini-starkware

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: implemented prevent_contract_segment_execution function [#2099](https://github.com/lambdaclass/cairo-vm/pull/2099)
+
 #### [2.1.0] - 2025-05-21
 
 * chore: bump pip `cairo-lang` 0.13.5 [#1959](https://github.com/lambdaclass/cairo-vm/pull/1959)

vm/src/vm/errors/memory_errors.rs

@@ -15,6 +15,8 @@ use crate::types::{
 
 #[derive(Debug, PartialEq, Error)]
 pub enum MemoryError {
+    #[error("Cell {0} has already been accessed; it cannot be execution-protected.")]
+    CellAlreadyAccessed(Relocatable),
     #[error(transparent)]
     Math(#[from] MathError),
     #[error(transparent)]

vm/src/vm/vm_core.rs

@@ -946,6 +946,17 @@ impl VirtualMachine {
         self.segments.memory.insert_value(key, val)
     }
 
+    /// Unset (delete) the first cell of the given contract segment, to prevent contract segment
+    /// execution. Returns an error if the first cell of the contract segment is already accessed.
+    pub fn prevent_contract_segment_execution(
+        &mut self,
+        contract_segment_start: Relocatable,
+    ) -> Result<(), MemoryError> {
+        self.segments
+            .memory
+            .prevent_contract_segment_execution(contract_segment_start)
+    }
+
     ///Writes data into the memory from address ptr and returns the first address after the data.
     pub fn load_data(
         &mut self,
@@ -4727,6 +4738,80 @@ mod tests {
         assert_eq!(vm.segments.compute_effective_sizes(), &vec![4]);
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn test_prevent_contract_segment_execution() {
+        let mut vm = vm!();
+
+        let segment0 = vm.segments.add();
+        let segment1 = vm.segments.add();
+        let segment2 = vm.segments.add();
+        let segment3 = vm.segments.add();
+        let tmp_segment = vm.add_temporary_segment();
+        assert_eq!(segment0.segment_index, 0);
+        assert_eq!(segment1.segment_index, 1);
+        assert_eq!(segment2.segment_index, 2);
+        assert_eq!(segment3.segment_index, 3);
+        assert_eq!(tmp_segment.segment_index, -1);
+        vm.segments.memory = memory![
+            ((0, 1), 1),
+            ((1, 0), 3),
+            ((1, 1), 4),
+            ((2, 0), 7),
+            ((3, 0), 7),
+            ((-1, 0), 5),
+            ((-1, 1), 5),
+            ((-1, 2), 5)
+        ];
+        vm.run_finished = true;
+
+        vm.mark_address_range_as_accessed((2, 0).into(), 1).unwrap();
+
+        let cell0 = Relocatable::from((0, 0));
+        let cell1 = Relocatable::from((1, 1));
+        let cell2 = Relocatable::from((2, 0));
+        let cell3 = Relocatable::from((3, 7));
+        let cell7 = Relocatable::from((7, 17));
+        let cell_tmp = Relocatable::from((-1, 1));
+        vm.prevent_contract_segment_execution(cell0).unwrap();
+        vm.prevent_contract_segment_execution(cell1).unwrap();
+        vm.prevent_contract_segment_execution(cell_tmp).unwrap();
+        vm.prevent_contract_segment_execution(cell3).unwrap();
+
+        // Check that the cells were set to NONE.
+        assert!(vm
+            .segments
+            .memory
+            .get_cell_for_testing(cell0)
+            .unwrap()
+            .is_none());
+        assert!(vm
+            .segments
+            .memory
+            .get_cell_for_testing(cell1)
+            .unwrap()
+            .is_none());
+        assert!(vm
+            .segments
+            .memory
+            .get_cell_for_testing(cell_tmp)
+            .unwrap()
+            .is_none());
+        // Segment 3 cell was out of offset range, so it should not be modified or allocated.
+        assert!(vm.segments.memory.get_cell_for_testing(cell3).is_none());
+        // Segment 2 cell was accessed, so attempting to unset the memory should result in error.
+        assert_matches!(
+            vm.prevent_contract_segment_execution(cell2).unwrap_err(),
+            MemoryError::CellAlreadyAccessed(relocatable) if relocatable == cell2
+        );
+        // Segment 7 was not allocated, so attempting to unset the memory should result in error.
+        assert_matches!(
+            vm.prevent_contract_segment_execution(cell7).unwrap_err(),
+            MemoryError::UnallocatedSegment(boxed)
+            if *boxed == (cell7.segment_index.try_into().unwrap(), cell7.offset)
+        );
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn mark_as_accessed() {

vm/src/vm/vm_memory/memory.rs

@@ -237,6 +237,39 @@ impl Memory {
         self.validate_memory_cell(key)
     }
 
+    /// Unset (delete) the first cell of the given contract segment, to prevent contract segment
+    /// execution. Returns an error if the first cell of the contract segment is already accessed.
+    pub fn prevent_contract_segment_execution(
+        &mut self,
+        contract_segment_start: Relocatable,
+    ) -> Result<(), MemoryError> {
+        let data = if contract_segment_start.segment_index.is_negative() {
+            &mut self.temp_data
+        } else {
+            &mut self.data
+        };
+        let (index, offset) = from_relocatable_to_indexes(contract_segment_start);
+
+        let data_len = data.len();
+        let segment = data
+            .get_mut(index)
+            .ok_or_else(|| MemoryError::UnallocatedSegment(Box::new((index, data_len))))?;
+
+        // Check if the offset exists; if not, there is nothing to unset.
+        if offset >= segment.len() {
+            return Ok(());
+        }
+
+        // Ensure the cell has not been accessed.
+        if segment[offset].is_accessed() {
+            return Err(MemoryError::CellAlreadyAccessed(contract_segment_start));
+        }
+
+        // Unset the cell.
+        segment[offset] = MemoryCell::NONE;
+        Ok(())
+    }
+
     /// Retrieve a value from memory (either normal or temporary) and apply relocation rules
     pub(crate) fn get<'a, 'b: 'a, K: 'a>(&'b self, key: &'a K) -> Option<Cow<'b, MaybeRelocatable>>
     where
@@ -671,6 +704,11 @@ impl Memory {
         data.get(i)?.get(j)
     }
 
+    #[cfg(test)]
+    pub(crate) fn get_cell_for_testing(&self, addr: Relocatable) -> Option<&MemoryCell> {
+        self.get_cell(addr)
+    }
+
     pub fn is_accessed(&self, addr: &Relocatable) -> Result<bool, MemoryError> {
         Ok(self
             .get_cell(*addr)