feat: allow tracers to stop/suspend execution

crates/vm2-interface/src/lib.rs

@@ -91,3 +91,41 @@ pub use self::{state_interface::*, tracer_interface::*};
 
 mod state_interface;
 mod tracer_interface;
+
+/// Returned from Tracer::after_instruction to indicate whether the VM should continue execution.
+#[derive(Debug)]
+pub enum ExecutionStatus {
+    /// Continue execution.
+    Running,
+    /// Stop or suspend execution.
+    Stopped(ExecutionEnd),
+}
+
+impl ExecutionStatus {
+    /// Combines ExecutionStatus values from two sources, choosing the most severe one.
+    /// So if one tracer wants to suspend but the other wants to panic, the VM will panic.
+    pub fn merge(self, other: Self) -> Self {
+        match (self, other) {
+            (
+                ExecutionStatus::Stopped(ExecutionEnd::SuspendedOnHook(_)),
+                ExecutionStatus::Stopped(end),
+            ) => ExecutionStatus::Stopped(end),
+            (ExecutionStatus::Stopped(end), _) => ExecutionStatus::Stopped(end),
+            (_, ExecutionStatus::Stopped(end)) => ExecutionStatus::Stopped(end),
+            _ => ExecutionStatus::Running,
+        }
+    }
+}
+
+/// VM stop reason and return value.
+#[derive(Debug, PartialEq)]
+pub enum ExecutionEnd {
+    /// The executed program has finished and returned the specified data.
+    ProgramFinished(Vec<u8>),
+    /// The executed program has reverted returning the specified data.
+    Reverted(Vec<u8>),
+    /// The executed program has panicked.
+    Panicked,
+    /// Returned when the bootloader writes to the heap location specified by [`hook_address`](crate::Settings.hook_address).
+    SuspendedOnHook(u32),
+}

crates/vm2-interface/src/tracer_interface.rs

@@ -1,4 +1,4 @@
-use crate::GlobalStateInterface;
+use crate::{ExecutionStatus, GlobalStateInterface};
 
 macro_rules! forall_simple_opcodes {
     ($m:ident) => {
@@ -263,8 +263,12 @@ pub trait Tracer {
     /// Executes logic after an instruction handler.
     ///
     /// The default implementation does nothing.
-    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {
+    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(
+        &mut self,
+        state: &mut S,
+    ) -> ExecutionStatus {
         let _ = state;
+        ExecutionStatus::Running
     }
 
     /// Provides cycle statistics for "complex" instructions from the prover perspective (mostly precompile calls).
@@ -302,9 +306,13 @@ impl<A: Tracer, B: Tracer> Tracer for (A, B) {
         self.1.before_instruction::<OP, S>(state);
     }
 
-    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {
-        self.0.after_instruction::<OP, S>(state);
-        self.1.after_instruction::<OP, S>(state);
+    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(
+        &mut self,
+        state: &mut S,
+    ) -> ExecutionStatus {
+        self.0
+            .after_instruction::<OP, S>(state)
+            .merge(self.1.after_instruction::<OP, S>(state))
     }
 
     fn on_extra_prover_cycles(&mut self, stats: CycleStats) {

crates/vm2/src/decode.rs

@@ -11,7 +11,7 @@ use zksync_vm2_interface::{
         self, Add, And, Div, Mul, Or, PointerAdd, PointerPack, PointerShrink, PointerSub,
         RotateLeft, RotateRight, ShiftLeft, ShiftRight, Sub, Xor,
     },
-    Tracer,
+    ExecutionEnd, ExecutionStatus, Tracer,
 };
 
 use crate::{
@@ -20,7 +20,6 @@ use crate::{
         Immediate1, Immediate2, Register, Register1, Register2, RegisterAndImmediate,
         RelativeStack, SourceWriter,
     },
-    instruction::{ExecutionEnd, ExecutionStatus},
     mode_requirements::ModeRequirements,
     Instruction, Predicate, VirtualMachine, World,
 };

crates/vm2/src/instruction.rs

@@ -1,5 +1,7 @@
 use std::fmt;
 
+pub(crate) use zksync_vm2_interface::ExecutionStatus;
+
 use crate::{addressing_modes::Arguments, vm::VirtualMachine};
 
 /// Single EraVM instruction (an opcode + [`Arguments`]).
@@ -21,22 +23,3 @@ impl<T, W> fmt::Debug for Instruction<T, W> {
 }
 
 pub(crate) type Handler<T, W> = fn(&mut VirtualMachine<T, W>, &mut W, &mut T) -> ExecutionStatus;
-
-#[derive(Debug)]
-pub(crate) enum ExecutionStatus {
-    Running,
-    Stopped(ExecutionEnd),
-}
-
-/// VM stop reason returned from [`VirtualMachine::run()`].
-#[derive(Debug, PartialEq)]
-pub enum ExecutionEnd {
-    /// The executed program has finished and returned the specified data.
-    ProgramFinished(Vec<u8>),
-    /// The executed program has reverted returning the specified data.
-    Reverted(Vec<u8>),
-    /// The executed program has panicked.
-    Panicked,
-    /// Returned when the bootloader writes to the heap location specified by [`hook_address`](crate::Settings.hook_address).
-    SuspendedOnHook(u32),
-}

crates/vm2/src/instruction_handlers/ret.rs

@@ -1,7 +1,7 @@
 use primitive_types::U256;
 use zksync_vm2_interface::{
     opcodes::{self, Normal, Panic, Revert, TypeLevelReturnType},
-    ReturnType, Tracer,
+    ExecutionEnd, ReturnType, Tracer,
 };
 
 use super::{
@@ -12,7 +12,7 @@ use super::{
 use crate::{
     addressing_modes::{Arguments, Immediate1, Register1, Source, INVALID_INSTRUCTION_COST},
     callframe::FrameRemnant,
-    instruction::{ExecutionEnd, ExecutionStatus},
+    instruction::ExecutionStatus,
     mode_requirements::ModeRequirements,
     predication::Flags,
     tracing::VmAndWorld,

crates/vm2/src/lib.rs

@@ -13,13 +13,14 @@ use zksync_vm2_interface::Tracer;
 pub(crate) use self::single_instruction_test::{heap, program, stack};
 pub use self::{
     fat_pointer::FatPointer,
-    instruction::{ExecutionEnd, Instruction},
+    instruction::Instruction,
     mode_requirements::ModeRequirements,
     predication::Predicate,
     program::Program,
     vm::{Settings, VirtualMachine},
     world_diff::{Snapshot, StorageChange, WorldDiff},
 };
+pub use zksync_vm2_interface::ExecutionEnd;
 
 pub mod addressing_modes;
 #[cfg(not(feature = "single_instruction_test"))]

crates/vm2/src/vm.rs

@@ -1,12 +1,13 @@
 use std::fmt;
 
 use primitive_types::H160;
-use zksync_vm2_interface::{opcodes::TypeLevelCallingMode, CallingMode, HeapId, Tracer};
+use zksync_vm2_interface::{
+    opcodes::TypeLevelCallingMode, CallingMode, ExecutionStatus, HeapId, Tracer,
+};
 
 use crate::{
     callframe::{Callframe, FrameRemnant},
     decommit::u256_into_address,
-    instruction::ExecutionStatus,
     stack::StackPool,
     state::{State, StateSnapshot},
     world_diff::{ExternalSnapshot, Snapshot, WorldDiff},