Implement tracers on vm run method

lambdaclass · Aug 22, 2024 · ff82b18 · ff82b18
1 parent f5a4b3f
commit ff82b18
Show file tree

Hide file tree

Showing 9 changed files with 89 additions and 34 deletions.
diff --git a/core/lib/multivm/src/versions/era_vm/mod.rs b/core/lib/multivm/src/versions/era_vm/mod.rs
@@ -8,5 +8,6 @@ mod refunds;
 mod snapshot;
 #[cfg(test)]
 mod tests;
+mod tracers;
 mod transaction_data;
 pub mod vm;
diff --git a/core/lib/multivm/src/versions/era_vm/tracers/circuits_tracer.rs b/core/lib/multivm/src/versions/era_vm/tracers/circuits_tracer.rs
@@ -2,6 +2,12 @@ use super::traits::{BootloaderTracer, Tracer};
 
 pub struct CircuitsTracer {}
 
+impl CircuitsTracer {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 impl Tracer for CircuitsTracer {
     fn before_decoding(
         &mut self,

diff --git a/core/lib/multivm/src/versions/era_vm/tracers/dispatcher.rs b/core/lib/multivm/src/versions/era_vm/tracers/dispatcher.rs
@@ -1,12 +1,19 @@
-use super::traits::{BootloaderTracer, VmTracer};
 use era_vm::{state::VMState, tracers::tracer::Tracer, Execution, Opcode};
 
+use super::traits::{BootloaderTracer, VmTracer};
+
 #[derive(Default)]
 // dispatcher calls to other tracers
 pub struct TracerDispatcher {
     tracers: Vec<Box<dyn VmTracer>>,
 }
 
+impl TracerDispatcher {
+    pub fn new(tracers: Vec<Box<dyn VmTracer>>) -> Self {
+        Self { tracers }
+    }
+}
+
 impl Tracer for TracerDispatcher {
     fn before_decoding(&mut self, execution: &mut Execution, state: &mut VMState) {
         for tracer in self.tracers.iter_mut() {

diff --git a/core/lib/multivm/src/versions/era_vm/tracers/manager.rs b/core/lib/multivm/src/versions/era_vm/tracers/manager.rs
@@ -2,24 +2,37 @@ use era_vm::{execution::Execution, opcode::Opcode, state::VMState, tracers::trac
 use zksync_state::{ReadStorage, StoragePtr};
 
 use super::{
-    circuits_tracer::CircuitsTracer,
-    dispatcher::TracerDispatcher,
-    pubdata_tracer::PubdataTracer,
-    refunds_tracer::RefundsTracer,
-    result_tracer::ResultTracer,
-    traits::{BootloaderTracer, VmTracer},
+    circuits_tracer::CircuitsTracer, dispatcher::TracerDispatcher, pubdata_tracer::PubdataTracer,
+    refunds_tracer::RefundsTracer, result_tracer::ResultTracer, traits::BootloaderTracer,
 };
 
 // this tracer manager is the one that gets called when running the vm
 pub struct VmTracerManager<S: ReadStorage> {
     dispatcher: TracerDispatcher,
     result_tracer: ResultTracer,
     refund_tracer: Option<RefundsTracer>,
-    pubdata_tracer: Option<PubdataTracer>,
+    pubdata_tracer: PubdataTracer,
     circuits_tracer: CircuitsTracer,
     storage: StoragePtr<S>,
 }
 
+impl<S: ReadStorage> VmTracerManager<S> {
+    pub fn new(
+        storage: StoragePtr<S>,
+        dispatcher: TracerDispatcher,
+        refund_tracer: Option<RefundsTracer>,
+    ) -> Self {
+        Self {
+            dispatcher,
+            refund_tracer,
+            circuits_tracer: CircuitsTracer::new(),
+            result_tracer: ResultTracer::new(),
+            pubdata_tracer: PubdataTracer::new(),
+            storage,
+        }
+    }
+}
+
 impl<S: ReadStorage> Tracer for VmTracerManager<S> {
     fn before_decoding(&mut self, execution: &mut Execution, state: &mut VMState) {
         // Call the dispatcher to handle all the tracers added to it
@@ -30,9 +43,7 @@ impl<S: ReadStorage> Tracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.before_decoding(execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.before_decoding(execution, state);
-        }
+        self.pubdata_tracer.before_decoding(execution, state);
         self.circuits_tracer.before_decoding(execution, state);
     }
 
@@ -45,9 +56,7 @@ impl<S: ReadStorage> Tracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.after_decoding(opcode, execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.after_decoding(opcode, execution, state);
-        }
+        self.pubdata_tracer.after_decoding(opcode, execution, state);
         self.circuits_tracer
             .after_decoding(opcode, execution, state);
     }
@@ -67,9 +76,8 @@ impl<S: ReadStorage> Tracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.before_execution(opcode, execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.before_execution(opcode, execution, state);
-        }
+        self.pubdata_tracer
+            .before_execution(opcode, execution, state);
         self.circuits_tracer
             .before_execution(opcode, execution, state);
     }
@@ -83,9 +91,8 @@ impl<S: ReadStorage> Tracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.after_execution(opcode, execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.after_execution(opcode, execution, state);
-        }
+        self.pubdata_tracer
+            .after_execution(opcode, execution, state);
         self.circuits_tracer
             .after_execution(opcode, execution, state);
     }
@@ -108,9 +115,8 @@ impl<S: ReadStorage> BootloaderTracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.before_bootloader_execution(opcode, execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.before_bootloader_execution(opcode, execution, state);
-        }
+        self.pubdata_tracer
+            .before_bootloader_execution(opcode, execution, state);
         self.circuits_tracer
             .before_bootloader_execution(opcode, execution, state);
     }
@@ -131,9 +137,8 @@ impl<S: ReadStorage> BootloaderTracer for VmTracerManager<S> {
         if let Some(refunds_tracer) = &mut self.refund_tracer {
             refunds_tracer.after_bootloader_execution(opcode, execution, state);
         }
-        if let Some(pubdata_tracer) = &mut self.pubdata_tracer {
-            pubdata_tracer.after_bootloader_execution(opcode, execution, state);
-        }
+        self.pubdata_tracer
+            .after_bootloader_execution(opcode, execution, state);
         self.circuits_tracer
             .after_bootloader_execution(opcode, execution, state);
     }

diff --git a/core/lib/multivm/src/versions/era_vm/tracers/mod.rs b/core/lib/multivm/src/versions/era_vm/tracers/mod.rs
@@ -1,7 +1,7 @@
-mod circuits_tracer;
-mod dispatcher;
+pub mod circuits_tracer;
+pub mod dispatcher;
 pub mod manager;
-mod pubdata_tracer;
-mod refunds_tracer;
-mod result_tracer;
+pub mod pubdata_tracer;
+pub mod refunds_tracer;
+pub mod result_tracer;
 mod traits;
diff --git a/core/lib/multivm/src/versions/era_vm/tracers/pubdata_tracer.rs b/core/lib/multivm/src/versions/era_vm/tracers/pubdata_tracer.rs
@@ -2,6 +2,12 @@ use super::traits::{BootloaderTracer, Tracer, VmTracer};
 
 pub struct PubdataTracer {}
 
+impl PubdataTracer {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 impl Tracer for PubdataTracer {
     fn before_decoding(
         &mut self,

diff --git a/core/lib/multivm/src/versions/era_vm/tracers/refunds_tracer.rs b/core/lib/multivm/src/versions/era_vm/tracers/refunds_tracer.rs
@@ -1,7 +1,13 @@
-use super::traits::{BootloaderTracer, Tracer, VmTracer};
+use super::traits::{BootloaderTracer, Tracer};
 
 pub struct RefundsTracer {}
 
+impl RefundsTracer {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 impl Tracer for RefundsTracer {
     fn before_decoding(
         &mut self,

diff --git a/core/lib/multivm/src/versions/era_vm/tracers/result_tracer.rs b/core/lib/multivm/src/versions/era_vm/tracers/result_tracer.rs
@@ -2,6 +2,12 @@ use super::traits::{BootloaderTracer, Tracer, VmTracer};
 
 pub struct ResultTracer {}
 
+impl ResultTracer {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
 impl Tracer for ResultTracer {
     fn before_decoding(
         &mut self,

diff --git a/core/lib/multivm/src/versions/era_vm/vm.rs b/core/lib/multivm/src/versions/era_vm/vm.rs
@@ -38,6 +38,9 @@ use super::{
     logs::IntoSystemLog,
     refunds::compute_refund,
     snapshot::VmSnapshot,
+    tracers::{
+        dispatcher::TracerDispatcher, manager::VmTracerManager, refunds_tracer::RefundsTracer,
+    },
 };
 use crate::{
     era_vm::{bytecode::compress_bytecodes, transaction_data::TransactionData},
@@ -157,6 +160,7 @@ impl<S: ReadStorage + 'static> Vm<S> {
     pub fn run(
         &mut self,
         execution_mode: VmExecutionMode,
+        tracer: TracerDispatcher,
         track_refunds: bool,
     ) -> (ExecutionResult, Refunds) {
         let mut refunds = Refunds {
@@ -166,8 +170,18 @@ impl<S: ReadStorage + 'static> Vm<S> {
         let mut pubdata_before = self.inner.state.pubdata() as u32;
         let mut last_tx_result = None;
 
+        let refund_tracer = if track_refunds {
+            Some(RefundsTracer::new())
+        } else {
+            None
+        };
+
+        let mut tracer = VmTracerManager::new(self.storage.clone(), tracer, refund_tracer);
+
         loop {
-            let (result, _blob_tracer) = self.inner.run_program_with_custom_bytecode();
+            let result = self
+                .inner
+                .run_program_with_custom_bytecode(Some(&mut tracer));
 
             let result = match result {
                 ExecutionOutput::Ok(output) => {
@@ -515,7 +529,11 @@ impl<S: ReadStorage + 'static> VmInterface for Vm<S> {
         }
 
         let snapshot = self.inner.state.snapshot();
-        let (result, refunds) = self.run(execution_mode, enable_refund_tracer);
+        let (result, refunds) = self.run(
+            execution_mode,
+            TracerDispatcher::new(vec![]),
+            enable_refund_tracer,
+        );
 
         let ignore_world_diff = matches!(execution_mode, VmExecutionMode::OneTx)
             && matches!(result, ExecutionResult::Halt { .. });