feat(vm): Improve tracer trait

core/bin/system-constants-generator/src/utils.rs

@@ -3,9 +3,9 @@ use std::cell::RefCell;
 use std::rc::Rc;
 use vm::constants::{BLOCK_GAS_LIMIT, BOOTLOADER_HEAP_PAGE};
 use vm::{
-    BootloaderState, BoxedTracer, DynTracer, ExecutionEndTracer, ExecutionProcessing,
-    HistoryEnabled, HistoryMode, L1BatchEnv, L2BlockEnv, SystemEnv, TxExecutionMode, Vm,
-    VmExecutionMode, VmExecutionStopReason, VmTracer, ZkSyncVmState,
+    BootloaderState, BoxedTracer, DynTracer, HistoryEnabled, HistoryMode, L1BatchEnv, L2BlockEnv,
+    SystemEnv, TxExecutionMode, Vm, VmExecutionMode, VmExecutionStopReason, VmTracer,
+    ZkSyncVmState,
 };
 use zksync_contracts::{
     load_sys_contract, read_bootloader_code, read_sys_contract_bytecode, read_zbin_bytecode,
@@ -33,9 +33,7 @@ struct SpecialBootloaderTracer {
 
 impl<S: WriteStorage, H: HistoryMode> DynTracer<S, H> for SpecialBootloaderTracer {}
 
-impl<H: HistoryMode> ExecutionEndTracer<H> for SpecialBootloaderTracer {}
-
-impl<S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H> for SpecialBootloaderTracer {
+impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for SpecialBootloaderTracer {
     fn initialize_tracer(&mut self, state: &mut ZkSyncVmState<S, H>) {
         state.memory.populate_page(
             BOOTLOADER_HEAP_PAGE as usize,
@@ -55,8 +53,6 @@ impl<S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H> for SpecialBootl
     }
 }
 
-impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for SpecialBootloaderTracer {}
-
 pub static GAS_TEST_SYSTEM_CONTRACTS: Lazy<BaseSystemContracts> = Lazy::new(|| {
     let bytecode = read_bootloader_code("gas_test");
     let hash = hash_bytecode(&bytecode);

core/lib/vm/src/implementation/execution.rs

@@ -6,9 +6,7 @@ use crate::old_vm::{
     utils::{vm_may_have_ended_inner, VmExecutionResult},
 };
 use crate::tracers::{
-    traits::{
-        BoxedTracer, ExecutionEndTracer, ExecutionProcessing, TracerExecutionStatus, VmTracer,
-    },
+    traits::{TracerExecutionStatus, VmTracer},
     DefaultExecutionTracer, RefundsTracer,
 };
 use crate::types::{inputs::VmExecutionMode, outputs::VmExecutionResultAndLogs};
@@ -18,16 +16,17 @@ use crate::VmExecutionStopReason;
 impl<S: WriteStorage, H: HistoryMode> Vm<S, H> {
     pub(crate) fn inspect_inner(
         &mut self,
-        mut tracers: Vec<Box<dyn VmTracer<S, H>>>,
+        tracers: Vec<Box<dyn VmTracer<S, H>>>,
         execution_mode: VmExecutionMode,
     ) -> VmExecutionResultAndLogs {
+        let mut enable_refund_tracer = false;
         if let VmExecutionMode::OneTx = execution_mode {
-            // For correct results we have to include refunds tracer to the desired tracers
-            tracers.push(RefundsTracer::new(self.batch_env.clone()).into_boxed());
             // Move the pointer to the next transaction
             self.bootloader_state.move_tx_to_execute_pointer();
+            enable_refund_tracer = true;
         }
-        let (_, result) = self.inspect_and_collect_results(tracers, execution_mode);
+        let (_, result) =
+            self.inspect_and_collect_results(tracers, execution_mode, enable_refund_tracer);
         result
     }
 
@@ -37,12 +36,16 @@ impl<S: WriteStorage, H: HistoryMode> Vm<S, H> {
         &mut self,
         tracers: Vec<Box<dyn VmTracer<S, H>>>,
         execution_mode: VmExecutionMode,
+        with_refund_tracer: bool,
     ) -> (VmExecutionStopReason, VmExecutionResultAndLogs) {
+        let refund_tracers =
+            with_refund_tracer.then_some(RefundsTracer::new(self.batch_env.clone()));
         let mut tx_tracer: DefaultExecutionTracer<S, H> = DefaultExecutionTracer::new(
             self.system_env.default_validation_computational_gas_limit,
             execution_mode,
             tracers,
             self.storage.clone(),
+            refund_tracers,
         );
 
         let timestamp_initial = Timestamp(self.state.local_state.timestamp);
@@ -68,16 +71,18 @@ impl<S: WriteStorage, H: HistoryMode> Vm<S, H> {
 
         let result = tx_tracer.result_tracer.into_result();
 
-        let mut result = VmExecutionResultAndLogs {
+        let refunds = tx_tracer
+            .refund_tracer
+            .map(|x| x.get_refunds())
+            .unwrap_or_default();
+
+        let result = VmExecutionResultAndLogs {
             result,
             logs,
             statistics,
-            refunds: Default::default(),
+            refunds,
         };
 
-        for tracer in tx_tracer.custom_tracers.iter_mut() {
-            tracer.save_results(&mut result);
-        }
         (stop_reason, result)
     }
 
@@ -96,19 +101,18 @@ impl<S: WriteStorage, H: HistoryMode> Vm<S, H> {
                 self.state
             );
 
-            tracer.before_cycle(&mut self.state);
             self.state
                 .cycle(tracer)
                 .expect("Failed execution VM cycle.");
 
-            tracer.after_cycle(&mut self.state, &mut self.bootloader_state);
+            if let TracerExecutionStatus::Stop(reason) =
+                tracer.finish_cycle(&mut self.state, &mut self.bootloader_state)
+            {
+                break VmExecutionStopReason::TracerRequestedStop(reason);
+            }
             if self.has_ended() {
                 break VmExecutionStopReason::VmFinished;
             }
-
-            if let TracerExecutionStatus::Stop(reason) = tracer.should_stop_execution() {
-                break VmExecutionStopReason::TracerRequestedStop(reason);
-            }
         };
         tracer.after_vm_execution(&mut self.state, &self.bootloader_state, result.clone());
         result

core/lib/vm/src/lib.rs

@@ -16,10 +16,7 @@ pub use errors::{
 
 pub use tracers::{
     call::CallTracer,
-    traits::{
-        BoxedTracer, DynTracer, ExecutionEndTracer, ExecutionProcessing, TracerExecutionStatus,
-        TracerExecutionStopReason, VmTracer,
-    },
+    traits::{BoxedTracer, DynTracer, TracerExecutionStatus, TracerExecutionStopReason, VmTracer},
     utils::VmExecutionStopReason,
     StorageInvocations, ValidationError, ValidationTracer, ValidationTracerParams,
 };

core/lib/vm/src/tests/rollbacks.rs

@@ -12,9 +12,8 @@ use crate::tests::tester::{
 use crate::tests::utils::read_test_contract;
 use crate::types::inputs::system_env::TxExecutionMode;
 use crate::{
-    BootloaderState, DynTracer, ExecutionEndTracer, ExecutionProcessing, HistoryEnabled,
-    HistoryMode, TracerExecutionStatus, TracerExecutionStopReason, VmExecutionMode, VmTracer,
-    ZkSyncVmState,
+    BootloaderState, DynTracer, HistoryEnabled, HistoryMode, TracerExecutionStatus,
+    TracerExecutionStopReason, VmExecutionMode, VmTracer, ZkSyncVmState,
 };
 
 #[test]
@@ -153,39 +152,28 @@ fn test_vm_loadnext_rollbacks() {
 // Testing tracer that does not allow the recursion to go deeper than a certain limit
 struct MaxRecursionTracer {
     max_recursion_depth: usize,
-    should_stop_execution: bool,
 }
 
 /// Tracer responsible for calculating the number of storage invocations and
 /// stopping the VM execution if the limit is reached.
 impl<S, H: HistoryMode> DynTracer<S, H> for MaxRecursionTracer {}
 
-impl<H: HistoryMode> ExecutionEndTracer<H> for MaxRecursionTracer {
-    fn should_stop_execution(&self) -> TracerExecutionStatus {
-        if self.should_stop_execution {
-            TracerExecutionStatus::Stop(TracerExecutionStopReason::Finish)
-        } else {
-            TracerExecutionStatus::Continue
-        }
-    }
-}
-
-impl<S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H> for MaxRecursionTracer {
-    fn after_cycle(
+impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for MaxRecursionTracer {
+    fn finish_cycle(
         &mut self,
         state: &mut ZkSyncVmState<S, H>,
         _bootloader_state: &mut BootloaderState,
-    ) {
+    ) -> TracerExecutionStatus {
         let current_depth = state.local_state.callstack.depth();
 
         if current_depth > self.max_recursion_depth {
-            self.should_stop_execution = true;
+            TracerExecutionStatus::Stop(TracerExecutionStopReason::Finish)
+        } else {
+            TracerExecutionStatus::Continue
         }
     }
 }
 
-impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for MaxRecursionTracer {}
-
 #[test]
 fn test_layered_rollback() {
     // This test checks that the layered rollbacks work correctly, i.e.
@@ -236,7 +224,6 @@ fn test_layered_rollback() {
     vm.vm.inspect(
         vec![Box::new(MaxRecursionTracer {
             max_recursion_depth: 15,
-            should_stop_execution: false,
         })],
         VmExecutionMode::OneTx,
     );

core/lib/vm/src/tracers/call.rs

@@ -16,8 +16,8 @@ use zksync_types::U256;
 use crate::errors::VmRevertReason;
 use crate::old_vm::history_recorder::HistoryMode;
 use crate::old_vm::memory::SimpleMemory;
-use crate::tracers::traits::{DynTracer, ExecutionEndTracer, ExecutionProcessing, VmTracer};
-use crate::types::outputs::VmExecutionResultAndLogs;
+use crate::tracers::traits::{DynTracer, VmTracer};
+use crate::{BootloaderState, VmExecutionStopReason, ZkSyncVmState};
 
 #[derive(Debug, Clone)]
 pub struct CallTracer<H: HistoryMode> {
@@ -88,12 +88,13 @@ impl<S, H: HistoryMode> DynTracer<S, H> for CallTracer<H> {
     }
 }
 
-impl<H: HistoryMode> ExecutionEndTracer<H> for CallTracer<H> {}
-
-impl<S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H> for CallTracer<H> {}
-
 impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for CallTracer<H> {
-    fn save_results(&mut self, _result: &mut VmExecutionResultAndLogs) {
+    fn after_vm_execution(
+        &mut self,
+        _state: &mut ZkSyncVmState<S, H>,
+        _bootloader_state: &BootloaderState,
+        _stop_reason: VmExecutionStopReason,
+    ) {
         self.result
             .set(
                 std::mem::take(&mut self.stack)