[Cairo1] Use a cheatcode to relocate all dicts + Make temporary segment usage configurable (#1776)

fmoletta · orizi · web-flow · commit c745eceb9bd1 · 2024-05-30T19:09:42.000Z
* Merged hint temporary segments by cheatcode hint.

* Added line removed by mistake.

* Remove std import

* Rename Cheatcode for clarity + remove unwrap

* use real or temporary based on flag

* Remove assert

* Remove unwrap

* Add flag to HintProcessor

* Update test util

* Add Changelog Entry

* Fmt + clarity

---------

Co-authored-by: Ori Ziv &lt;oriz@starkware.co&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,13 @@
 
 #### Upcoming Changes
 
+* feat(BREAKING): Use a cheatcode to relocate all dicts + Make temporary segment usage configurable [#1776](https://github.com/lambdaclass/cairo-vm/pull/1776)
+  * Add the flags `segment_arena_validation` & `use_temporary_segments` to the `Cairo1HintProcessor` & `DictManagerExecScope` respectively. These flags will determine if real segments or temporary segments will be used when creating dictionaries.
+  * `DictManagerExecScope::finalize_segment` no longer performs relocation and is ignored if `use_temporary_segments` is set to false.
+  * Add method `DictManagerExecScope::relocate_all_dictionaries` that adds relocation rules for all tracked dictionaries, relocating them one next to the other in a new segment.
+  * Add cheatcode `RelocateAllDictionaries` to the `Cairo1HintProcessor`, which calls the aforementioned method.
+  * Add casm instruction to call the aforementioned cheatcode in `create_entry_code` if either `proof_mode` or `append_return_values` are set to true, and segment arena is present.
+
 * Bump `starknet-types-core` version + Use the lib's pedersen hash [#1734](https://github.com/lambdaclass/cairo-vm/pull/1734)
 
 * refactor: Add boolean method Cairo1RunConfig::copy_to_output + Update Doc [#1778](https://github.com/lambdaclass/cairo-vm/pull/1778)
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cairo1-run/Cargo.toml b/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"]
diff --git a/cairo1-run/src/cairo_run.rs b/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};
 
@@ -195,7 +198,11 @@ pub fn cairo_run_program(
 
     let (processor_hints, program_hints) = build_hints_vec(instructions.clone());
 
-    let mut hint_processor = Cairo1HintProcessor::new(&processor_hints, RunResources::default());
+    let mut hint_processor = Cairo1HintProcessor::new(
+        &processor_hints,
+        RunResources::default(),
+        cairo_run_config.copy_to_output(),
+    );
 
     let data: Vec<MaybeRelocatable> = instructions
         .flat_map(|inst| inst.assemble().encode())
@@ -742,6 +749,23 @@ fn create_entry_code(
             // len(builtins) + len(builtins - output) + segment_arena_ptr + info_segment + 0
             let off = 2 * builtins.len() + 2;
             let segment_arena_ptr = ctx.add_var(CellExpression::Deref(deref!([fp + off as i16])));
+            // Call the hint that will relocate all dictionaries
+            ctx.add_hint(
+                |[ignored_in], [ignored_out]| StarknetHint::Cheatcode {
+                    selector: BigIntAsHex {
+                        value: BigInt::from_bytes_be(
+                            Sign::Plus,
+                            "RelocateAllDictionaries".as_bytes(),
+                        ),
+                    },
+                    input_start: ignored_in.clone(),
+                    input_end: ignored_in,
+                    output_start: ignored_out,
+                    output_end: ignored_out,
+                },
+                [segment_arena_ptr],
+                [segment_arena_ptr],
+            );
             // Validating the segment arena's segments are one after the other.
             casm_build_extend! {ctx,
                 tempvar n_segments = segment_arena_ptr[-2];
diff --git a/vm/src/hint_processor/cairo_1_hint_processor/dict_manager.rs b/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.
@@ -25,6 +25,9 @@ pub struct DictManagerExecScope {
     segment_to_tracker: HashMap<isize, usize>,
     /// The actual trackers of the dictionaries, in the order of allocation.
     trackers: Vec<DictTrackerExecScope>,
+    // If set to true, dictionaries will be created on temporary segments which can then be relocated into a single segment by the end of the run
+    // If set to false, each dictionary will use a single real segment
+    use_temporary_segments: bool,
 }
 
 impl DictTrackerExecScope {
@@ -33,29 +36,39 @@ impl DictTrackerExecScope {
         Self {
             data: HashMap::default(),
             start,
-            next_start: None,
+            end: None,
         }
     }
 }
 
 impl DictManagerExecScope {
     pub const DICT_DEFAULT_VALUE: usize = 0;
 
+    // Creates a new DictManagerExecScope
+    pub fn new(use_temporary_segments: bool) -> Self {
+        Self {
+            use_temporary_segments,
+            ..Default::default()
+        }
+    }
+
     /// 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 = if self.use_temporary_segments {
+            vm.add_temporary_segment()
+        } else {
+            vm.add_memory_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());
+        if self
+            .segment_to_tracker
+            .insert(dict_segment.segment_index, self.trackers.len())
+            .is_some()
+        {
+            return Err(HintError::CantCreateDictionaryOnTakenSegment(
+                dict_segment.segment_index,
+            ));
+        }
 
         self.trackers.push(tracker);
         Ok(dict_segment)
@@ -85,31 +98,35 @@ impl DictManagerExecScope {
     }
 
     /// Finalizes a segment of a dictionary.
-    pub fn finalize_segment(
-        &mut self,
-        vm: &mut VirtualMachine,
-        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 {
-            return Err(HintError::CustomHint(
-                format!(
-                    "The segment is already finalized. \
-                    Attempting to override next start {prev}, with: {next_start}.",
-                )
-                .into_boxed_str(),
-            ));
+    /// Does nothing if use_temporary_segments is set to false
+    pub fn finalize_segment(&mut self, dict_end: Relocatable) -> Result<(), HintError> {
+        if self.use_temporary_segments {
+            let tracker_idx = self.get_dict_infos_index(dict_end)?;
+            let tracker = &mut self.trackers[tracker_idx];
+            if let Some(prev) = tracker.end {
+                return Err(HintError::CustomHint(
+                    format!(
+                        "The segment is already finalized. \
+                    Attempting to override next start {prev}, with: {dict_end}.",
+                    )
+                    .into_boxed_str(),
+                ));
+            }
+            tracker.end = Some(dict_end);
         }
-        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);
+        Ok(())
+    }
+
+    /// Relocates all dictionaries into a single segment
+    /// Does nothing if use_temporary_segments is set to false
+    pub fn relocate_all_dictionaries(&mut self, vm: &mut VirtualMachine) -> Result<(), HintError> {
+        if self.use_temporary_segments {
+            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_or_default() - tracker.start)?;
+                prev_end += 1;
+            }
         }
         Ok(())
     }
diff --git a/vm/src/hint_processor/cairo_1_hint_processor/hint_processor.rs b/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},
@@ -54,13 +54,21 @@ fn get_beta() -> Felt252 {
 pub struct Cairo1HintProcessor {
     hints: HashMap<usize, Vec<Hint>>,
     run_resources: RunResources,
+    /// If set to true, uses a single segment for dictionaries to aid in segment arena validations
+    /// WARNING: The program must call the "RelocateAllDictionaries" Cheatcode if the flag is enabled
+    segment_arena_validations: bool,
 }
 
 impl Cairo1HintProcessor {
-    pub fn new(hints: &[(usize, Vec<Hint>)], run_resources: RunResources) -> Self {
+    pub fn new(
+        hints: &[(usize, Vec<Hint>)],
+        run_resources: RunResources,
+        segment_arena_validations: bool,
+    ) -> Self {
         Self {
             hints: hints.iter().cloned().collect(),
             run_resources,
+            segment_arena_validations,
         }
     }
     // Runs a single Hint
@@ -166,7 +174,7 @@ impl Cairo1HintProcessor {
             })) => self.linear_split(vm, value, scalar, max_x, x, y),
 
             Hint::Core(CoreHintBase::Core(CoreHint::AllocFelt252Dict { segment_arena_ptr })) => {
-                self.alloc_felt_256_dict(vm, segment_arena_ptr, exec_scopes)
+                self.alloc_felt_252_dict(vm, segment_arena_ptr, exec_scopes)
             }
 
             Hint::Core(CoreHintBase::Core(CoreHint::AssertLeFindSmallArcs {
@@ -266,6 +274,20 @@ impl Cairo1HintProcessor {
                 t_or_k0,
                 t_or_k1,
             ),
+            Hint::Starknet(StarknetHint::Cheatcode { selector, .. }) => {
+                let selector = &selector.value.to_bytes_be().1;
+                let selector = crate::stdlib::str::from_utf8(selector).map_err(|_| {
+                    HintError::CustomHint(Box::from("failed to parse selector".to_string()))
+                })?;
+                match selector {
+                    "RelocateAllDictionaries" => {
+                        let dict_manager_exec_scope = exec_scopes
+                            .get_mut_ref::<DictManagerExecScope>("dict_manager_exec_scope")?;
+                        dict_manager_exec_scope.relocate_all_dictionaries(vm)
+                    }
+                    _ => Err(HintError::UnknownHint(selector.into())),
+                }
+            }
 
             hint => Err(HintError::UnknownHint(
                 format!("{:?}", hint).into_boxed_str(),
@@ -418,7 +440,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)]
@@ -548,7 +570,7 @@ impl Cairo1HintProcessor {
         Err(HintError::KeyNotFound)
     }
 
-    fn alloc_felt_256_dict(
+    fn alloc_felt_252_dict(
         &self,
         vm: &mut VirtualMachine,
         segment_arena_ptr: &ResOperand,
@@ -577,7 +599,7 @@ impl Cairo1HintProcessor {
                 Err(_) => {
                     exec_scopes.assign_or_update_variable(
                         "dict_manager_exec_scope",
-                        Box::<DictManagerExecScope>::default(),
+                        Box::new(DictManagerExecScope::new(self.segment_arena_validations)),
                     );
                     exec_scopes.get_mut_ref::<DictManagerExecScope>("dict_manager_exec_scope")?
                 }
diff --git a/vm/src/tests/cairo_1_run_from_entrypoint_tests.rs b/vm/src/tests/cairo_1_run_from_entrypoint_tests.rs
@@ -604,7 +604,7 @@ fn fibonacci_with_run_resources_ok() {
     let contract_class: CasmContractClass = serde_json::from_slice(program_data).unwrap();
     // Program takes 621 steps
     let mut hint_processor =
-        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(621));
+        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(621), false);
     assert_matches!(
         run_cairo_1_entrypoint_with_run_resources(
             serde_json::from_slice(program_data.as_slice()).unwrap(),
@@ -625,7 +625,7 @@ fn fibonacci_with_run_resources_2_ok() {
     let contract_class: CasmContractClass = serde_json::from_slice(program_data).unwrap();
     // Program takes 621 steps
     let mut hint_processor =
-        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(1000));
+        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(1000), false);
     assert_matches!(
         run_cairo_1_entrypoint_with_run_resources(
             contract_class,
@@ -648,7 +648,7 @@ fn fibonacci_with_run_resources_error() {
     let contract_class: CasmContractClass = serde_json::from_slice(program_data).unwrap();
     // Program takes 621 steps
     let mut hint_processor =
-        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(100));
+        Cairo1HintProcessor::new(&contract_class.hints, RunResources::new(100), false);
     assert!(run_cairo_1_entrypoint_with_run_resources(
         contract_class,
         0,
diff --git a/vm/src/tests/mod.rs b/vm/src/tests/mod.rs
@@ -108,7 +108,7 @@ fn run_cairo_1_entrypoint(
 ) {
     let contract_class: CasmContractClass = serde_json::from_slice(program_content).unwrap();
     let mut hint_processor =
-        Cairo1HintProcessor::new(&contract_class.hints, RunResources::default());
+        Cairo1HintProcessor::new(&contract_class.hints, RunResources::default(), false);
 
     let mut runner = CairoRunner::new(
         &(contract_class.clone().try_into().unwrap()),
