feat: add VM hooks as rust conditional feature (#761)

doc: improve doc support for optional features

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"

Makefile

@@ -165,8 +165,9 @@ compare_trace_proof: $(CAIRO_RS_TRACE_PROOF) $(CAIRO_TRACE_PROOF)
 compare_memory_proof: $(CAIRO_RS_MEM_PROOF) $(CAIRO_MEM_PROOF)
 	cd tests; ./compare_vm_state.sh memory proof_mode
 
+# Run with nightly enable the `doc_cfg` feature wich let us provide clear explaination about which parts of the code are behind a feature flag
 docs:
-	cargo doc --verbose --release --locked --no-deps
+	RUSTDOCFLAGS="--cfg docsrs" cargo +nightly doc --verbose --release --locked --no-deps --all-features --open
 
 clean:
 	rm -f $(TEST_DIR)/*.json

src/hint_processor/builtin_hint_processor/mod.rs

@@ -17,6 +17,8 @@ pub mod segments;
 pub mod set;
 pub mod sha256_utils;
 pub mod signature;
+#[cfg(feature = "skip_next_instruction_hint")]
+#[cfg_attr(docsrs, doc(cfg(feature = "skip_next_instruction_hint")))]
 pub mod skip_next_instruction;
 pub mod squash_dict_utils;
 pub mod uint256_utils;

src/hint_processor/builtin_hint_processor/skip_next_instruction.rs

@@ -1,14 +1,10 @@
-#[cfg(feature = "skip_next_instruction_hint")]
 use crate::vm::errors::hint_errors::HintError;
-#[cfg(feature = "skip_next_instruction_hint")]
 use crate::vm::vm_core::VirtualMachine;
 
-/*
-This hint doesn't belong to the Cairo common library
-It's only added for testing proposes
-*/
-
-#[cfg(feature = "skip_next_instruction_hint")]
+/// Prevent the execution of the next instruction
+///
+/// This hint doesn't belong to the Cairo common library
+/// It's only added for testing purposes
 pub fn skip_next_instruction(vm: &mut VirtualMachine) -> Result<(), HintError> {
     vm.skip_next_instruction_execution();
     Ok(())

src/hint_processor/hint_processor_utils.rs

@@ -58,7 +58,7 @@ pub fn get_ptr_from_reference(
     }
 }
 
-//Returns the value given by a reference as an Option<MaybeRelocatable>
+///Returns the value given by a reference as [MaybeRelocatable]
 pub fn get_maybe_relocatable_from_reference(
     vm: &VirtualMachine,
     hint_reference: &HintReference,
@@ -80,7 +80,7 @@ pub fn get_maybe_relocatable_from_reference(
     value.ok_or(HintError::FailedToGetIds)
 }
 
-///Computes the memory address of the ids variable indicated by the HintReference as a Relocatable
+///Computes the memory address of the ids variable indicated by the HintReference as a [Relocatable]
 pub fn compute_addr_from_reference(
     //Reference data of the ids variable
     hint_reference: &HintReference,

src/lib.rs

@@ -1,3 +1,11 @@
+//! An implementation of the Cairo virtual machine
+//!
+//! # Feature Flags
+//! - `skip_next_instruction_hint`: Enable the `skip_next_instruction()` hint. Not enabled by default.
+//! - `hooks`: Enable [Hooks](vm::hooks) support for the [VirtualMachine](vm::vm_core::VirtualMachine). Not enabled by default.
+//! - `with_mimalloc`: Use [MiMalloc](https://crates.io/crates/mimalloc) as the program global allocator.
+
+#![cfg_attr(docsrs, feature(doc_cfg))]
 #![deny(warnings)]
 pub mod cairo_run;
 pub mod hint_processor;

src/vm/hooks.rs

@@ -0,0 +1,266 @@
+//! VM hooks
+//!
+//! Make it possible to execute custom arbitrary code at different stages of the VM execution
+//!
+//! If added to the VM, hooks function will be called during the VM execution at specific stages.
+//!
+//! Available hooks:
+//! - before_first_step, executed before entering the execution loop in [run_until_pc](CairoRunner::run_until_pc)
+//! - pre_step_instruction, executed before each instruction_step in [step](VirtualMachine::step)
+//! - post_step_instruction, executed after each instruction_step in [step](VirtualMachine::step)
+
+use std::{any::Any, collections::HashMap, sync::Arc};
+
+use felt::Felt;
+
+use crate::{
+    hint_processor::hint_processor_definition::HintProcessor, types::exec_scope::ExecutionScopes,
+};
+
+use super::{
+    errors::vm_errors::VirtualMachineError, runners::cairo_runner::CairoRunner,
+    vm_core::VirtualMachine,
+};
+
+type BeforeFirstStepHookFunc = Arc<
+    dyn Fn(
+            &mut VirtualMachine,
+            &mut CairoRunner,
+            &HashMap<usize, Vec<Box<dyn Any>>>,
+        ) -> Result<(), VirtualMachineError>
+        + Sync
+        + Send,
+>;
+
+type StepHookFunc = Arc<
+    dyn Fn(
+            &mut VirtualMachine,
+            &mut dyn HintProcessor,
+            &mut ExecutionScopes,
+            &HashMap<usize, Vec<Box<dyn Any>>>,
+            &HashMap<String, Felt>,
+        ) -> Result<(), VirtualMachineError>
+        + Sync
+        + Send,
+>;
+
+/// The hooks to be executed during the VM run
+///
+/// They can be individually ignored by setting them to [None]
+#[derive(Clone, Default)]
+pub struct Hooks {
+    before_first_step: Option<BeforeFirstStepHookFunc>,
+    pre_step_instruction: Option<StepHookFunc>,
+    post_step_instruction: Option<StepHookFunc>,
+}
+
+impl Hooks {
+    pub fn new(
+        before_first_step: Option<BeforeFirstStepHookFunc>,
+        pre_step_instruction: Option<StepHookFunc>,
+        post_step_instruction: Option<StepHookFunc>,
+    ) -> Self {
+        Hooks {
+            before_first_step,
+            pre_step_instruction,
+            post_step_instruction,
+        }
+    }
+}
+
+impl VirtualMachine {
+    pub fn execute_before_first_step(
+        &mut self,
+        runner: &mut CairoRunner,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.hooks.clone().before_first_step {
+            (hook_func)(self, runner, hint_data_dictionary)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn execute_pre_step_instruction(
+        &mut self,
+        hint_executor: &mut dyn HintProcessor,
+        exec_scope: &mut ExecutionScopes,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+        constants: &HashMap<String, Felt>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.hooks.clone().pre_step_instruction {
+            (hook_func)(
+                self,
+                hint_executor,
+                exec_scope,
+                hint_data_dictionary,
+                constants,
+            )?;
+        }
+
+        Ok(())
+    }
+
+    pub fn execute_post_step_instruction(
+        &mut self,
+        hint_executor: &mut dyn HintProcessor,
+        exec_scope: &mut ExecutionScopes,
+        hint_data_dictionary: &HashMap<usize, Vec<Box<dyn Any>>>,
+        constants: &HashMap<String, Felt>,
+    ) -> Result<(), VirtualMachineError> {
+        if let Some(hook_func) = self.hooks.clone().post_step_instruction {
+            (hook_func)(
+                self,
+                hint_executor,
+                exec_scope,
+                hint_data_dictionary,
+                constants,
+            )?;
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use super::*;
+    use crate::{
+        hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor,
+        types::program::Program,
+        utils::test_utils::{cairo_runner, vm},
+    };
+
+    #[test]
+    fn empty_hooks() {
+        let program = Program::from_file(Path::new("cairo_programs/sqrt.json"), Some("main"))
+            .expect("Call to `Program::from_file()` failed.");
+
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program);
+        let mut vm = vm!();
+        vm.hooks = Hooks::new(None, None, None);
+
+        let end = cairo_runner.initialize(&mut vm).unwrap();
+        assert!(cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .is_ok());
+    }
+
+    #[test]
+    fn hook_failure() {
+        let program = Program::from_file(Path::new("cairo_programs/sqrt.json"), Some("main"))
+            .expect("Call to `Program::from_file()` failed.");
+
+        fn before_first_step_hook(
+            _vm: &mut VirtualMachine,
+            _runner: &mut CairoRunner,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+        ) -> Result<(), VirtualMachineError> {
+            Err(VirtualMachineError::Unexpected)
+        }
+
+        fn pre_step_hook(
+            _vm: &mut VirtualMachine,
+            _hint_processor: &mut dyn HintProcessor,
+            _exec_scope: &mut ExecutionScopes,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+            _constants: &HashMap<String, Felt>,
+        ) -> Result<(), VirtualMachineError> {
+            Err(VirtualMachineError::Unexpected)
+        }
+
+        fn post_step_hook(
+            _vm: &mut VirtualMachine,
+            _hint_processor: &mut dyn HintProcessor,
+            _exec_scope: &mut ExecutionScopes,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+            _constants: &HashMap<String, Felt>,
+        ) -> Result<(), VirtualMachineError> {
+            Err(VirtualMachineError::Unexpected)
+        }
+
+        // Before first fail
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program);
+        let mut vm = vm!();
+        vm.hooks = Hooks::new(Some(Arc::new(before_first_step_hook)), None, None);
+
+        let end = cairo_runner.initialize(&mut vm).unwrap();
+        assert!(cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .is_err());
+
+        // Pre step fail
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program);
+        let mut vm = vm!();
+        vm.hooks = Hooks::new(None, Some(Arc::new(pre_step_hook)), None);
+
+        let end = cairo_runner.initialize(&mut vm).unwrap();
+        assert!(cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .is_err());
+
+        // Post step fail
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program);
+        let mut vm = vm!();
+        vm.hooks = Hooks::new(None, None, Some(Arc::new(post_step_hook)));
+
+        let end = cairo_runner.initialize(&mut vm).unwrap();
+        assert!(cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .is_err());
+    }
+
+    #[test]
+    fn hook_success() {
+        let program = Program::from_file(Path::new("cairo_programs/sqrt.json"), Some("main"))
+            .expect("Call to `Program::from_file()` failed.");
+
+        fn before_first_step_hook(
+            _vm: &mut VirtualMachine,
+            _runner: &mut CairoRunner,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+        ) -> Result<(), VirtualMachineError> {
+            Ok(())
+        }
+
+        fn pre_step_hook(
+            _vm: &mut VirtualMachine,
+            _hint_processor: &mut dyn HintProcessor,
+            _exec_scope: &mut ExecutionScopes,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+            _constants: &HashMap<String, Felt>,
+        ) -> Result<(), VirtualMachineError> {
+            Ok(())
+        }
+
+        fn post_step_hook(
+            _vm: &mut VirtualMachine,
+            _hint_processor: &mut dyn HintProcessor,
+            _exec_scope: &mut ExecutionScopes,
+            _hint_data: &HashMap<usize, Vec<Box<dyn Any>>>,
+            _constants: &HashMap<String, Felt>,
+        ) -> Result<(), VirtualMachineError> {
+            Ok(())
+        }
+
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program);
+        let mut vm = vm!();
+        vm.hooks = 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(&mut vm).unwrap();
+        assert!(cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .is_ok());
+    }
+}

src/vm/mod.rs

@@ -6,3 +6,7 @@ pub mod security;
 pub mod trace;
 pub mod vm_core;
 pub mod vm_memory;
+
+#[cfg(any(feature = "hooks"))]
+#[cfg_attr(docsrs, doc(cfg(feature = "hooks")))]
+pub mod hooks;

src/vm/runners/cairo_runner.rs

@@ -526,6 +526,8 @@ impl CairoRunner {
     ) -> Result<(), VirtualMachineError> {
         let references = self.get_reference_list();
         let hint_data_dictionary = self.get_hint_data_dictionary(&references, hint_processor)?;
+        #[cfg(feature = "hooks")]
+        vm.execute_before_first_step(self, &hint_data_dictionary)?;
         while vm.run_context.pc != address {
             vm.step(
                 hint_processor,