Add support for caching program segment

vm/src/vm/runners/cairo_runner.rs

@@ -13,6 +13,7 @@ use crate::{
     vm::{
         runners::builtin_runner::SegmentArenaBuiltinRunner,
         trace::trace_entry::{relocate_trace_register, RelocatedTraceEntry, TraceEntry},
+        vm_memory::memory::MemoryCell,
     },
     Felt252,
 };
@@ -153,6 +154,7 @@ pub struct CairoRunner {
     initial_fp: Option<Relocatable>,
     initial_pc: Option<Relocatable>,
     run_ended: bool,
+    loaded_program: bool,
     segments_finalized: bool,
     execution_public_memory: Option<Vec<usize>>,
     pub(crate) runner_mode: RunnerMode,
@@ -210,6 +212,7 @@ impl CairoRunner {
             initial_fp: None,
             initial_pc: None,
             run_ended: false,
+            loaded_program: false,
             segments_finalized: false,
             runner_mode: mode.clone(),
             relocated_memory: Vec::new(),
@@ -486,18 +489,47 @@ impl CairoRunner {
         let prog_base = self.program_base.ok_or(RunnerError::NoProgBase)?;
         let exec_base = self.execution_base.ok_or(RunnerError::NoExecBase)?;
         self.initial_pc = Some((prog_base + entrypoint)?);
+        if !self.loaded_program {
+            self.load_program()?;
+        }
         self.vm
-            .load_data(prog_base, &self.program.shared_program_data.data)
+            .segments
+            .load_data(exec_base, &stack)
             .map_err(RunnerError::MemoryInitializationError)?;
+        Ok(())
+    }
 
+    /// Loads the program in the program segment
+    ///
+    /// If this method is not called, the program is loaded automatically before
+    /// execution, in `initialize_state`.
+    pub fn load_program(&mut self) -> Result<(), RunnerError> {
+        let prog_base = self.program_base.ok_or(RunnerError::NoProgBase)?;
+        self.vm
+            .load_data(prog_base, &self.program.shared_program_data.data)
+            .map_err(RunnerError::MemoryInitializationError)?;
         // Mark all addresses from the program segment as accessed
         for i in 0..self.program.shared_program_data.data.len() {
             self.vm.segments.memory.mark_as_accessed((prog_base + i)?);
         }
+        self.loaded_program = true;
+        Ok(())
+    }
+
+    /// Loads the given program in the program segment.
+    ///
+    /// If this method is not called, the program is loaded automatically before
+    /// execution, in `initialize_state`.
+    ///
+    /// # Safety
+    ///
+    /// The given program must be the same as the one defined in the runner.
+    pub fn load_cached_program(&mut self, data: Vec<MemoryCell>) -> Result<(), RunnerError> {
         self.vm
             .segments
-            .load_data(exec_base, &stack)
-            .map_err(RunnerError::MemoryInitializationError)?;
+            .memory
+            .replace_segment(self.program_base.unwrap(), data)?;
+        self.loaded_program = true;
         Ok(())
     }
 

vm/src/vm/vm_memory/memory.rs

@@ -1,4 +1,4 @@
-use crate::stdlib::{borrow::Cow, collections::HashMap, fmt, prelude::*};
+use crate::stdlib::{borrow::Cow, collections::HashMap, fmt, mem::replace, prelude::*};
 
 use crate::types::errors::math_errors::MathError;
 use crate::vm::runners::cairo_pie::CairoPieMemory;
@@ -38,7 +38,7 @@ pub struct ValidationRule(
 ///   and the 4th word storing the offset.
 #[derive(Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Debug)]
 #[repr(align(32))]
-pub(crate) struct MemoryCell([u64; 4]);
+pub struct MemoryCell([u64; 4]);
 
 impl MemoryCell {
     pub const NONE_MASK: u64 = 1 << 63;
@@ -106,6 +106,7 @@ impl From<MemoryCell> for MaybeRelocatable {
     }
 }
 
+#[derive(Clone)]
 pub struct AddressSet(Vec<bv::BitVec>);
 
 impl AddressSet {
@@ -185,7 +186,28 @@ impl Memory {
         }
     }
 
-    fn get_segment(&mut self, key: Relocatable) -> Result<&mut Vec<MemoryCell>, MemoryError> {
+    pub fn replace_segment(
+        &mut self,
+        key: Relocatable,
+        segment: Vec<MemoryCell>,
+    ) -> Result<(), MemoryError> {
+        let (value_index, _) = from_relocatable_to_indexes(key);
+        let data = if key.segment_index.is_negative() {
+            &mut self.temp_data
+        } else {
+            &mut self.data
+        };
+        let data_len = data.len();
+        let _ = replace(
+            data.get_mut(value_index).ok_or_else(|| {
+                MemoryError::UnallocatedSegment(Box::new((value_index, data_len)))
+            })?,
+            segment,
+        );
+        Ok(())
+    }
+
+    pub fn get_segment(&mut self, key: Relocatable) -> Result<&mut Vec<MemoryCell>, MemoryError> {
         let (value_index, _) = from_relocatable_to_indexes(key);
         let data = if key.segment_index.is_negative() {
             &mut self.temp_data