feat: add VM hooks as rust conditional feature

Cargo.toml

@@ -13,10 +13,11 @@ description = "Blazing fast Cairo interpreter"
 [features]
 default = ["with_mimalloc"]
 with_mimalloc = ["mimalloc"]
-skip_next_instruction_hint = []
 # This feature will reference every test-oriented feature.
 # Note that these features are not retro-compatible with the cairo Python VM.
-test_utils = ["skip_next_instruction_hint"]
+test_utils = ["skip_next_instruction_hint", "hooks"]
+skip_next_instruction_hint = []
+hooks = []
 
 [dependencies]
 mimalloc = { version = "0.1.29", default-features = false, optional = true }
@@ -34,7 +35,7 @@ sha3 = "0.10.1"
 rand_core = "0.6.4"
 lazy_static = "1.4.0"
 nom = "7"
-sha2 = {version="0.10.2", features=["compress"]}
+sha2 = { version = "0.10.2", features = ["compress"] }
 thiserror = "1.0.32"
 generic-array = "0.14.6"
 keccak = "0.1.2"

src/vm/hooks.rs

@@ -0,0 +1,80 @@
+use std::{collections::HashMap, sync::Arc};
+
+use felt::Felt;
+
+use crate::{
+    hint_processor::hint_processor_definition::HintReference,
+    serde::deserialize_program::ApTracking, types::exec_scope::ExecutionScopes,
+};
+
+use super::{errors::vm_errors::VirtualMachineError, vm_core::VirtualMachine};
+
+type HookFunc = Arc<
+    dyn Fn(
+            &mut CairoRunner,
+            &mut VirtualMachine,
+            &HashMap<usize, Vec<Box<dyn Any>>>,
+        ) -> Result<(), VirtualMachineError>
+        + Sync
+        + Send,
+>;
+
+#[derive(Clone)]
+pub struct Hooks {
+    before_first_step: Option<HookFunc>,
+    pre_step_instruction: Option<HookFunc>,
+    post_step_instruction: Option<HookFunc>,
+}
+
+impl Hooks {
+    pub fn new(
+        before_first_step: Option<HookFunc>,
+        pre_step_instruction: Option<HookFunc>,
+        post_step_instruction: Option<HookFunc>,
+    ) -> Self {
+        Hooks {
+            before_first_step,
+            pre_step_instruction,
+            post_step_instruction,
+        }
+    }
+
+    pub fn execute_pre_step_instruction(
+        &self,
+        runner: &mut CairoRunner,
+        vm: &mut VirtualMachine,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.pre_step_instruction {
+            (hook_func)(runner, vm, hint_data_dictionary)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn execute_post_step_instruction(
+        &self,
+        runner: &mut CairoRunner,
+        vm: &mut VirtualMachine,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.post_step_instruction {
+            (hook_func)(runner, vm, hint_data_dictionary)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn execute_before_first_step(
+        &self,
+        runner: &mut CairoRunner,
+        vm: &mut VirtualMachine,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.before_first_step {
+            (hook_func)(runner, vm, hint_data_dictionary)?;
+        }
+
+        Ok(())
+    }
+}

src/vm/mod.rs

@@ -6,3 +6,6 @@ pub mod security;
 pub mod trace;
 pub mod vm_core;
 pub mod vm_memory;
+
+#[cfg(hooks)]
+pub mod hooks;

src/vm/runners/cairo_runner.rs

@@ -520,15 +520,20 @@ impl CairoRunner {
         address: Relocatable,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
         let references = self.get_reference_list();
         let hint_data_dictionary = self.get_hint_data_dictionary(&references, hint_processor)?;
+        #[cfg(hooks)]
+        hooks.execute_before_first_step(&mut self, vm, &hint_data_dictionary);
         while vm.run_context.pc != address {
             vm.step(
                 hint_processor,
                 &mut self.exec_scopes,
                 &hint_data_dictionary,
                 &self.program.constants,
+                #[cfg(hooks)]
+                hooks,
             )?;
         }
         Ok(())
@@ -540,6 +545,7 @@ impl CairoRunner {
         steps: usize,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
         let references = self.get_reference_list();
         let hint_data_dictionary = self.get_hint_data_dictionary(&references, hint_processor)?;
@@ -554,6 +560,8 @@ impl CairoRunner {
                 &mut self.exec_scopes,
                 &hint_data_dictionary,
                 &self.program.constants,
+                #[cfg(hooks)]
+                hooks,
             )?;
         }
 
@@ -566,17 +574,31 @@ impl CairoRunner {
         steps: usize,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
-        self.run_for_steps(steps.saturating_sub(vm.current_step), vm, hint_processor)
+        self.run_for_steps(
+            steps.saturating_sub(vm.current_step),
+            vm,
+            hint_processor,
+            #[cfg(hooks)]
+            hooks,
+        )
     }
 
     /// Execute steps until the step counter reaches a power of two.
     pub fn run_until_next_power_of_2(
         &mut self,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
-        self.run_until_steps(vm.current_step.next_power_of_two(), vm, hint_processor)
+        self.run_until_steps(
+            vm.current_step.next_power_of_two(),
+            vm,
+            hint_processor,
+            #[cfg(hooks)]
+            hooks,
+        )
     }
 
     pub fn get_perm_range_check_limits(
@@ -700,6 +722,7 @@ impl CairoRunner {
         disable_finalize_all: bool,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
         if self.run_ended {
             return Err(RunnerError::RunAlreadyFinished.into());
@@ -714,7 +737,12 @@ impl CairoRunner {
 
         vm.segments.compute_effective_sizes(&vm.memory);
         if self.proof_mode && !disable_trace_padding {
-            self.run_until_next_power_of_2(vm, hint_processor)?;
+            self.run_until_next_power_of_2(
+                vm,
+                hint_processor,
+                #[cfg(hooks)]
+                hooks,
+            )?;
             loop {
                 match self.check_used_cells(vm) {
                     Ok(_) => break,
@@ -726,8 +754,19 @@ impl CairoRunner {
                     },
                 }
 
-                self.run_for_steps(1, vm, hint_processor)?;
-                self.run_until_next_power_of_2(vm, hint_processor)?;
+                self.run_for_steps(
+                    1,
+                    vm,
+                    hint_processor,
+                    #[cfg(hooks)]
+                    hooks,
+                )?;
+                self.run_until_next_power_of_2(
+                    vm,
+                    hint_processor,
+                    #[cfg(hooks)]
+                    hooks,
+                )?;
             }
         }
 
@@ -959,6 +998,7 @@ impl CairoRunner {
         verify_secure: bool,
         vm: &mut VirtualMachine,
         hint_processor: &mut dyn HintProcessor,
+        #[cfg(hooks)] hooks: crate::vm::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
         let stack = args
             .iter()
@@ -969,8 +1009,21 @@ impl CairoRunner {
 
         self.initialize_vm(vm)?;
 
-        self.run_until_pc(end, vm, hint_processor)?;
-        self.end_run(true, false, vm, hint_processor)?;
+        self.run_until_pc(
+            end,
+            vm,
+            hint_processor,
+            #[cfg(hooks)]
+            hooks,
+        )?;
+        self.end_run(
+            true,
+            false,
+            vm,
+            hint_processor,
+            #[cfg(hooks)]
+            hooks,
+        )?;
 
         if verify_secure {
             verify_secure_runner(self, false, vm)?;

src/vm/vm_core.rs

@@ -522,9 +522,27 @@ impl VirtualMachine {
         exec_scopes: &mut ExecutionScopes,
         hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
         constants: &HashMap<String, Felt>,
+        #[cfg(hooks)] hooks: super::hooks::Hooks,
     ) -> Result<(), VirtualMachineError> {
         self.step_hint(hint_executor, exec_scopes, hint_data_dictionary, constants)?;
-        self.step_instruction()
+
+        #[cfg(hooks)]
+        hooks.execute_pre_step_instruction(
+            &mut self,
+            exec_scopes,
+            hint_data_dictionary,
+            constants,
+        )?;
+        self.step_instruction()?;
+        #[cfg(hooks)]
+        hooks.execute_post_step_instruction(
+            &mut self,
+            exec_scopes,
+            hint_data_dictionary,
+            constants,
+        )?;
+
+        Ok(())
     }
 
     fn compute_op0_deductions(