feat: allow tracers to stop execution

crates/vm2-interface/src/lib.rs

@@ -28,6 +28,7 @@
 //! # use zksync_vm2_interface as zksync_vm2_interface_v1;
 //! use zksync_vm2_interface_v1::{
 //!     StateInterface as StateInterfaceV1, GlobalStateInterface as GlobalStateInterfaceV1, Tracer as TracerV1, opcodes::NearCall,
+//!     ShouldStop,
 //! };
 //!
 //! trait StateInterface: StateInterfaceV1 {
@@ -60,7 +61,9 @@
 //!
 //! trait Tracer {
 //!     fn before_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {}
-//!     fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {}
+//!     fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) -> ShouldStop {
+//!         ShouldStop::Continue
+//!     }
 //! }
 //!
 //! impl<T: TracerV1> Tracer for T {
@@ -73,7 +76,9 @@
 //!             }
 //!         }
 //!     }
-//!     fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {}
+//!     fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) -> ShouldStop {
+//!         todo!()
+//!     }
 //! }
 //!
 //! // Now you can use the new features by implementing TracerV2

crates/vm2-interface/src/tracer_interface.rs

@@ -254,17 +254,26 @@ impl<T: opcodes::TypeLevelReturnType> OpcodeType for opcodes::Ret<T> {
 /// }
 /// ```
 pub trait Tracer {
-    /// Executes logic before an instruction handler.
+    /// This method is executed before an instruction handler.
     ///
     /// The default implementation does nothing.
     fn before_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {
         let _ = state;
     }
-    /// Executes logic after an instruction handler.
+    /// This method is executed after an instruction handler.
+    ///
+    /// The return value indicates whether the VM should continue or stop execution.
+    /// The tracer's return value takes precedence over the VM but only if it is at least as severe.
+    /// For example, if the VM wants to stop and the tracer wants to suspend, the VM will still stop.
     ///
     /// The default implementation does nothing.
-    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(&mut self, state: &mut S) {
+    #[must_use]
+    fn after_instruction<OP: OpcodeType, S: GlobalStateInterface>(
+        &mut self,
+        state: &mut S,
+    ) -> ShouldStop {
         let _ = state;
+        ShouldStop::Continue
     }
 
     /// Provides cycle statistics for "complex" instructions from the prover perspective (mostly precompile calls).
@@ -273,6 +282,26 @@ pub trait Tracer {
     fn on_extra_prover_cycles(&mut self, _stats: CycleStats) {}
 }
 
+/// Returned from [`Tracer::after_instruction`] to indicate if the VM should stop.
+#[derive(Debug)]
+pub enum ShouldStop {
+    /// The VM should stop.
+    Stop,
+    /// The VM should continue.
+    Continue,
+}
+
+impl ShouldStop {
+    #[must_use]
+    #[inline(always)]
+    fn merge(self, other: ShouldStop) -> ShouldStop {
+        match (self, other) {
+            (ShouldStop::Continue, ShouldStop::Continue) => ShouldStop::Continue,
+            _ => ShouldStop::Stop,
+        }
+    }
+}
+
 /// Cycle statistics emitted by the VM and supplied to [`Tracer::on_extra_prover_cycles()`].
 #[derive(Debug, Clone, Copy)]
 pub enum CycleStats {
@@ -302,9 +331,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,
+    ) -> ShouldStop {
+        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/instruction.rs

@@ -1,5 +1,7 @@
 use std::fmt;
 
+use zksync_vm2_interface::ShouldStop;
+
 use crate::{addressing_modes::Arguments, vm::VirtualMachine};
 
 /// Single EraVM instruction (an opcode + [`Arguments`]).
@@ -28,6 +30,26 @@ pub(crate) enum ExecutionStatus {
     Stopped(ExecutionEnd),
 }
 
+impl ExecutionStatus {
+    #[must_use]
+    #[inline(always)]
+    pub(crate) fn merge_tracer(self, should_stop: ShouldStop) -> Self {
+        match (&self, should_stop) {
+            (Self::Running, ShouldStop::Stop) => Self::Stopped(ExecutionEnd::StoppedByTracer),
+            _ => self,
+        }
+    }
+}
+
+impl From<ShouldStop> for ExecutionStatus {
+    fn from(should_stop: ShouldStop) -> Self {
+        match should_stop {
+            ShouldStop::Stop => Self::Stopped(ExecutionEnd::StoppedByTracer),
+            ShouldStop::Continue => Self::Running,
+        }
+    }
+}
+
 /// VM stop reason returned from [`VirtualMachine::run()`].
 #[derive(Debug, PartialEq)]
 pub enum ExecutionEnd {
@@ -39,4 +61,6 @@ pub enum ExecutionEnd {
     Panicked,
     /// Returned when the bootloader writes to the heap location specified by [`hook_address`](crate::Settings.hook_address).
     SuspendedOnHook(u32),
+    /// One of the tracers decided it is time to stop the VM.
+    StoppedByTracer,
 }

crates/vm2/src/instruction_handlers/common.rs

@@ -58,13 +58,13 @@ pub(crate) fn full_boilerplate<Opcode: OpcodeType, T: Tracer, W: World<T>>(
     if args.predicate().satisfied(&vm.state.flags) {
         tracer.before_instruction::<Opcode, _>(&mut VmAndWorld { vm, world });
         vm.state.current_frame.pc = unsafe { vm.state.current_frame.pc.add(1) };
-        let result = business_logic(vm, args, world, tracer);
-        tracer.after_instruction::<Opcode, _>(&mut VmAndWorld { vm, world });
-        result
+        business_logic(vm, args, world, tracer)
+            .merge_tracer(tracer.after_instruction::<Opcode, _>(&mut VmAndWorld { vm, world }))
     } else {
         tracer.before_instruction::<opcodes::Nop, _>(&mut VmAndWorld { vm, world });
         vm.state.current_frame.pc = unsafe { vm.state.current_frame.pc.add(1) };
-        tracer.after_instruction::<opcodes::Nop, _>(&mut VmAndWorld { vm, world });
-        ExecutionStatus::Running
+        tracer
+            .after_instruction::<opcodes::Nop, _>(&mut VmAndWorld { vm, world })
+            .into()
     }
 }

crates/vm2/src/instruction_handlers/far_call.rs

@@ -9,7 +9,7 @@ use zksync_vm2_interface::{
 };
 
 use super::{
-    common::boilerplate_ext,
+    common::full_boilerplate,
     heap_access::grow_heap,
     monomorphization::{match_boolean, monomorphize, parameterize},
     ret::{panic_from_failed_far_call, RETURN_COST},
@@ -44,7 +44,7 @@ where
     W: World<T>,
     M: TypeLevelCallingMode,
 {
-    boilerplate_ext::<FarCall<M>, _, _>(vm, world, tracer, |vm, args, world, tracer| {
+    full_boilerplate::<FarCall<M>, _, _>(vm, world, tracer, |vm, args, world, tracer| {
         let (raw_abi, raw_abi_is_pointer) = Register1::get_with_pointer_flag(args, &mut vm.state);
 
         let address_mask: U256 = U256::MAX >> (256 - 160);
@@ -109,8 +109,7 @@ where
 
         let Some((calldata, program, is_evm_interpreter)) = failing_part else {
             vm.state.current_frame.gas += new_frame_gas.saturating_sub(RETURN_COST);
-            panic_from_failed_far_call(vm, world, tracer, exception_handler);
-            return;
+            return panic_from_failed_far_call(vm, world, tracer, exception_handler);
         };
 
         let stipend = if is_evm_interpreter {
@@ -155,6 +154,8 @@ where
             | u8::from(abi.is_constructor_call);
 
         vm.state.registers[2] = call_type.into();
+
+        ExecutionStatus::Running
     })
 }
 

crates/vm2/src/instruction_handlers/ret.rs

@@ -146,12 +146,11 @@ pub(crate) fn free_panic<T: Tracer, W: World<T>>(
 ) -> ExecutionStatus {
     tracer.before_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
     // args aren't used for panics unless TO_LABEL
-    let result = naked_ret::<T, W, Panic, false>(
+    naked_ret::<T, W, Panic, false>(
         vm,
         &Arguments::new(Predicate::Always, 0, ModeRequirements::none()),
-    );
-    tracer.after_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
-    result
+    )
+    .merge_tracer(tracer.after_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world }))
 }
 
 /// Formally, a far call pushes a new frame and returns from it immediately if it panics.
@@ -161,7 +160,7 @@ pub(crate) fn panic_from_failed_far_call<T: Tracer, W: World<T>>(
     world: &mut W,
     tracer: &mut T,
     exception_handler: u16,
-) {
+) -> ExecutionStatus {
     tracer.before_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
 
     // Gas is already subtracted in the far call code.
@@ -172,7 +171,9 @@ pub(crate) fn panic_from_failed_far_call<T: Tracer, W: World<T>>(
     vm.state.flags = Flags::new(true, false, false);
     vm.state.current_frame.set_pc_from_u16(exception_handler);
 
-    tracer.after_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
+    tracer
+        .after_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world })
+        .into()
 }
 
 fn invalid<T: Tracer, W: World<T>>(