From 0336b707e4784d53d057b340afd2124a056e43ae Mon Sep 17 00:00:00 2001
From: Ori Ziv <oriz@starkware.co>
Date: Sun, 10 Nov 2024 20:33:49 +0200
Subject: [PATCH] Preventing temporary segments generation for testing.

commit-id:c7159ee0
---
 .../src/casm_run/dict_manager.rs                 |  8 ++++++--
 crates/cairo-lang-runner/src/casm_run/mod.rs     | 16 +++++++++++++---
 crates/cairo-lang-runner/src/casm_run/test.rs    |  2 ++
 crates/cairo-lang-runner/src/lib.rs              |  1 +
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/crates/cairo-lang-runner/src/casm_run/dict_manager.rs b/crates/cairo-lang-runner/src/casm_run/dict_manager.rs
index 69aa5bd8ce6..f1e2e9cf71c 100644
--- a/crates/cairo-lang-runner/src/casm_run/dict_manager.rs
+++ b/crates/cairo-lang-runner/src/casm_run/dict_manager.rs
@@ -31,10 +31,14 @@ impl DictManagerExecScope {
     pub const DICT_DEFAULT_VALUE: usize = 0;
 
     /// Allocates a new segment for a new dictionary and return the start of the segment.
-    pub fn new_default_dict(&mut self, vm: &mut VirtualMachine) -> Relocatable {
+    pub fn new_default_dict(
+        &mut self,
+        vm: &mut VirtualMachine,
+        no_temporary_segments: bool,
+    ) -> Relocatable {
         // If we are not on the first segment - using a temporary segments to later be merged into
         // the previous segments.
-        let dict_segment = if self.trackers.is_empty() {
+        let dict_segment = if self.trackers.is_empty() || no_temporary_segments {
             vm.add_memory_segment()
         } else {
             vm.add_temporary_segment()
diff --git a/crates/cairo-lang-runner/src/casm_run/mod.rs b/crates/cairo-lang-runner/src/casm_run/mod.rs
index cb4bc96ba84..111c2dfa58e 100644
--- a/crates/cairo-lang-runner/src/casm_run/mod.rs
+++ b/crates/cairo-lang-runner/src/casm_run/mod.rs
@@ -100,6 +100,8 @@ pub struct CairoHintProcessor<'a> {
     /// Resources used during syscalls - does not include resources used during the current VM run.
     /// At the end of the run - adding both would result in the actual expected resource usage.
     pub syscalls_used_resources: StarknetExecutionResources,
+    /// Avoid allocating memory segments so finalization of segment arena may not occur.
+    pub no_temporary_segments: bool,
 }
 
 pub fn cell_ref_to_relocatable(cell_ref: &CellRef, vm: &VirtualMachine) -> Relocatable {
@@ -421,7 +423,12 @@ impl HintProcessorLogic for CairoHintProcessor<'_> {
         let hint = match hint {
             Hint::Starknet(hint) => hint,
             Hint::Core(core_hint_base) => {
-                return execute_core_hint_base(vm, exec_scopes, core_hint_base);
+                return execute_core_hint_base(
+                    vm,
+                    exec_scopes,
+                    core_hint_base,
+                    self.no_temporary_segments,
+                );
             }
             Hint::External(hint) => {
                 return self.execute_external_hint(vm, hint);
@@ -1649,10 +1656,11 @@ pub fn execute_core_hint_base(
     vm: &mut VirtualMachine,
     exec_scopes: &mut ExecutionScopes,
     core_hint_base: &cairo_lang_casm::hints::CoreHintBase,
+    no_temporary_segments: bool,
 ) -> Result<(), HintError> {
     match core_hint_base {
         cairo_lang_casm::hints::CoreHintBase::Core(core_hint) => {
-            execute_core_hint(vm, exec_scopes, core_hint)
+            execute_core_hint(vm, exec_scopes, core_hint, no_temporary_segments)
         }
         cairo_lang_casm::hints::CoreHintBase::Deprecated(deprecated_hint) => {
             execute_deprecated_hint(vm, exec_scopes, deprecated_hint)
@@ -1723,6 +1731,7 @@ pub fn execute_core_hint(
     vm: &mut VirtualMachine,
     exec_scopes: &mut ExecutionScopes,
     core_hint: &CoreHint,
+    no_temporary_segments: bool,
 ) -> Result<(), HintError> {
     match core_hint {
         CoreHint::AllocSegment { dst } => {
@@ -1921,7 +1930,8 @@ pub fn execute_core_hint(
                     exec_scopes.get_mut_ref::<DictManagerExecScope>("dict_manager_exec_scope")?
                 }
             };
-            let new_dict_segment = dict_manager_exec_scope.new_default_dict(vm);
+            let new_dict_segment =
+                dict_manager_exec_scope.new_default_dict(vm, no_temporary_segments);
             vm.insert_value((dict_infos_base + 3 * n_dicts)?, new_dict_segment)?;
         }
         CoreHint::Felt252DictEntryInit { dict_ptr, key } => {
diff --git a/crates/cairo-lang-runner/src/casm_run/test.rs b/crates/cairo-lang-runner/src/casm_run/test.rs
index 764fee12aa2..bf32b65e00b 100644
--- a/crates/cairo-lang-runner/src/casm_run/test.rs
+++ b/crates/cairo-lang-runner/src/casm_run/test.rs
@@ -129,6 +129,7 @@ fn test_runner(function: CasmContext, n_returns: usize, expected: &[i128]) {
         starknet_state: StarknetState::default(),
         run_resources: RunResources::default(),
         syscalls_used_resources: Default::default(),
+        no_temporary_segments: true,
     };
 
     let RunFunctionResult { ap, memory, .. } =
@@ -157,6 +158,7 @@ fn test_allocate_segment() {
         starknet_state: StarknetState::default(),
         run_resources: RunResources::default(),
         syscalls_used_resources: Default::default(),
+        no_temporary_segments: true,
     };
 
     let RunFunctionResult { ap, memory, .. } =
diff --git a/crates/cairo-lang-runner/src/lib.rs b/crates/cairo-lang-runner/src/lib.rs
index 867e9800184..42c8da36e93 100644
--- a/crates/cairo-lang-runner/src/lib.rs
+++ b/crates/cairo-lang-runner/src/lib.rs
@@ -205,6 +205,7 @@ impl SierraCasmRunner {
             string_to_hint,
             run_resources: RunResources::default(),
             syscalls_used_resources: Default::default(),
+            no_temporary_segments: true,
         };
         let RunResult { gas_counter, memory, value, used_resources, profiling_info } = self
             .run_function(