feat(vm): Implement call tracing for fast VM (#2905)

## What ❔

- Implements call tracing for the fast VM.
- Enables call tracing for the batch VM executor (i.e., one used in the
state keeper) if it is requested. Compares produced traces in the shadow
VM mode.

## Why ❔

Call tracing is the last large missing part of the fast VM.

---------

Co-authored-by: Alex Ostrovski <[email protected]>

core/lib/multivm/Cargo.toml

@@ -40,5 +40,5 @@ assert_matches.workspace = true
 pretty_assertions.workspace = true
 rand.workspace = true
 test-casing.workspace = true
-zksync_test_contracts.workspace = true
 zksync_eth_signer.workspace = true
+zksync_test_contracts.workspace = true

core/lib/multivm/src/tracers/call_tracer/mod.rs

@@ -17,6 +17,8 @@ pub mod vm_virtual_blocks;
 #[derive(Debug, Clone)]
 pub struct CallTracer {
     stack: Vec<FarcallAndNearCallCount>,
+    finished_calls: Vec<Call>,
+
     result: Arc<OnceCell<Vec<Call>>>,
 
     max_stack_depth: usize,
@@ -41,6 +43,7 @@ impl CallTracer {
     pub fn new(result: Arc<OnceCell<Vec<Call>>>) -> Self {
         Self {
             stack: vec![],
+            finished_calls: vec![],
             result,
             max_stack_depth: 0,
             max_near_calls: 0,

core/lib/multivm/src/tracers/call_tracer/vm_latest/mod.rs

@@ -67,7 +67,9 @@ impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for CallTracer {
         _bootloader_state: &BootloaderState,
         _stop_reason: VmExecutionStopReason,
     ) {
-        self.store_result()
+        let result = std::mem::take(&mut self.finished_calls);
+        let cell = self.result.as_ref();
+        cell.set(result).unwrap();
     }
 }
 
@@ -191,15 +193,14 @@ impl CallTracer {
             .farcall
             .parent_gas
             .saturating_sub(state.vm_local_state.callstack.current.ergs_remaining as u64);
-
         self.save_output_latest(state, memory, ret_opcode, &mut current_call.farcall);
 
         // If there is a parent call, push the current call to it
         // Otherwise, push the current call to the stack, because it's the top level call
         if let Some(parent_call) = self.stack.last_mut() {
             parent_call.farcall.calls.push(current_call.farcall);
         } else {
-            self.push_call_and_update_stats(current_call.farcall, current_call.near_calls_after);
+            self.finished_calls.push(current_call.farcall);
         }
     }
 }

core/lib/multivm/src/versions/shadow/mod.rs

@@ -25,7 +25,7 @@ use crate::{
 mod tests;
 
 type ReferenceVm<S = InMemoryStorage> = vm_latest::Vm<StorageView<S>, HistoryEnabled>;
-type ShadowedFastVm<S = InMemoryStorage> = crate::vm_instance::ShadowedFastVm<S, (), ()>;
+type ShadowedFastVm<S = InMemoryStorage, Tr = ()> = crate::vm_instance::ShadowedFastVm<S, Tr, ()>;
 
 fn hash_block(block_env: L2BlockEnv, tx_hashes: &[H256]) -> H256 {
     let mut hasher = L2BlockHasher::new(

core/lib/multivm/src/versions/shadow/tests.rs

@@ -1,20 +1,30 @@
 //! Unit tests from the `testonly` test suite.
 
-use std::{collections::HashSet, rc::Rc};
+use std::{collections::HashSet, fmt, rc::Rc};
 
 use zksync_types::{writes::StateDiffRecord, StorageKey, Transaction, H256, U256};
+use zksync_vm2::interface::Tracer;
+use zksync_vm_interface::{
+    utils::{CheckDivergence, DivergenceErrors},
+    Call,
+};
 
 use super::ShadowedFastVm;
 use crate::{
     interface::{
         pubdata::{PubdataBuilder, PubdataInput},
+        storage::InMemoryStorage,
         utils::{ShadowMut, ShadowRef},
         CurrentExecutionState, L2BlockEnv, VmExecutionResultAndLogs,
     },
-    versions::testonly::TestedVm,
+    versions::testonly::{TestedVm, TestedVmWithCallTracer},
+    vm_fast,
 };
 
-impl TestedVm for ShadowedFastVm {
+impl<Tr> TestedVm for ShadowedFastVm<InMemoryStorage, Tr>
+where
+    Tr: Tracer + Default + fmt::Debug + 'static,
+{
     type StateDump = ();
 
     fn dump_state(&self) -> Self::StateDump {
@@ -135,6 +145,44 @@ impl TestedVm for ShadowedFastVm {
     }
 }
 
+#[derive(Debug)]
+struct ExecutionResultAndTraces {
+    result: VmExecutionResultAndLogs,
+    traces: Vec<Call>,
+}
+
+impl From<(VmExecutionResultAndLogs, Vec<Call>)> for ExecutionResultAndTraces {
+    fn from((result, traces): (VmExecutionResultAndLogs, Vec<Call>)) -> Self {
+        Self { result, traces }
+    }
+}
+
+impl From<ExecutionResultAndTraces> for (VmExecutionResultAndLogs, Vec<Call>) {
+    fn from(value: ExecutionResultAndTraces) -> Self {
+        (value.result, value.traces)
+    }
+}
+
+impl CheckDivergence for ExecutionResultAndTraces {
+    fn check_divergence(&self, other: &Self) -> DivergenceErrors {
+        let mut errors = self.result.check_divergence(&other.result);
+        errors.extend(self.traces.check_divergence(&other.traces));
+        errors
+    }
+}
+
+impl TestedVmWithCallTracer for ShadowedFastVm<InMemoryStorage, vm_fast::CallTracer> {
+    fn inspect_with_call_tracer(&mut self) -> (VmExecutionResultAndLogs, Vec<Call>) {
+        self.get_custom_mut("inspect_with_call_tracer", |r| {
+            ExecutionResultAndTraces::from(match r {
+                ShadowMut::Main(vm) => vm.inspect_with_call_tracer(),
+                ShadowMut::Shadow(vm) => vm.inspect_with_call_tracer(),
+            })
+        })
+        .into()
+    }
+}
+
 mod block_tip {
     use crate::versions::testonly::block_tip::*;
 
@@ -167,6 +215,40 @@ mod bytecode_publishing {
     }
 }
 
+mod call_tracer {
+    use crate::versions::testonly::call_tracer::*;
+
+    #[test]
+    fn basic_behavior() {
+        test_basic_behavior::<super::ShadowedFastVm<_, _>>();
+    }
+
+    #[test]
+    fn transfer() {
+        test_transfer::<super::ShadowedFastVm<_, _>>();
+    }
+
+    #[test]
+    fn reverted_tx() {
+        test_reverted_tx::<super::ShadowedFastVm<_, _>>();
+    }
+
+    #[test]
+    fn reverted_deployment() {
+        test_reverted_deployment_tx::<super::ShadowedFastVm<_, _>>();
+    }
+
+    #[test]
+    fn out_of_gas() {
+        test_out_of_gas::<super::ShadowedFastVm<_, _>>();
+    }
+
+    #[test]
+    fn recursive_tx() {
+        test_recursive_tx::<super::ShadowedFastVm<_, _>>();
+    }
+}
+
 mod circuits {
     use crate::versions::testonly::circuits::*;