Fix MemorySegmentManager::add_zero_segment

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 #### Upcoming Changes
 
+* fix(BREAKING): [#1818](https://github.com/lambdaclass/cairo-vm/pull/1818): 
+    * Fix `MemorySegmentManager::add_zero_segment` function when resizing a segment
+    * Signature change(BREAKING): `MemorySegmentManager::get_memory_holes` now receives `builtin_segment_indexes: HashSet<usize>`
+
 * fix(BREAKING): Replace `CairoRunner` method `initialize_all_builtins` with `initialize_program_builtins`. Now it only initializes program builtins instead of all of them
 
 #### [1.0.0] - 2024-08-01

vm/src/vm/runners/cairo_runner.rs

@@ -791,14 +791,16 @@ impl CairoRunner {
 
     /// Count the number of holes present in the segments.
     pub fn get_memory_holes(&self) -> Result<usize, MemoryError> {
-        let output_builtin_index = self
+        // Grab builtin segment indexes, except for the output builtin
+        let builtin_segment_indexes: HashSet<usize> = self
             .vm
             .builtin_runners
             .iter()
-            .position(|b| b.name() == BuiltinName::output);
-        self.vm
-            .segments
-            .get_memory_holes(self.vm.builtin_runners.len(), output_builtin_index)
+            .filter(|b| b.name() != BuiltinName::output)
+            .map(|b| b.base())
+            .collect();
+
+        self.vm.segments.get_memory_holes(builtin_segment_indexes)
     }
 
     /// Check if there are enough trace cells to fill the entire diluted checks.

vm/src/vm/vm_memory/memory_segments.rs

@@ -1,3 +1,4 @@
+use crate::stdlib::collections::HashSet;
 use core::cmp::max;
 use core::fmt;
 
@@ -200,27 +201,18 @@ impl MemorySegmentManager {
     }
 
     /// Counts the memory holes (aka unaccessed memory cells) in memory
-    /// Receives the amount of builtins in the vm and the position of the output builtin within the builtins list if present
+    /// # Parameters
+    /// - `builtin_segment_indexes`: Set representing the segments indexes of the builtins initialized in the VM, except for the output builtin.
     pub fn get_memory_holes(
         &self,
-        builtin_count: usize,
-        output_builtin_index: Option<usize>,
+        builtin_segment_indexes: HashSet<usize>,
     ) -> Result<usize, MemoryError> {
         let data = &self.memory.data;
         let mut memory_holes = 0;
-        let builtin_segments_start = 1; // program segment + execution segment
-        let builtin_segments_end = builtin_segments_start + builtin_count;
-        let output_segment_index =
-            output_builtin_index.map(|i| i + 2 /*program segment + execution segment*/);
-        // Count the memory holes for each segment by substracting the amount of accessed_addresses from the segment's size
-        // Segments without accesses addresses are not accounted for when counting memory holes
         for i in 0..data.len() {
             // Instead of marking all of the builtin segment's address as accessed, we just skip them when counting memory holes
             // Output builtin is extempt from this behaviour
-            if i > builtin_segments_start
-                && i <= builtin_segments_end
-                && !output_segment_index.is_some_and(|output_index| output_index == i)
-            {
+            if builtin_segment_indexes.contains(&i) {
                 continue;
             }
             let accessed_amount =
@@ -293,8 +285,9 @@ impl MemorySegmentManager {
         if self.zero_segment_index.is_zero() {
             self.zero_segment_index = self.add().segment_index as usize;
         }
+
         // Fil zero segment with zero values until size is reached
-        for _ in 0..self.zero_segment_size.saturating_sub(size) {
+        for _ in 0..(size.saturating_sub(self.zero_segment_size)) {
             // As zero_segment_index is only accessible to the segment manager
             // we can asume that it is always valid and index direcly into it
             self.memory.data[self.zero_segment_index].push(MemoryCell::new(Felt252::ZERO.into()))
@@ -690,7 +683,7 @@ mod tests {
             .memory
             .mark_as_accessed((0, 0).into());
         assert_eq!(
-            memory_segment_manager.get_memory_holes(0, None),
+            memory_segment_manager.get_memory_holes(HashSet::new()),
             Err(MemoryError::MissingSegmentUsedSizes),
         );
     }
@@ -707,7 +700,7 @@ mod tests {
                 .mark_as_accessed((0, i).into());
         }
         assert_eq!(
-            memory_segment_manager.get_memory_holes(0, None),
+            memory_segment_manager.get_memory_holes(HashSet::new()),
             Err(MemoryError::SegmentHasMoreAccessedAddressesThanSize(
                 Box::new((0, 3, 2))
             )),
@@ -719,15 +712,21 @@ mod tests {
     fn get_memory_holes_empty() {
         let mut memory_segment_manager = MemorySegmentManager::new();
         memory_segment_manager.segment_used_sizes = Some(Vec::new());
-        assert_eq!(memory_segment_manager.get_memory_holes(0, None), Ok(0),);
+        assert_eq!(
+            memory_segment_manager.get_memory_holes(HashSet::new()),
+            Ok(0),
+        );
     }
 
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn get_memory_holes_empty2() {
         let mut memory_segment_manager = MemorySegmentManager::new();
         memory_segment_manager.segment_used_sizes = Some(vec![4]);
-        assert_eq!(memory_segment_manager.get_memory_holes(0, None), Ok(0),);
+        assert_eq!(
+            memory_segment_manager.get_memory_holes(HashSet::new()),
+            Ok(0),
+        );
     }
 
     #[test]
@@ -750,7 +749,10 @@ mod tests {
                 .memory
                 .mark_as_accessed((0, i).into());
         }
-        assert_eq!(memory_segment_manager.get_memory_holes(0, None), Ok(2),);
+        assert_eq!(
+            memory_segment_manager.get_memory_holes(HashSet::new()),
+            Ok(2),
+        );
     }
 
     #[test]
@@ -775,7 +777,10 @@ mod tests {
                 .memory
                 .mark_as_accessed((0, i).into());
         }
-        assert_eq!(memory_segment_manager.get_memory_holes(0, None), Ok(7),);
+        assert_eq!(
+            memory_segment_manager.get_memory_holes(HashSet::new()),
+            Ok(7),
+        );
     }
 
     #[test]
@@ -1000,4 +1005,42 @@ mod tests {
             Ok(x) if x == mayberelocatable!(2, 0)
         );
     }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn test_add_zero_segment() {
+        // Create MemorySegmentManager with program and execution segments
+        let mut memory_segment_manager = MemorySegmentManager::new();
+        memory_segment_manager.add();
+        memory_segment_manager.add();
+
+        // Add zero segment
+        memory_segment_manager.add_zero_segment(3);
+        assert_eq!(memory_segment_manager.zero_segment_index, 2);
+        assert_eq!(memory_segment_manager.zero_segment_size, 3);
+        assert_eq!(
+            &memory_segment_manager.memory.data[2],
+            &Vec::from([
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0))
+            ])
+        );
+
+        // Resize zero segment
+        memory_segment_manager.add_zero_segment(5);
+        assert_eq!(memory_segment_manager.zero_segment_index, 2);
+        assert_eq!(memory_segment_manager.zero_segment_size, 5);
+
+        assert_eq!(
+            &memory_segment_manager.memory.data[2],
+            &Vec::from([
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0)),
+                MemoryCell::new(MaybeRelocatable::from(0))
+            ])
+        );
+    }
 }