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_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("Segment {0} has already been accessed; it cannot be execution-protected.")]
+    SegmentAlreadyAccessed(usize),
     #[error(transparent)]
     Math(#[from] MathError),
     #[error(transparent)]

vm/src/vm/vm_core.rs

@@ -946,6 +946,14 @@ impl VirtualMachine {
         self.segments.memory.insert_value(key, val)
     }
 
+    /// Set the first cell of the given segment to a value that is not a valid opcode, to prevent
+    /// segment execution. Returns an error if the first cell of the segment is already accessed.
+    pub fn prevent_segment_execution(&mut self, segment_index: isize) -> Result<(), MemoryError> {
+        self.segments
+            .memory
+            .prevent_segment_execution(segment_index)
+    }
+
     ///Writes data into the memory from address ptr and returns the first address after the data.
     pub fn load_data(
         &mut self,
@@ -1359,6 +1367,7 @@ mod tests {
     use crate::types::instruction::OpcodeExtension;
     use crate::types::layout_name::LayoutName;
     use crate::types::program::Program;
+    use crate::vm::vm_memory::memory::MemoryCell;
     use crate::{
         any_box,
         hint_processor::builtin_hint_processor::builtin_hint_processor_definition::{
@@ -4727,6 +4736,48 @@ mod tests {
         assert_eq!(vm.segments.compute_effective_sizes(), &vec![4]);
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn test_prevent_segment_execution() {
+        let mut vm = vm!();
+
+        let segment0 = vm.segments.add();
+        let segment1 = vm.segments.add();
+        let segment2 = 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!(tmp_segment.segment_index, -1);
+        vm.segments.memory = memory![
+            ((0, 1), 1),
+            ((1, 0), 3),
+            ((1, 1), 4),
+            ((2, 0), 7),
+            ((-1, 0), 5)
+        ];
+
+        vm.mark_address_range_as_accessed((2, 0).into(), 1).unwrap();
+
+        vm.prevent_segment_execution(0).unwrap();
+        vm.prevent_segment_execution(1).unwrap();
+        vm.prevent_segment_execution(-1).unwrap();
+
+        let non_opcode = MemoryCell::new(MaybeRelocatable::Int(Felt252::from(-1)));
+
+        assert_eq!(vm.segments.memory.data[0][0], non_opcode);
+        assert_eq!(vm.segments.memory.data[1][0], non_opcode);
+        assert_eq!(vm.segments.memory.temp_data[0][0], non_opcode);
+        assert_matches!(
+            vm.prevent_segment_execution(2).unwrap_err(),
+            MemoryError::SegmentAlreadyAccessed(idx) if idx == 2
+        );
+        assert_matches!(
+            vm.prevent_segment_execution(7).unwrap_err(),
+            MemoryError::UnallocatedSegment(boxed) if *boxed == (7, 3)
+        );
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn mark_as_accessed() {

vm/src/vm/vm_memory/memory.rs

@@ -237,6 +237,37 @@ impl Memory {
         self.validate_memory_cell(key)
     }
 
+    /// Set the first cell of the given segment to a value that is not a valid opcode, to prevent
+    /// segment execution. Returns an error if the first cell of the segment is already accessed.
+    pub fn prevent_segment_execution(&mut self, segment_index: isize) -> Result<(), MemoryError> {
+        let key = Relocatable::from((segment_index, 0));
+        let data = if segment_index.is_negative() {
+            &mut self.temp_data
+        } else {
+            &mut self.data
+        };
+        let (segment_uindex, _) = from_relocatable_to_indexes(key);
+
+        let data_len = data.len();
+        let segment = data
+            .get_mut(segment_uindex)
+            .ok_or_else(|| MemoryError::UnallocatedSegment(Box::new((segment_uindex, data_len))))?;
+
+        // Check if the segment has at least one cell. If so, check that the first cell has not been
+        // accessed.
+        if !segment.is_empty() && segment[0].is_accessed() {
+            return Err(MemoryError::SegmentAlreadyAccessed(segment_uindex));
+        }
+
+        // Either the segment is empty, or the first cell is not accessed. Set the first cell to
+        // a value that is not a valid opcode (-1).
+        if segment.is_empty() {
+            segment.push(MemoryCell::NONE);
+        }
+        segment[0] = MemoryCell::new(MaybeRelocatable::Int(Felt252::from(-1)));
+        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