enable segment arena

lambdaclass · Apr 17, 2024 · bae0dd0 · bae0dd0
1 parent da0383d
commit bae0dd0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 26 deletions.
diff --git a/cairo1-run/src/cairo_run.rs b/cairo1-run/src/cairo_run.rs
@@ -380,29 +380,9 @@ fn create_entry_code(
         ap_offset += (1 + values.len()) as i16;
     }
     let mut array_args_data_iter = array_args_data.iter();
-    let after_arrays_data_offset = ap_offset;
     let mut arg_iter = config.args.iter().enumerate();
     let mut param_index = 0;
     let mut expected_arguments_size = 0;
-    if signature.param_types.iter().any(|ty| {
-        get_info(sierra_program_registry, ty)
-            .map(|x| x.long_id.generic_id == SegmentArenaType::ID)
-            .unwrap_or_default()
-    }) {
-        casm_extend! {ctx,
-            // SegmentArena segment.
-            %{ memory[ap + 0] = segments.add() %}
-            // Infos segment.
-            %{ memory[ap + 1] = segments.add() %}
-            ap += 2;
-            [ap + 0] = 0, ap++;
-            // Write Infos segment, n_constructed (0), and n_destructed (0) to the segment.
-            [ap - 2] = [[ap - 3]];
-            [ap - 1] = [[ap - 3] + 1];
-            [ap - 1] = [[ap - 3] + 2];
-        }
-        ap_offset += 3;
-    }
 
     for ty in &signature.param_types {
         let info = get_info(sierra_program_registry, ty)
@@ -420,10 +400,6 @@ fn create_entry_code(
         } else if generic_ty == &GasBuiltinType::ID {
             casm_extend!(ctx, [ap + 0] = initial_gas, ap++;);
             ap_offset += 1;
-        } else if generic_ty == &SegmentArenaType::ID {
-            let offset = -ap_offset + after_arrays_data_offset;
-            casm_extend!(ctx, [ap + 0] = [ap + offset] + 3, ap++;);
-            ap_offset += 1;
         } else {
             let ty_size = type_sizes[ty];
             let param_ap_offset_end = ap_offset + ty_size;

diff --git a/vm/src/vm/runners/builtin_runner/segment_arena.rs b/vm/src/vm/runners/builtin_runner/segment_arena.rs
@@ -46,7 +46,7 @@ impl SegmentArenaBuiltinRunner {
         if used < INITIAL_SEGMENT_SIZE {
             return Err(MemoryError::InvalidUsedSizeSegmentArena);
         }
-        Ok(used - INITIAL_SEGMENT_SIZE)
+        Ok(used)
     }
 
     pub fn initial_stack(&self) -> Vec<MaybeRelocatable> {
@@ -324,7 +324,7 @@ mod tests {
         let builtin = BuiltinRunner::SegmentArena(SegmentArenaBuiltinRunner::new(true));
         let mut memory_segment_manager = MemorySegmentManager::new();
         memory_segment_manager.segment_used_sizes = Some(vec![6]);
-        // (SIZE(6) - INITIAL_SIZE(3)) / CELLS_PER_INSTANCE(3)
+        // SIZE(6) / CELLS_PER_INSTANCE(3)
         assert_eq!(builtin.get_used_instances(&memory_segment_manager), Ok(1));
     }
 

diff --git a/vm/src/vm/runners/cairo_runner.rs b/vm/src/vm/runners/cairo_runner.rs
@@ -182,6 +182,14 @@ impl RunnerMode {
             RunnerMode::ProofModeCanonical | RunnerMode::ProofModeCairo1
         )
     }
+
+    /// Returns whether the runner mode is for Cairo 1 programs
+    pub fn is_cairo1(&self) -> bool {
+        matches!(
+            self,
+            RunnerMode::ExecutionModeCairo1 | RunnerMode::ProofModeCairo1
+        )
+    }
 }
 
 impl CairoRunner {
@@ -334,6 +342,13 @@ impl CairoRunner {
             }
         }
 
+        if self.runner_mode.is_cairo1() {
+            let included = program_builtins.remove(&BuiltinName::segment_arena);
+            if included || self.is_proof_mode() {
+                builtin_runners.push(SegmentArenaBuiltinRunner::new(included).into());
+            }
+        }
+
         if let Some(instance_def) = self.layout.builtins.keccak.as_ref() {
             let included = program_builtins.remove(&BuiltinName::keccak);
             if included || self.is_proof_mode() {