Extracted entry-code creation into a basic configuration. (#6640)

crates/cairo-lang-runnable-utils/src/builder.rs

@@ -159,17 +159,22 @@ impl RunnableBuilder {
     pub fn assemble_function_program(
         &self,
         func: &Function,
+        config: EntryCodeConfig,
     ) -> Result<(AssembledCairoProgram, Vec<BuiltinName>), BuildError> {
-        let (header, builtins) = self.create_entry_code(func)?;
+        let (header, builtins) = self.create_entry_code(func, config)?;
         let footer = create_code_footer();
 
         let assembled_cairo_program = self.casm_program.assemble_ex(&header, &footer);
         Ok((assembled_cairo_program, builtins))
     }
 
     /// CASM style string representation of the program.
-    pub fn casm_function_program(&self, func: &Function) -> Result<String, BuildError> {
-        let (header, builtins) = self.create_entry_code(func)?;
+    pub fn casm_function_program(
+        &self,
+        func: &Function,
+        config: EntryCodeConfig,
+    ) -> Result<String, BuildError> {
+        let (header, builtins) = self.create_entry_code(func, config)?;
         let footer = create_code_footer();
 
         Ok(chain!(
@@ -193,6 +198,7 @@ impl RunnableBuilder {
     fn create_entry_code(
         &self,
         func: &Function,
+        config: EntryCodeConfig,
     ) -> Result<(Vec<Instruction>, Vec<BuiltinName>), BuildError> {
         let param_types = self.generic_id_and_size_from_concrete(&func.signature.param_types);
         let return_types = self.generic_id_and_size_from_concrete(&func.signature.ret_types);
@@ -201,14 +207,18 @@ impl RunnableBuilder {
         let code_offset =
             self.casm_program.debug_info.sierra_statement_info[entry_point].start_offset;
         // Finalizing for proof only if all returned values are builtins or droppable.
-        // (To handle cases such as a function returning a non-squashed dict, which can't be
-        // provable by itself)
-        let finalize_for_proof = func.signature.ret_types.iter().all(|ty| {
+        let droppable_return_value = func.signature.ret_types.iter().all(|ty| {
             let info = self.type_info(ty);
             info.droppable || !self.is_user_arg_type(&info.long_id.generic_id)
         });
+        if !droppable_return_value {
+            assert!(
+                !config.finalize_segment_arena,
+                "Cannot finalize the segment arena when returning non-droppable values."
+            );
+        }
 
-        create_entry_code_from_params(&param_types, &return_types, code_offset, finalize_for_proof)
+        create_entry_code_from_params(&param_types, &return_types, code_offset, config)
     }
 
     /// Converts array of `ConcreteTypeId`s into corresponding `GenericTypeId`s and their sizes
@@ -228,6 +238,27 @@ impl RunnableBuilder {
     }
 }
 
+/// Configuration for the entry code creation.
+#[derive(Clone, Debug)]
+pub struct EntryCodeConfig {
+    /// Whether to finalize the segment arena after calling the function.
+    pub finalize_segment_arena: bool,
+}
+impl EntryCodeConfig {
+    /// Returns a configuration for testing purposes.
+    ///
+    /// This configuration will not finalize the segment arena after calling the function, to
+    /// prevent failure in case of functions returning values.
+    pub fn testing() -> Self {
+        Self { finalize_segment_arena: false }
+    }
+
+    /// Returns a configuration for proving purposes.
+    pub fn provable() -> Self {
+        Self { finalize_segment_arena: true }
+    }
+}
+
 /// Returns the entry code to call the function with `param_types` as its inputs and
 /// `return_types` as outputs, located at `code_offset`. If `finalize_for_proof` is true,
 /// will make sure to remove the segment arena after calling the function. For testing purposes,
@@ -236,7 +267,7 @@ pub fn create_entry_code_from_params(
     param_types: &[(GenericTypeId, i16)],
     return_types: &[(GenericTypeId, i16)],
     code_offset: usize,
-    finalize_for_proof: bool,
+    config: EntryCodeConfig,
 ) -> Result<(Vec<Instruction>, Vec<BuiltinName>), BuildError> {
     let mut ctx = CasmBuilder::default();
     let mut builtin_offset = 3;
@@ -267,7 +298,7 @@ pub fn create_entry_code_from_params(
     let emulated_builtins = UnorderedHashSet::<_>::from_iter([SystemType::ID]);
 
     let got_segment_arena = param_types.iter().any(|(ty, _)| ty == &SegmentArenaType::ID);
-    let has_post_calculation_loop = got_segment_arena && finalize_for_proof;
+    let has_post_calculation_loop = got_segment_arena && config.finalize_segment_arena;
 
     let mut local_exprs = vec![];
     if has_post_calculation_loop {
@@ -352,7 +383,7 @@ pub fn create_entry_code_from_params(
             casm_build_extend!(ctx, assert local_cell = cell;);
         }
     }
-    if got_segment_arena && finalize_for_proof {
+    if got_segment_arena && config.finalize_segment_arena {
         let segment_arena = builtin_vars[&SegmentArenaType::ID];
         // Validating the segment arena's segments are one after the other.
         casm_build_extend! {ctx,

crates/cairo-lang-runnable/src/compile.rs

@@ -7,7 +7,7 @@ use cairo_lang_compiler::diagnostics::DiagnosticsReporter;
 use cairo_lang_compiler::project::setup_project;
 use cairo_lang_filesystem::ids::CrateId;
 use cairo_lang_lowering::ids::ConcreteFunctionWithBodyId;
-use cairo_lang_runnable_utils::builder::RunnableBuilder;
+use cairo_lang_runnable_utils::builder::{EntryCodeConfig, RunnableBuilder};
 use cairo_lang_semantic::plugin::PluginSuite;
 use cairo_lang_sierra_generator::db::SierraGenGroup;
 use cairo_lang_sierra_generator::executables::find_executable_function_ids;
@@ -99,5 +99,5 @@ pub fn compile_runnable_function_in_prepared_db(
     );
     let runnable_func = sierra_program.funcs[0].clone();
     let builder = RunnableBuilder::new(sierra_program, None)?;
-    Ok(builder.casm_function_program(&runnable_func)?)
+    Ok(builder.casm_function_program(&runnable_func, EntryCodeConfig::provable())?)
 }

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()

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 } => {

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, .. } =

crates/cairo-lang-runner/src/lib.rs

@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use std::ops::{Add, Sub};
 
 use cairo_lang_casm::hints::Hint;
-use cairo_lang_runnable_utils::builder::{BuildError, RunnableBuilder};
+use cairo_lang_runnable_utils::builder::{BuildError, EntryCodeConfig, RunnableBuilder};
 use cairo_lang_sierra::extensions::NamedType;
 use cairo_lang_sierra::extensions::core::CoreConcreteLibfunc;
 use cairo_lang_sierra::extensions::enm::EnumType;
@@ -195,7 +195,8 @@ impl SierraCasmRunner {
         available_gas: Option<usize>,
         starknet_state: StarknetState,
     ) -> Result<RunResultStarknet, RunnerError> {
-        let (assembled_program, builtins) = self.builder.assemble_function_program(func)?;
+        let (assembled_program, builtins) =
+            self.builder.assemble_function_program(func, EntryCodeConfig::testing())?;
         let (hints_dict, string_to_hint) = build_hints_dict(&assembled_program.hints);
         let user_args = self.prepare_args(func, available_gas, args)?;
         let mut hint_processor = CairoHintProcessor {
@@ -205,6 +206,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(