Skip to content

Commit

Permalink
pass World separately to make tracer interface work without wrapper
Browse files Browse the repository at this point in the history
  • Loading branch information
joonazan committed Oct 16, 2024
1 parent 2fead2d commit f27d9ff
Show file tree
Hide file tree
Showing 6 changed files with 101 additions and 78 deletions.
12 changes: 6 additions & 6 deletions crates/vm2-interface/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,27 +57,27 @@
//! }
//!
//! trait Tracer {
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, _state: &mut S) {}
//! fn after_instruction<OP: OpcodeType, S: StateInterface>(&mut self, _state: &mut S) {}
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, storage: &mut S::StorageInterface) {}
//! fn after_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, storage: &mut S::StorageInterface) {}
//! }
//!
//! impl<T: TracerV1> Tracer for T {
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, storage: &mut S::StorageInterface) {
//! match OP::VALUE {
//! Opcode::NewOpcode => {}
//! // Do this for every old opcode
//! Opcode::NearCall => {
//! <Self as TracerV1>::before_instruction::<NearCall, _>(self, state)
//! <Self as TracerV1>::before_instruction::<NearCall, _>(self, state, storage)
//! }
//! }
//! }
//! fn after_instruction<OP: OpcodeType, S: StateInterface>(&mut self, _state: &mut S) {}
//! fn after_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, storage: &mut S::StorageInterface) {}
//! }
//!
//! // Now you can use the new features by implementing TracerV2
//! struct MyTracer;
//! impl Tracer for MyTracer {
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
//! fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, _: &mut S::StorageInterface) {
//! if OP::VALUE == Opcode::NewOpcode {
//! state.get_some_new_field();
//! }
Expand Down
14 changes: 12 additions & 2 deletions crates/vm2-interface/src/state_interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@ use primitive_types::{H160, U256};

/// Public interface of the VM state. Encompasses both read and write methods.
pub trait StateInterface {
/// Storage interface required for operations that read storage.
type StorageInterface;

/// Reads a register with the specified zero-based index. Returns a value together with a pointer flag.
fn read_register(&self, register: u8) -> (U256, bool);
/// Sets a register with the specified zero-based index
Expand Down Expand Up @@ -42,7 +45,12 @@ pub trait StateInterface {
/// Iterates over storage slots read or written during VM execution.
fn get_storage_state(&self) -> impl Iterator<Item = ((H160, U256), U256)>;
/// Gets value of the specified storage slot.
fn get_storage(&mut self, address: H160, slot: U256) -> U256;
fn get_storage(
&mut self,
storage: &mut Self::StorageInterface,
address: H160,
slot: U256,
) -> U256;

/// Iterates over all transient storage slots set during VM execution.
fn get_transient_storage_state(&self) -> impl Iterator<Item = ((H160, U256), U256)>;
Expand Down Expand Up @@ -217,6 +225,8 @@ pub struct DummyState;

#[cfg(test)]
impl StateInterface for DummyState {
type StorageInterface = ();

fn read_register(&self, _: u8) -> (U256, bool) {
unimplemented!()
}
Expand Down Expand Up @@ -277,7 +287,7 @@ impl StateInterface for DummyState {
std::iter::empty()
}

fn get_storage(&mut self, _: H160, _: U256) -> U256 {
fn get_storage(&mut self, _: &mut Self::StorageInterface, _: H160, _: U256) -> U256 {
unimplemented!()
}

Expand Down
52 changes: 37 additions & 15 deletions crates/vm2-interface/src/tracer_interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ impl<T: opcodes::TypeLevelReturnType> OpcodeType for opcodes::Ret<T> {
/// struct FarCallCounter(usize);
///
/// impl Tracer for FarCallCounter {
/// fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
/// fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S, _: &mut S::StorageInterface) {
/// match OP::VALUE {
/// Opcode::FarCall(_) => self.0 += 1,
/// _ => {}
Expand All @@ -257,14 +257,24 @@ pub trait Tracer {
/// Executes logic before an instruction handler.
///
/// The default implementation does nothing.
fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
fn before_instruction<OP: OpcodeType, S: StateInterface>(
&mut self,
state: &mut S,
storage: &mut S::StorageInterface,
) {
let _ = state;
let _ = storage;
}
/// Executes logic after an instruction handler.
///
/// The default implementation does nothing.
fn after_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
fn after_instruction<OP: OpcodeType, S: StateInterface>(
&mut self,
state: &mut S,
storage: &mut S::StorageInterface,
) {
let _ = state;
let _ = storage;
}

/// Provides cycle statistics for "complex" instructions from the prover perspective (mostly precompile calls).
Expand Down Expand Up @@ -297,14 +307,22 @@ impl Tracer for () {}

// Multiple tracers can be combined by building a linked list out of tuples.
impl<A: Tracer, B: Tracer> Tracer for (A, B) {
fn before_instruction<OP: OpcodeType, S: StateInterface>(&mut self, state: &mut S) {
self.0.before_instruction::<OP, S>(state);
self.1.before_instruction::<OP, S>(state);
fn before_instruction<OP: OpcodeType, S: StateInterface>(
&mut self,
state: &mut S,
storage: &mut S::StorageInterface,
) {
self.0.before_instruction::<OP, S>(state, storage);
self.1.before_instruction::<OP, S>(state, storage);
}

fn after_instruction<OP: OpcodeType, S: StateInterface>(&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: StateInterface>(
&mut self,
state: &mut S,
storage: &mut S::StorageInterface,
) {
self.0.after_instruction::<OP, S>(state, storage);
self.1.after_instruction::<OP, S>(state, storage);
}

fn on_extra_prover_cycles(&mut self, stats: CycleStats) {
Expand All @@ -321,7 +339,11 @@ mod tests {
struct FarCallCounter(usize);

impl Tracer for FarCallCounter {
fn before_instruction<OP: OpcodeType, S: crate::StateInterface>(&mut self, _: &mut S) {
fn before_instruction<OP: OpcodeType, S: crate::StateInterface>(
&mut self,
_: &mut S,
_: &mut S::StorageInterface,
) {
if let super::Opcode::FarCall(CallingMode::Normal) = OP::VALUE {
self.0 += 1;
}
Expand All @@ -332,26 +354,26 @@ mod tests {
fn test_tracer() {
let mut tracer = FarCallCounter(0);

tracer.before_instruction::<opcodes::Nop, _>(&mut DummyState);
tracer.before_instruction::<opcodes::Nop, _>(&mut DummyState, &mut ());
assert_eq!(tracer.0, 0);

tracer.before_instruction::<opcodes::FarCall<opcodes::Normal>, _>(&mut DummyState);
tracer.before_instruction::<opcodes::FarCall<opcodes::Normal>, _>(&mut DummyState, &mut ());
assert_eq!(tracer.0, 1);

tracer.before_instruction::<opcodes::FarCall<opcodes::Mimic>, _>(&mut DummyState);
tracer.before_instruction::<opcodes::FarCall<opcodes::Mimic>, _>(&mut DummyState, &mut ());
assert_eq!(tracer.0, 1);
}

#[test]
fn test_aggregate_tracer() {
let mut tracer = (FarCallCounter(0), (FarCallCounter(0), FarCallCounter(0)));

tracer.before_instruction::<opcodes::Nop, _>(&mut DummyState);
tracer.before_instruction::<opcodes::Nop, _>(&mut DummyState, &mut ());
assert_eq!(tracer.0 .0, 0);
assert_eq!(tracer.1 .0 .0, 0);
assert_eq!(tracer.1 .1 .0, 0);

tracer.before_instruction::<opcodes::FarCall<opcodes::Normal>, _>(&mut DummyState);
tracer.before_instruction::<opcodes::FarCall<opcodes::Normal>, _>(&mut DummyState, &mut ());
assert_eq!(tracer.0 .0, 1);
assert_eq!(tracer.1 .0 .0, 1);
assert_eq!(tracer.1 .1 .0, 1);
Expand Down
13 changes: 5 additions & 8 deletions crates/vm2/src/instruction_handlers/common.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
use zksync_vm2_interface::{opcodes, OpcodeType, Tracer};

use super::ret::free_panic;
use crate::{
addressing_modes::Arguments, instruction::ExecutionStatus, tracing::VmAndWorld, VirtualMachine,
World,
};
use crate::{addressing_modes::Arguments, instruction::ExecutionStatus, VirtualMachine, World};

#[inline(always)]
pub(crate) fn boilerplate<Opcode: OpcodeType, T: Tracer, W: World<T>>(
Expand Down Expand Up @@ -56,15 +53,15 @@ 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 });
tracer.before_instruction::<Opcode, _>(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 });
tracer.after_instruction::<Opcode, _>(vm, world);
result
} else {
tracer.before_instruction::<opcodes::Nop, _>(&mut VmAndWorld { vm, world });
tracer.before_instruction::<opcodes::Nop, _>(vm, world);
vm.state.current_frame.pc = unsafe { vm.state.current_frame.pc.add(1) };
tracer.after_instruction::<opcodes::Nop, _>(&mut VmAndWorld { vm, world });
tracer.after_instruction::<opcodes::Nop, _>(vm, world);
ExecutionStatus::Running
}
}
9 changes: 4 additions & 5 deletions crates/vm2/src/instruction_handlers/ret.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ use crate::{
instruction::{ExecutionEnd, ExecutionStatus},
mode_requirements::ModeRequirements,
predication::Flags,
tracing::VmAndWorld,
Instruction, Predicate, VirtualMachine, World,
};

Expand Down Expand Up @@ -144,13 +143,13 @@ pub(crate) fn free_panic<T: Tracer, W: World<T>>(
world: &mut W,
tracer: &mut T,
) -> ExecutionStatus {
tracer.before_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
tracer.before_instruction::<opcodes::Ret<Panic>, _>(vm, world);
// args aren't used for panics unless TO_LABEL
let result = naked_ret::<T, W, Panic, false>(
vm,
&Arguments::new(Predicate::Always, 0, ModeRequirements::none()),
);
tracer.after_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
tracer.after_instruction::<opcodes::Ret<Panic>, _>(vm, world);
result
}

Expand All @@ -162,7 +161,7 @@ pub(crate) fn panic_from_failed_far_call<T: Tracer, W: World<T>>(
tracer: &mut T,
exception_handler: u16,
) {
tracer.before_instruction::<opcodes::Ret<Panic>, _>(&mut VmAndWorld { vm, world });
tracer.before_instruction::<opcodes::Ret<Panic>, _>(vm, world);

// Gas is already subtracted in the far call code.
// No need to roll back, as no changes are made in this "frame".
Expand All @@ -172,7 +171,7 @@ 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>, _>(vm, world);
}

fn invalid<T: Tracer, W: World<T>>(
Expand Down
Loading

0 comments on commit f27d9ff

Please sign in to comment.