fix: restore instruction count functionality

core/tests/vm-benchmark/Cargo.toml

@@ -11,6 +11,7 @@ zksync_multivm.workspace = true
 zksync_types.workspace = true
 zksync_utils.workspace = true
 zksync_vlog.workspace = true
+zksync_vm2.workspace = true
 
 criterion.workspace = true
 once_cell.workspace = true

core/tests/vm-benchmark/src/bin/instruction_counts.rs

@@ -1,11 +1,16 @@
 //! Runs all benchmarks and prints out the number of zkEVM opcodes each one executed.
 
-use vm_benchmark::{BenchmarkingVm, BYTECODES};
+use vm_benchmark::{BenchmarkingVmFactory, Fast, Legacy, BYTECODES};
 
 fn main() {
     for bytecode in BYTECODES {
         let tx = bytecode.deploy_tx();
         let name = bytecode.name;
-        println!("{name} {}", BenchmarkingVm::new().instruction_count(&tx));
+        println!("{name} {}", Fast::<()>::count_instructions(&tx));
+        println!(
+            "{} {}",
+            name.to_string() + "_legacy",
+            Legacy::count_instructions(&tx)
+        );
     }
 }

core/tests/vm-benchmark/src/instruction_counter.rs

@@ -13,7 +13,6 @@ pub struct InstructionCounter {
 
 /// A tracer that counts the number of instructions executed by the VM.
 impl InstructionCounter {
-    #[allow(dead_code)] // FIXME: re-enable instruction counting once new tracers are merged
     pub fn new(output: Rc<RefCell<usize>>) -> Self {
         Self { count: 0, output }
     }

core/tests/vm-benchmark/src/vm.rs

@@ -9,8 +9,8 @@ use zksync_multivm::{
         VmExecutionResultAndLogs, VmFactory, VmInterface, VmInterfaceExt,
         VmInterfaceHistoryEnabled,
     },
-    vm_fast, vm_latest,
-    vm_latest::{constants::BATCH_COMPUTATIONAL_GAS_LIMIT, HistoryEnabled},
+    vm_fast,
+    vm_latest::{self, constants::BATCH_COMPUTATIONAL_GAS_LIMIT, HistoryEnabled, ToTracerPointer},
     zk_evm_latest::ethereum_types::{Address, U256},
 };
 use zksync_types::{
@@ -20,7 +20,7 @@ use zksync_types::{
 };
 use zksync_utils::bytecode::hash_bytecode;
 
-use crate::transaction::PRIVATE_KEY;
+use crate::{instruction_counter::InstructionCounter, transaction::PRIVATE_KEY};
 
 static SYSTEM_CONTRACTS: Lazy<BaseSystemContracts> = Lazy::new(BaseSystemContracts::load_from_disk);
 
@@ -72,16 +72,19 @@ pub trait BenchmarkingVmFactory {
         system_env: SystemEnv,
         storage: &'static InMemoryStorage,
     ) -> Self::Instance;
+
+    /// Counts instructions executed by the VM while processing the transaction.
+    fn count_instructions(tx: &Transaction) -> usize;
 }
 
 /// Factory for the new / fast VM.
 #[derive(Debug)]
-pub struct Fast(());
+pub struct Fast<Tr = ()>(Tr);
 
-impl BenchmarkingVmFactory for Fast {
+impl<Tr: vm_fast::Tracer + Default + 'static> BenchmarkingVmFactory for Fast<Tr> {
     const LABEL: VmLabel = VmLabel::Fast;
 
-    type Instance = vm_fast::Vm<&'static InMemoryStorage>;
+    type Instance = vm_fast::Vm<&'static InMemoryStorage, Tr>;
 
     fn create(
         batch_env: L1BatchEnv,
@@ -90,6 +93,29 @@ impl BenchmarkingVmFactory for Fast {
     ) -> Self::Instance {
         vm_fast::Vm::custom(batch_env, system_env, storage)
     }
+
+    fn count_instructions(tx: &Transaction) -> usize {
+        let mut vm = BenchmarkingVm::<Fast<InstructionCount>>::default();
+        vm.0.push_transaction(tx.clone());
+
+        #[derive(Default)]
+        struct InstructionCount(usize);
+        impl vm_fast::Tracer for InstructionCount {
+            fn before_instruction<
+                OP: zksync_vm2::interface::OpcodeType,
+                S: zksync_vm2::interface::StateInterface,
+            >(
+                &mut self,
+                _: &mut S,
+            ) {
+                self.0 += 1;
+            }
+        }
+        let mut tracer = InstructionCount(0);
+
+        vm.0.inspect(&mut tracer, VmExecutionMode::OneTx);
+        tracer.0
+    }
 }
 
 /// Factory for the legacy VM (latest version).
@@ -109,6 +135,19 @@ impl BenchmarkingVmFactory for Legacy {
         let storage = StorageView::new(storage).to_rc_ptr();
         vm_latest::Vm::new(batch_env, system_env, storage)
     }
+
+    fn count_instructions(tx: &Transaction) -> usize {
+        let mut vm = BenchmarkingVm::<Self>::default();
+        vm.0.push_transaction(tx.clone());
+        let count = Rc::new(RefCell::new(0));
+        vm.0.inspect(
+            &mut InstructionCounter::new(count.clone())
+                .into_tracer_pointer()
+                .into(),
+            VmExecutionMode::OneTx,
+        );
+        count.take()
+    }
 }
 
 #[derive(Debug)]
@@ -169,13 +208,6 @@ impl<VM: BenchmarkingVmFactory> BenchmarkingVm<VM> {
         }
         tx_result
     }
-
-    pub fn instruction_count(&mut self, tx: &Transaction) -> usize {
-        self.0.push_transaction(tx.clone());
-        let count = Rc::new(RefCell::new(0));
-        self.0.execute(VmExecutionMode::OneTx); // FIXME: re-enable instruction counting once new tracers are merged
-        count.take()
-    }
 }
 
 impl BenchmarkingVm<Fast> {