Enable segment arena

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* fix(BREAKING): Enable segment arena builtin on cairo1-run.
+
 * fix(BREAKING): Remove unsafe impl of `Add<usize> for &'a Relocatable`[#1718](https://github.com/lambdaclass/cairo-vm/pull/1718)
 
 * fix(BREAKING): Handle triple dereference references[#1708](https://github.com/lambdaclass/cairo-vm/pull/1708)

cairo1-run/src/cairo_run.rs

@@ -38,7 +38,8 @@
         runners::{
             builtin_runner::{
                 BITWISE_BUILTIN_NAME, EC_OP_BUILTIN_NAME, HASH_BUILTIN_NAME, OUTPUT_BUILTIN_NAME,
-                POSEIDON_BUILTIN_NAME, RANGE_CHECK_BUILTIN_NAME, SIGNATURE_BUILTIN_NAME,
+                POSEIDON_BUILTIN_NAME, RANGE_CHECK_BUILTIN_NAME, SEGMENT_ARENA_BUILTIN_NAME,
+                SIGNATURE_BUILTIN_NAME,
             },
             cairo_runner::{CairoRunner, RunResources, RunnerMode},
         },
@@ -188,7 +189,7 @@
     let runner_mode = if cairo_run_config.proof_mode {
         RunnerMode::ProofModeCairo1
     } else {
-        RunnerMode::ExecutionMode
+        RunnerMode::ExecutionModeCairo1
     };
 
     let mut runner = CairoRunner::new_v2(&program, cairo_run_config.layout, runner_mode)?;
@@ -380,29 +381,9 @@
         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)
@@ -421,9 +402,7 @@
             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;
+            continue;
         } else {
             let ty_size = type_sizes[ty];
             let param_ap_offset_end = ap_offset + ty_size;
@@ -620,6 +599,11 @@
     let mut builtin_offset: HashMap<cairo_lang_sierra::ids::GenericTypeId, i16> = HashMap::new();
     let mut current_offset = 3;
     for (debug_name, builtin_name, sierra_id) in [
+        (
+            "SegmentArena",
+            BuiltinName::segment_arena,
+            SegmentArenaType::ID,
+        ),
         ("Poseidon", BuiltinName::poseidon, PoseidonType::ID),
         ("EcOp", BuiltinName::ec_op, EcOpType::ID),
         ("Bitwise", BuiltinName::bitwise, BitwiseType::ID),
@@ -729,6 +713,7 @@
                 "EcOp" => EC_OP_BUILTIN_NAME,
                 "Bitwise" => BITWISE_BUILTIN_NAME,
                 "Pedersen" => HASH_BUILTIN_NAME,
+                "SegmentArena" => SEGMENT_ARENA_BUILTIN_NAME,
                 "Output" => OUTPUT_BUILTIN_NAME,
                 "Ecdsa" => SIGNATURE_BUILTIN_NAME,
                 _ => {

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> {
@@ -119,7 +119,7 @@ mod tests {
         let mut vm = vm!();
         vm.segments.segment_used_sizes = Some(vec![3]);
 
-        assert_eq!(builtin.get_used_instances(&vm.segments), Ok(0));
+        assert_eq!(builtin.get_used_instances(&vm.segments), Ok(1));
     }
 
     #[test]
@@ -130,7 +130,7 @@ mod tests {
         let mut vm = vm!();
         vm.segments.segment_used_sizes = Some(vec![3]);
 
-        assert_eq!(builtin.get_used_instances(&vm.segments), Ok(0));
+        assert_eq!(builtin.get_used_instances(&vm.segments), Ok(1));
     }
 
     #[test]
@@ -154,7 +154,7 @@ mod tests {
             builtin.final_stack(&vm.segments, pointer),
             Err(RunnerError::InvalidStopPointer(Box::new((
                 SEGMENT_ARENA_BUILTIN_NAME,
-                relocatable!(0, 3),
+                relocatable!(0, 6),
                 relocatable!(0, 0)
             ))))
         );
@@ -245,7 +245,7 @@ mod tests {
 
         assert_eq!(
             builtin.get_used_cells_and_allocated_size(&vm),
-            Ok((0_usize, 0))
+            Ok((3_usize, 3))
         );
     }
 
@@ -303,7 +303,7 @@ mod tests {
         let mut vm = vm!();
 
         vm.segments.segment_used_sizes = Some(vec![4]);
-        assert_eq!(builtin.get_used_cells(&vm.segments), Ok(1));
+        assert_eq!(builtin.get_used_cells(&vm.segments), Ok(1 + 3));
     }
 
     #[test]
@@ -324,8 +324,8 @@ 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)
-        assert_eq!(builtin.get_used_instances(&memory_segment_manager), Ok(1));
+        // SIZE(6) / CELLS_PER_INSTANCE(3)
+        assert_eq!(builtin.get_used_instances(&memory_segment_manager), Ok(2));
     }
 
     #[test]

vm/src/vm/runners/cairo_runner.rs

@@ -164,11 +164,34 @@ pub struct CairoRunner {
 
 #[derive(Clone, Debug, PartialEq)]
 pub enum RunnerMode {
-    ExecutionMode,
+    /// Execution mode for Cairo 0 programs
+    ExecutionModeCairo0,
+    /// Execution mode for Cairo 1 programs
+    ExecutionModeCairo1,
+    /// Proof mode for Cairo 0 programs
     ProofModeCanonical,
+    /// Proof mode for Cairo 1 programs
     ProofModeCairo1,
 }
 
+impl RunnerMode {
+    /// Returns whether the runner mode is proof mode
+    pub fn is_proof_mode(&self) -> bool {
+        matches!(
+            self,
+            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 {
     pub fn new_v2(
         program: &Program,
@@ -208,7 +231,7 @@ impl CairoRunner {
             original_steps: None,
             relocated_memory: Vec::new(),
             exec_scopes: ExecutionScopes::new(),
-            execution_public_memory: if mode != RunnerMode::ExecutionMode {
+            execution_public_memory: if mode.is_proof_mode() {
                 Some(Vec::new())
             } else {
                 None
@@ -225,7 +248,7 @@ impl CairoRunner {
         if proof_mode {
             Self::new_v2(program, layout, RunnerMode::ProofModeCanonical)
         } else {
-            Self::new_v2(program, layout, RunnerMode::ExecutionMode)
+            Self::new_v2(program, layout, RunnerMode::ExecutionModeCairo0)
         }
     }
 
@@ -259,6 +282,7 @@ impl CairoRunner {
             BuiltinName::ec_op,
             BuiltinName::keccak,
             BuiltinName::poseidon,
+            BuiltinName::segment_arena,
             BuiltinName::range_check96,
             BuiltinName::add_mod,
             BuiltinName::mul_mod,
@@ -319,6 +343,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() {
@@ -367,8 +398,7 @@ impl CairoRunner {
     }
 
     fn is_proof_mode(&self) -> bool {
-        self.runner_mode == RunnerMode::ProofModeCanonical
-            || self.runner_mode == RunnerMode::ProofModeCairo1
+        self.runner_mode.is_proof_mode()
     }
 
     // Initialize all the builtins. Values used are the original one from the CairoFunctionRunner