Merged hint temporary segments by cheatcode hint.

cairo1-run/Cargo.toml

@@ -29,6 +29,7 @@ assert_matches = "1.5.0"
 rstest = "0.17.0"
 mimalloc = { version = "0.1.37", default-features = false, optional = true }
 num-traits = { version = "0.2", default-features = false }
+num-bigint.workspace = true
 
 [features]
 default = ["with_mimalloc"]

cairo1-run/src/cairo_run.rs

@@ -4,7 +4,7 @@ use cairo_lang_casm::{
     casm, casm_build_extend,
     cell_expression::CellExpression,
     deref, deref_or_immediate,
-    hints::Hint,
+    hints::{Hint, StarknetHint},
     inline::CasmContext,
     instructions::{Instruction, InstructionBody},
 };
@@ -31,7 +31,9 @@ use cairo_lang_sierra_to_casm::{
     metadata::calc_metadata_ap_change_only,
 };
 use cairo_lang_sierra_type_size::get_type_size_map;
-use cairo_lang_utils::{casts::IntoOrPanic, unordered_hash_map::UnorderedHashMap};
+use cairo_lang_utils::{
+    bigint::BigIntAsHex, casts::IntoOrPanic, unordered_hash_map::UnorderedHashMap,
+};
 use cairo_vm::{
     hint_processor::cairo_1_hint_processor::hint_processor::Cairo1HintProcessor,
     math_utils::signed_felt,
@@ -48,6 +50,7 @@ use cairo_vm::{
     Felt252,
 };
 use itertools::{chain, Itertools};
+use num_bigint::{BigInt, Sign};
 use num_traits::{cast::ToPrimitive, Zero};
 use std::{collections::HashMap, iter::Peekable};
 
@@ -214,7 +217,6 @@ pub fn cairo_run_program(
         cairo_run_config.trace_enabled,
     )?;
     let end = runner.initialize(cairo_run_config.proof_mode)?;
-    load_arguments(&mut runner, &cairo_run_config, main_func)?;
 
     // Run it until the end / infinite loop in proof_mode
     runner.run_until_pc(end, &mut hint_processor)?;
@@ -608,6 +610,25 @@ fn create_entry_code(
         | BuiltinName::add_mod
         | BuiltinName::mul_mod => unreachable!(),
     };
+    if got_segment_arena {
+        let ptr = get_var(&BuiltinName::segment_arena);
+        ctx.add_hint(
+            |[ignored_in], [ignored_out]| StarknetHint::Cheatcode {
+                selector: BigIntAsHex {
+                    value: BigInt::from_bytes_be(
+                        Sign::Plus,
+                        "FinalizeDictionarySegments".as_bytes(),
+                    ),
+                },
+                input_start: ignored_in.clone(),
+                input_end: ignored_in,
+                output_start: ignored_out,
+                output_end: ignored_out,
+            },
+            [ptr],
+            [ptr],
+        );
+    }
     if copy_to_output_builtin && got_segment_arena {
         // Copying the final builtins into a local variables.
         for (i, builtin) in builtins.iter().enumerate() {

vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs

@@ -14,8 +14,8 @@ pub struct DictTrackerExecScope {
     data: HashMap<Felt252, MaybeRelocatable>,
     /// The start of the segment of the dictionary.
     start: Relocatable,
-    /// The start of the next segment in the segment arena, if finalized.
-    next_start: Option<Relocatable>,
+    /// The end of this segment, if finalized.
+    end: Option<Relocatable>,
 }
 
 /// Helper object to allocate, track and destruct all dictionaries in the run.
@@ -33,7 +33,7 @@ impl DictTrackerExecScope {
         Self {
             data: HashMap::default(),
             start,
-            next_start: None,
+            end: None,
         }
     }
 }
@@ -43,19 +43,12 @@ impl DictManagerExecScope {
 
     /// 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) -> Result<Relocatable, HintError> {
-        let dict_segment = match self.trackers.last() {
-            // This is the first dict - a totally new segment is required.
-            None => vm.add_memory_segment(),
-            // New dict segment should be appended to the last segment.
-            // Appending by a temporary segment, if the last segment is not finalized.
-            Some(last) => last
-                .next_start
-                .unwrap_or_else(|| vm.add_temporary_segment()),
-        };
+        let dict_segment = vm.add_temporary_segment();
         let tracker = DictTrackerExecScope::new(dict_segment);
-        // Not checking if overriding - since overriding is allowed.
-        self.segment_to_tracker
-            .insert(dict_segment.segment_index, self.trackers.len());
+        assert!(self
+            .segment_to_tracker
+            .insert(dict_segment.segment_index, self.trackers.len())
+            .is_none());
 
         self.trackers.push(tracker);
         Ok(dict_segment)
@@ -85,31 +78,29 @@ impl DictManagerExecScope {
     }
 
     /// Finalizes a segment of a dictionary.
-    pub fn finalize_segment(
-        &mut self,
-        vm: &mut VirtualMachine,
-        dict_end: Relocatable,
-    ) -> Result<(), HintError> {
+    pub fn finalize_segment(&mut self, dict_end: Relocatable) -> Result<(), HintError> {
         let tracker_idx = self.get_dict_infos_index(dict_end).unwrap();
         let tracker = &mut self.trackers[tracker_idx];
-        let next_start = (dict_end + 1u32).unwrap();
-        if let Some(prev) = tracker.next_start {
+        if let Some(prev) = tracker.end {
             return Err(HintError::CustomHint(
                 format!(
                     "The segment is already finalized. \
-                    Attempting to override next start {prev}, with: {next_start}.",
+                    Attempting to override next start {prev}, with: {dict_end}.",
                 )
                 .into_boxed_str(),
             ));
         }
-        tracker.next_start = Some(next_start);
-        if let Some(next) = self.trackers.get(tracker_idx + 1) {
-            // Merging the next temporary segment with the closed segment.
-            vm.add_relocation_rule(next.start, next_start).unwrap();
-            // Updating the segment to point to tracker the next segment points to.
-            let next_tracker_idx = self.segment_to_tracker[&next.start.segment_index];
-            self.segment_to_tracker
-                .insert(dict_end.segment_index, next_tracker_idx);
+        tracker.end = Some(dict_end);
+        Ok(())
+    }
+
+    /// Finalizes all segments of dictionaries by adding relocation rules to merge them.
+    pub fn finalize_all_segments(&mut self, vm: &mut VirtualMachine) -> Result<(), HintError> {
+        let mut prev_end = vm.add_memory_segment();
+        for tracker in &self.trackers {
+            vm.add_relocation_rule(tracker.start, prev_end)?;
+            prev_end += (tracker.end.unwrap() - tracker.start)?;
+            prev_end += 1;
         }
         Ok(())
     }

vm/src/hint_processor/cairo_1_hint_processor/hint_processor.rs

@@ -17,7 +17,7 @@ use crate::{
 use ark_ff::fields::{Fp256, MontBackend, MontConfig};
 use ark_ff::{Field, PrimeField};
 use ark_std::UniformRand;
-use cairo_lang_casm::hints::{CoreHintBase, DeprecatedHint};
+use cairo_lang_casm::hints::{CoreHintBase, DeprecatedHint, StarknetHint};
 use cairo_lang_casm::{
     hints::{CoreHint, Hint},
     operand::{CellRef, ResOperand},
@@ -266,6 +266,20 @@ impl Cairo1HintProcessor {
                 t_or_k0,
                 t_or_k1,
             ),
+            Hint::Starknet(StarknetHint::Cheatcode { selector, .. }) => {
+                let selector = &selector.value.to_bytes_be().1;
+                let selector = std::str::from_utf8(selector).map_err(|_| {
+                    HintError::CustomHint(Box::from("failed to parse selector".to_string()))
+                })?;
+                match selector {
+                    "FinalizeDictionarySegments" => {
+                        let dict_manager_exec_scope = exec_scopes
+                            .get_mut_ref::<DictManagerExecScope>("dict_manager_exec_scope")?;
+                        dict_manager_exec_scope.finalize_all_segments(vm)
+                    }
+                    _ => Err(HintError::UnknownHint(selector.into())),
+                }
+            }
 
             hint => Err(HintError::UnknownHint(
                 format!("{:?}", hint).into_boxed_str(),
@@ -418,7 +432,7 @@ impl Cairo1HintProcessor {
         vm.insert_value(cell_ref_to_relocatable(dict_index, vm)?, dict_infos_index)
             .map_err(HintError::from)?;
         // The hint is only for dictionary finalization, so can be called.
-        dict_manager_exec_scope.finalize_segment(vm, dict_address)
+        dict_manager_exec_scope.finalize_segment(dict_address)
     }
 
     #[allow(clippy::too_many_arguments)]