[CHECK PERF] Instruction cache

src/types/instruction.rs

@@ -10,7 +10,7 @@ pub enum Register {
     FP,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Instruction {
     pub off0: isize,
     pub off1: isize,
@@ -25,46 +25,46 @@ pub struct Instruction {
     pub opcode: Opcode,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Op1Addr {
     Imm,
     AP,
     FP,
     Op0,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Res {
     Op1,
     Add,
     Mul,
     Unconstrained,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum PcUpdate {
     Regular,
     Jump,
     JumpRel,
     Jnz,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum ApUpdate {
     Regular,
     Add,
     Add1,
     Add2,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum FpUpdate {
     Regular,
     APPlus2,
     Dst,
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Opcode {
     NOp,
     AssertEq,

src/vm/runners/cairo_runner.rs

@@ -329,8 +329,7 @@ impl CairoRunner {
                 offset: prog_base.offset + entrypoint,
             };
             self.initial_pc = Some(initial_pc);
-            vm.segments
-                .load_data(prog_base, &self.program.shared_program_data.data)
+            vm.load_data(prog_base, &self.program.shared_program_data.data)
                 .map_err(RunnerError::MemoryInitializationError)?;
 
             // Mark all addresses from the program segment as accessed
@@ -342,8 +341,7 @@ impl CairoRunner {
             }
         }
         if let Some(exec_base) = self.execution_base {
-            vm.segments
-                .load_data(exec_base, &stack)
+            vm.load_data(exec_base, &stack)
                 .map_err(RunnerError::MemoryInitializationError)?;
         } else {
             return Err(RunnerError::NoProgBase);

src/vm/vm_core.rs

@@ -83,6 +83,7 @@ pub struct VirtualMachine {
     run_finished: bool,
     #[cfg(feature = "hooks")]
     pub(crate) hooks: crate::vm::hooks::Hooks,
+    instruction_cache: Vec<Option<Instruction>>,
 }
 
 impl VirtualMachine {
@@ -110,6 +111,7 @@ impl VirtualMachine {
             trace_relocated: false,
             #[cfg(feature = "hooks")]
             hooks: Default::default(),
+            instruction_cache: Vec::new(),
         }
     }
 
@@ -183,12 +185,12 @@ impl VirtualMachine {
 
     fn update_registers(
         &mut self,
-        instruction: Instruction,
+        instruction: &Instruction,
         operands: Operands,
     ) -> Result<(), VirtualMachineError> {
-        self.update_fp(&instruction, &operands)?;
-        self.update_ap(&instruction, &operands)?;
-        self.update_pc(&instruction, &operands)?;
+        self.update_fp(instruction, &operands)?;
+        self.update_ap(instruction, &operands)?;
+        self.update_pc(instruction, &operands)?;
         Ok(())
     }
 
@@ -381,7 +383,7 @@ impl VirtualMachine {
         Ok(())
     }
 
-    fn run_instruction(&mut self, instruction: Instruction) -> Result<(), VirtualMachineError> {
+    fn run_instruction(&mut self, instruction: &Instruction) -> Result<(), VirtualMachineError> {
         let (operands, operands_addresses, deduced_operands) =
             self.compute_operands(&instruction)?;
         self.insert_deduced_operands(deduced_operands, &operands, &operands_addresses)?;
@@ -410,16 +412,6 @@ impl VirtualMachine {
         Ok(())
     }
 
-    fn decode_current_instruction(&self) -> Result<Instruction, VirtualMachineError> {
-        let instruction = self
-            .segments
-            .memory
-            .get_integer(self.run_context.pc)?
-            .to_u64()
-            .ok_or(VirtualMachineError::InvalidInstructionEncoding)?;
-        decode_instruction(instruction)
-    }
-
     pub fn step_hint(
         &mut self,
         hint_executor: &mut dyn HintProcessor,
@@ -438,13 +430,17 @@ impl VirtualMachine {
     }
 
     pub fn step_instruction(&mut self) -> Result<(), VirtualMachineError> {
-        let instruction = self.decode_current_instruction()?;
+        let inst_cache = core::mem::take(&mut self.instruction_cache);
+        let instruction = inst_cache[self.run_context.pc.offset]
+            .as_ref()
+            .ok_or(VirtualMachineError::InvalidInstructionEncoding)?;
         if !self.skip_instruction_execution {
             self.run_instruction(instruction)?;
         } else {
             self.run_context.pc += instruction.size();
             self.skip_instruction_execution = false;
         }
+        self.instruction_cache = inst_cache;
         Ok(())
     }
 
@@ -786,12 +782,27 @@ impl VirtualMachine {
         self.segments.memory.insert_value(key, val)
     }
 
+    fn fill_instruction_cache(&mut self, base: usize, data: &[MaybeRelocatable]) {
+        self.instruction_cache.resize(base + data.len(), None);
+        for (i, v) in data.iter().enumerate() {
+            let MaybeRelocatable::Int(v) = v else {
+                continue;
+            };
+            let instr = v.to_u64().map(|v| decode_instruction(v).ok()).flatten();
+            self.instruction_cache[base + i] = instr;
+        }
+    }
+
     ///Writes data into the memory from address ptr and returns the first address after the data.
     pub fn load_data(
         &mut self,
         ptr: Relocatable,
         data: &Vec<MaybeRelocatable>,
     ) -> Result<Relocatable, MemoryError> {
+        // Hack to avoid decoding instructions at each step
+        if ptr.segment_index == 0 {
+            self.fill_instruction_cache(ptr.offset, data);
+        }
         self.segments.load_data(ptr, data)
     }
 
@@ -1069,6 +1080,7 @@ impl VirtualMachineBuilder {
             trace_relocated: false,
             #[cfg(feature = "hooks")]
             hooks: self.hooks,
+            instruction_cache: Vec::new(),
         }
     }
 }
@@ -1750,7 +1762,7 @@ mod tests {
         vm.run_context.fp = 6;
 
         assert_matches!(
-            vm.update_registers(instruction, operands),
+            vm.update_registers(&instruction, operands),
             Ok::<(), VirtualMachineError>(())
         );
         assert_eq!(vm.run_context.pc, Relocatable::from((0, 5)));
@@ -1786,7 +1798,7 @@ mod tests {
         run_context!(vm, 4, 5, 6);
 
         assert_matches!(
-            vm.update_registers(instruction, operands),
+            vm.update_registers(&instruction, operands),
             Ok::<(), VirtualMachineError>(())
         );
         assert_eq!(vm.run_context.pc, Relocatable::from((0, 12)));
@@ -3911,17 +3923,6 @@ mod tests {
         );
     }
 
-    #[test]
-    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
-    fn decode_current_instruction_invalid_encoding() {
-        let mut vm = vm!();
-        vm.segments = segments![((0, 0), ("112233445566778899", 16))];
-        assert_matches!(
-            vm.decode_current_instruction(),
-            Err(VirtualMachineError::InvalidInstructionEncoding)
-        );
-    }
-
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn add_relocation_rule_test() {