Introduce StepHooks trait (#2295)

piotmag769 · web-flow · commit c9c508c1681c · 2026-01-19T14:15:45.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,7 @@
 #### Upcoming Changes
 
 * fix: Make cairo1-run to conditionally relocate memory and trace [#2241](https://github.com/lambdaclass/cairo-vm/pull/2241)
+* feat(breaking): Introduce `StepHooks` trait which allows hooks to be stateful [#2295](https://github.com/lambdaclass/cairo-vm/pull/2295)
 
 #### [3.0.1] - 2025-12-22
 
diff --git a/vm/src/vm/hooks.rs b/vm/src/vm/hooks.rs
@@ -13,12 +13,12 @@ use crate::stdlib::{any::Any, collections::HashMap, prelude::*, sync::Arc};
 
 use crate::Felt252;
 
+use super::{errors::vm_errors::VirtualMachineError, vm_core::VirtualMachine};
+use crate::vm::runners::cairo_runner::CairoRunner;
 use crate::{
     hint_processor::hint_processor_definition::HintProcessor, types::exec_scope::ExecutionScopes,
 };
 
-use super::{errors::vm_errors::VirtualMachineError, vm_core::VirtualMachine};
-
 type BeforeFirstStepHookFunc = Arc<
     dyn Fn(&mut VirtualMachine, &[Box<dyn Any>]) -> Result<(), VirtualMachineError> + Sync + Send,
 >;
@@ -35,9 +35,36 @@ type StepHookFunc = Arc<
         + Send,
 >;
 
-/// The hooks to be executed during the VM run
+/// The hooks to be executed during the VM run.
+pub trait StepHooks {
+    fn before_first_step(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hints_data: &[Box<dyn Any>],
+    ) -> Result<(), VirtualMachineError>;
+
+    fn pre_step_instruction(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hint_processor: &mut dyn HintProcessor,
+        exec_scopes: &mut ExecutionScopes,
+        hints_data: &[Box<dyn Any>],
+        constants: &HashMap<String, Felt252>,
+    ) -> Result<(), VirtualMachineError>;
+
+    fn post_step_instruction(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hint_processor: &mut dyn HintProcessor,
+        exec_scopes: &mut ExecutionScopes,
+        hints_data: &[Box<dyn Any>],
+        constants: &HashMap<String, Felt252>,
+    ) -> Result<(), VirtualMachineError>;
+}
+
+/// The hooks to be executed during the VM run.
 ///
-/// They can be individually ignored by setting them to [None]
+/// They can be individually ignored by setting them to [None].
 #[derive(Clone, Default)]
 pub struct Hooks {
     before_first_step: Option<BeforeFirstStepHookFunc>,
@@ -59,13 +86,58 @@ impl Hooks {
     }
 }
 
+impl StepHooks for Hooks {
+    fn before_first_step(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hint_data: &[Box<dyn Any>],
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(before_first_step) = &self.before_first_step {
+            return before_first_step(vm, hint_data);
+        }
+        Ok(())
+    }
+
+    fn pre_step_instruction(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hint_executor: &mut dyn HintProcessor,
+        exec_scope: &mut ExecutionScopes,
+        hints_data: &[Box<dyn Any>],
+        constants: &HashMap<String, Felt252>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(pre_step_instruction) = &self.pre_step_instruction {
+            return pre_step_instruction(vm, hint_executor, exec_scope, hints_data, constants);
+        }
+
+        Ok(())
+    }
+
+    fn post_step_instruction(
+        &mut self,
+        vm: &mut VirtualMachine,
+        hint_executor: &mut dyn HintProcessor,
+        exec_scope: &mut ExecutionScopes,
+        hints_data: &[Box<dyn Any>],
+        constants: &HashMap<String, Felt252>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(post_step_instruction) = &self.post_step_instruction {
+            return post_step_instruction(vm, hint_executor, exec_scope, hints_data, constants);
+        }
+
+        Ok(())
+    }
+}
+
 impl VirtualMachine {
     pub fn execute_before_first_step(
         &mut self,
         hint_data: &[Box<dyn Any>],
     ) -> Result<(), VirtualMachineError> {
-        if let Some(hook_func) = self.hooks.clone().before_first_step {
-            (hook_func)(self, hint_data)?;
+        if let Some(mut hooks) = self.hooks.take() {
+            let result = hooks.before_first_step(self, hint_data);
+            self.hooks = Some(hooks);
+            return result;
         }
 
         Ok(())
@@ -78,8 +150,11 @@ impl VirtualMachine {
         hint_data: &[Box<dyn Any>],
         constants: &HashMap<String, Felt252>,
     ) -> Result<(), VirtualMachineError> {
-        if let Some(hook_func) = self.hooks.clone().pre_step_instruction {
-            (hook_func)(self, hint_executor, exec_scope, hint_data, constants)?;
+        if let Some(mut hooks) = self.hooks.take() {
+            let result =
+                hooks.pre_step_instruction(self, hint_executor, exec_scope, hint_data, constants);
+            self.hooks = Some(hooks);
+            return result;
         }
 
         Ok(())
@@ -92,14 +167,23 @@ impl VirtualMachine {
         hint_data: &[Box<dyn Any>],
         constants: &HashMap<String, Felt252>,
     ) -> Result<(), VirtualMachineError> {
-        if let Some(hook_func) = self.hooks.clone().post_step_instruction {
-            (hook_func)(self, hint_executor, exec_scope, hint_data, constants)?;
+        if let Some(mut hooks) = self.hooks.take() {
+            let result =
+                hooks.post_step_instruction(self, hint_executor, exec_scope, hint_data, constants);
+            self.hooks = Some(hooks);
+            return result;
         }
 
         Ok(())
     }
 }
 
+impl CairoRunner {
+    pub fn set_vm_hooks(&mut self, hooks: Box<dyn StepHooks>) {
+        self.vm.hooks = Some(hooks);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -117,7 +201,7 @@ mod tests {
 
         let mut hint_processor = BuiltinHintProcessor::new_empty();
         let mut cairo_runner = cairo_runner!(program);
-        cairo_runner.vm.hooks = Hooks::new(None, None, None);
+        cairo_runner.vm.hooks = Some(Box::new(Hooks::new(None, None, None)));
 
         let end = cairo_runner.initialize(false).unwrap();
         assert!(cairo_runner.run_until_pc(end, &mut hint_processor).is_ok());
@@ -161,23 +245,35 @@ mod tests {
         // Before first fail
         let mut hint_processor = BuiltinHintProcessor::new_empty();
         let mut cairo_runner = cairo_runner!(program);
-        cairo_runner.vm.hooks = Hooks::new(Some(Arc::new(before_first_step_hook)), None, None);
+        cairo_runner.vm.hooks = Some(Box::new(Hooks::new(
+            Some(Arc::new(before_first_step_hook)),
+            None,
+            None,
+        )));
 
         let end = cairo_runner.initialize(false).unwrap();
         assert!(cairo_runner.run_until_pc(end, &mut hint_processor).is_err());
 
         // Pre step fail
         let mut hint_processor = BuiltinHintProcessor::new_empty();
         let mut cairo_runner = cairo_runner!(program);
-        cairo_runner.vm.hooks = Hooks::new(None, Some(Arc::new(pre_step_hook)), None);
+        cairo_runner.vm.hooks = Some(Box::new(Hooks::new(
+            None,
+            Some(Arc::new(pre_step_hook)),
+            None,
+        )));
 
         let end = cairo_runner.initialize(false).unwrap();
         assert!(cairo_runner.run_until_pc(end, &mut hint_processor).is_err());
 
         // Post step fail
         let mut hint_processor = BuiltinHintProcessor::new_empty();
         let mut cairo_runner = cairo_runner!(program);
-        cairo_runner.vm.hooks = Hooks::new(None, None, Some(Arc::new(post_step_hook)));
+        cairo_runner.vm.hooks = Some(Box::new(Hooks::new(
+            None,
+            None,
+            Some(Arc::new(post_step_hook)),
+        )));
 
         let end = cairo_runner.initialize(false).unwrap();
         assert!(cairo_runner.run_until_pc(end, &mut hint_processor).is_err());
@@ -220,11 +316,11 @@ mod tests {
 
         let mut hint_processor = BuiltinHintProcessor::new_empty();
         let mut cairo_runner = cairo_runner!(program);
-        cairo_runner.vm.hooks = Hooks::new(
+        cairo_runner.vm.hooks = Some(Box::new(Hooks::new(
             Some(Arc::new(before_first_step_hook)),
             Some(Arc::new(pre_step_hook)),
             Some(Arc::new(post_step_hook)),
-        );
+        )));
 
         let end = cairo_runner.initialize(false).unwrap();
         assert!(cairo_runner.run_until_pc(end, &mut hint_processor).is_ok());
diff --git a/vm/src/vm/vm_core.rs b/vm/src/vm/vm_core.rs
@@ -110,7 +110,7 @@ pub struct VirtualMachine {
     pub(crate) disable_trace_padding: bool,
     instruction_cache: Vec<Option<Instruction>>,
     #[cfg(feature = "test_utils")]
-    pub(crate) hooks: crate::vm::hooks::Hooks,
+    pub(crate) hooks: Option<Box<dyn crate::vm::hooks::StepHooks>>,
     pub(crate) relocation_table: Option<Vec<usize>>,
 }
 
@@ -141,7 +141,7 @@ impl VirtualMachine {
             disable_trace_padding,
             instruction_cache: Vec::new(),
             #[cfg(feature = "test_utils")]
-            hooks: Default::default(),
+            hooks: None,
             relocation_table: None,
         }
     }
@@ -1351,7 +1351,7 @@ pub struct VirtualMachineBuilder {
     skip_instruction_execution: bool,
     run_finished: bool,
     #[cfg(feature = "test_utils")]
-    pub(crate) hooks: crate::vm::hooks::Hooks,
+    pub(crate) hooks: Option<Box<dyn crate::vm::hooks::StepHooks>>,
 }
 
 impl Default for VirtualMachineBuilder {
@@ -1371,7 +1371,7 @@ impl Default for VirtualMachineBuilder {
             segments: MemorySegmentManager::new(),
             run_finished: false,
             #[cfg(feature = "test_utils")]
-            hooks: Default::default(),
+            hooks: None,
         }
     }
 }
@@ -1416,8 +1416,8 @@ impl VirtualMachineBuilder {
     }
 
     #[cfg(feature = "test_utils")]
-    pub fn hooks(mut self, hooks: crate::vm::hooks::Hooks) -> VirtualMachineBuilder {
-        self.hooks = hooks;
+    pub fn hooks(mut self, hooks: Box<dyn crate::vm::hooks::StepHooks>) -> VirtualMachineBuilder {
+        self.hooks = Some(hooks);
         self
     }
 
@@ -5345,11 +5345,12 @@ mod tests {
             Err(VirtualMachineError::Unexpected)
         }
         #[cfg(feature = "test_utils")]
-        let virtual_machine_builder = virtual_machine_builder.hooks(crate::vm::hooks::Hooks::new(
-            Some(std::sync::Arc::new(before_first_step_hook)),
-            None,
-            None,
-        ));
+        let virtual_machine_builder =
+            virtual_machine_builder.hooks(Box::new(crate::vm::hooks::Hooks::new(
+                Some(std::sync::Arc::new(before_first_step_hook)),
+                None,
+                None,
+            )));
 
         #[allow(unused_mut)]
         let mut virtual_machine_from_builder = virtual_machine_builder.build();
