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/benches/iai.rs

@@ -31,4 +31,5 @@ make_functions_and_main!(
     write_and_decode => write_and_decode_legacy,
     event_spam => event_spam_legacy,
     slot_hash_collision => slot_hash_collision_legacy,
+    heap_read_write => heap_read_write_legacy,
 );

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

@@ -25,14 +25,7 @@ fn main() {
         .keys()
         .collect::<HashSet<_>>()
         .intersection(&iai_after.keys().collect())
-        .filter_map(|&name| {
-            let diff = percent_difference(iai_before[name], iai_after[name]);
-            if diff.abs() > 2. {
-                Some((name, format!("{:+.1}%", diff)))
-            } else {
-                None
-            }
-        })
+        .map(|&name| (name, percent_difference(iai_before[name], iai_after[name])))
         .collect::<HashMap<_, _>>();
 
     let duration_changes = opcodes_before
@@ -47,12 +40,17 @@ fn main() {
 
     let mut nonzero_diff = false;
 
-    for name in perf_changes.keys().collect::<HashSet<_>>().union(
-        &duration_changes
-            .iter()
-            .filter_map(|(key, value)| (*value != 0).then_some(key))
-            .collect(),
-    ) {
+    for name in perf_changes
+        .iter()
+        .filter_map(|(key, value)| (value.abs() > 2.).then_some(key))
+        .collect::<HashSet<_>>()
+        .union(
+            &duration_changes
+                .iter()
+                .filter_map(|(key, value)| (*value != 0).then_some(key))
+                .collect(),
+        )
+    {
         // write the header before writing the first line of diff
         if !nonzero_diff {
             println!("Benchmark name | change in estimated runtime | change in number of opcodes executed \n--- | --- | ---");
@@ -63,7 +61,10 @@ fn main() {
         println!(
             "{} | {} | {}",
             name,
-            perf_changes.get(**name).unwrap_or(&n_a.clone()),
+            perf_changes
+                .get(**name)
+                .map(|percent| format!("{:+.1}%", percent))
+                .unwrap_or(n_a.clone()),
             duration_changes
                 .get(**name)
                 .map(|abs_diff| format!(

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> {