From 457d8a7eea9093af9440662e33e598c13ba41633 Mon Sep 17 00:00:00 2001 From: Joonatan Saarhelo Date: Thu, 31 Oct 2024 12:20:09 +0000 Subject: [PATCH] feat: allow tracers to stop/suspend execution (#78) Tracers can now request to a success / revert / panic in `Tracer::after_instruction` Makes validation tracer implementation nicer. --- crates/vm2-interface/src/lib.rs | 9 +++- crates/vm2-interface/src/tracer_interface.rs | 45 ++++++++++++++++--- crates/vm2/src/instruction.rs | 24 ++++++++++ crates/vm2/src/instruction_handlers/common.rs | 10 ++--- .../vm2/src/instruction_handlers/far_call.rs | 9 ++-- crates/vm2/src/instruction_handlers/ret.rs | 13 +++--- 6 files changed, 87 insertions(+), 23 deletions(-) diff --git a/crates/vm2-interface/src/lib.rs b/crates/vm2-interface/src/lib.rs index 29461be..dfea319 100644 --- a/crates/vm2-interface/src/lib.rs +++ b/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(&mut self, state: &mut S) {} -//! fn after_instruction(&mut self, state: &mut S) {} +//! fn after_instruction(&mut self, state: &mut S) -> ShouldStop { +//! ShouldStop::Continue +//! } //! } //! //! impl Tracer for T { @@ -73,7 +76,9 @@ //! } //! } //! } -//! fn after_instruction(&mut self, state: &mut S) {} +//! fn after_instruction(&mut self, state: &mut S) -> ShouldStop { +//! todo!() +//! } //! } //! //! // Now you can use the new features by implementing TracerV2 diff --git a/crates/vm2-interface/src/tracer_interface.rs b/crates/vm2-interface/src/tracer_interface.rs index ab1abfb..29539b9 100644 --- a/crates/vm2-interface/src/tracer_interface.rs +++ b/crates/vm2-interface/src/tracer_interface.rs @@ -254,17 +254,26 @@ impl OpcodeType for opcodes::Ret { /// } /// ``` 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(&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(&mut self, state: &mut S) { + #[must_use] + fn after_instruction( + &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 Tracer for (A, B) { self.1.before_instruction::(state); } - fn after_instruction(&mut self, state: &mut S) { - self.0.after_instruction::(state); - self.1.after_instruction::(state); + fn after_instruction( + &mut self, + state: &mut S, + ) -> ShouldStop { + self.0 + .after_instruction::(state) + .merge(self.1.after_instruction::(state)) } fn on_extra_prover_cycles(&mut self, stats: CycleStats) { diff --git a/crates/vm2/src/instruction.rs b/crates/vm2/src/instruction.rs index 530ae00..230c7a1 100644 --- a/crates/vm2/src/instruction.rs +++ b/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 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, } diff --git a/crates/vm2/src/instruction_handlers/common.rs b/crates/vm2/src/instruction_handlers/common.rs index c46ef47..2576e25 100644 --- a/crates/vm2/src/instruction_handlers/common.rs +++ b/crates/vm2/src/instruction_handlers/common.rs @@ -58,13 +58,13 @@ pub(crate) fn full_boilerplate>( if args.predicate().satisfied(&vm.state.flags) { tracer.before_instruction::(&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::(&mut VmAndWorld { vm, world }); - result + business_logic(vm, args, world, tracer) + .merge_tracer(tracer.after_instruction::(&mut VmAndWorld { vm, world })) } else { tracer.before_instruction::(&mut VmAndWorld { vm, world }); vm.state.current_frame.pc = unsafe { vm.state.current_frame.pc.add(1) }; - tracer.after_instruction::(&mut VmAndWorld { vm, world }); - ExecutionStatus::Running + tracer + .after_instruction::(&mut VmAndWorld { vm, world }) + .into() } } diff --git a/crates/vm2/src/instruction_handlers/far_call.rs b/crates/vm2/src/instruction_handlers/far_call.rs index 2c000b1..494d252 100644 --- a/crates/vm2/src/instruction_handlers/far_call.rs +++ b/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, M: TypeLevelCallingMode, { - boilerplate_ext::, _, _>(vm, world, tracer, |vm, args, world, tracer| { + full_boilerplate::, _, _>(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 }) } diff --git a/crates/vm2/src/instruction_handlers/ret.rs b/crates/vm2/src/instruction_handlers/ret.rs index 382949d..17c3699 100644 --- a/crates/vm2/src/instruction_handlers/ret.rs +++ b/crates/vm2/src/instruction_handlers/ret.rs @@ -146,12 +146,11 @@ pub(crate) fn free_panic>( ) -> ExecutionStatus { tracer.before_instruction::, _>(&mut VmAndWorld { vm, world }); // args aren't used for panics unless TO_LABEL - let result = naked_ret::( + naked_ret::( vm, &Arguments::new(Predicate::Always, 0, ModeRequirements::none()), - ); - tracer.after_instruction::, _>(&mut VmAndWorld { vm, world }); - result + ) + .merge_tracer(tracer.after_instruction::, _>(&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>( world: &mut W, tracer: &mut T, exception_handler: u16, -) { +) -> ExecutionStatus { tracer.before_instruction::, _>(&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>( vm.state.flags = Flags::new(true, false, false); vm.state.current_frame.set_pc_from_u16(exception_handler); - tracer.after_instruction::, _>(&mut VmAndWorld { vm, world }); + tracer + .after_instruction::, _>(&mut VmAndWorld { vm, world }) + .into() } fn invalid>(