feat: Introduced Initializers to trace inspector

crates/evm/src/tracing.rs

@@ -1,23 +1,78 @@
 //! Helpers for tracing.
 
 use crate::{Evm, IntoTxEnv};
-use core::{fmt::Debug, iter::Peekable};
+use core::{fmt::Debug, iter::Peekable, marker};
 use revm::{
     context::result::{ExecutionResult, ResultAndState},
     state::EvmState,
     DatabaseCommit,
 };
 
+/// Trait for inspector initialization.
+pub trait InspectorInitializer<T, E> {
+    /// Create a new inspector instance.
+    fn initialize(&mut self) -> Result<T, E>;
+}
+
+/// A simple initializer that clones the inspector.
+#[derive(Debug, Clone)]
+pub struct CloneInitializer<T: Clone, E>(T, marker::PhantomData<E>);
+
+impl<T: Clone, E> CloneInitializer<T, E> {
+    /// Create a new clone initializer.
+    pub fn new(value: T) -> Self {
+        Self(value, marker::PhantomData)
+    }
+}
+
+impl<T: Clone, E> InspectorInitializer<T, E> for CloneInitializer<T, E> {
+    fn initialize(&mut self) -> Result<T, E> {
+        Ok(self.0.clone())
+    }
+}
+
+/// An initializer that creates inspectors using a closure.
+#[derive(Debug)]
+pub struct FnInitializer<T, E, F: FnMut() -> Result<T, E>>(F);
+
+impl<T, E, F: FnMut() -> Result<T, E>> FnInitializer<T, E, F> {
+    /// Create a new function initializer.
+    pub fn new(f: F) -> Self {
+        Self(f)
+    }
+}
+
+impl<T, E, F: FnMut() -> Result<T, E>> InspectorInitializer<T, E> for FnInitializer<T, E, F> {
+    fn initialize(&mut self) -> Result<T, E> {
+        (self.0)()
+    }
+}
+
 /// A helper type for tracing transactions.
 #[derive(Debug, Clone)]
-pub struct TxTracer<E: Evm> {
+pub struct TxTracer<
+    E: Evm,
+    Init: InspectorInitializer<E::Inspector, E::Error> = CloneInitializer<
+        <E as Evm>::Inspector,
+        <E as Evm>::Error,
+    >,
+> {
     evm: E,
     fused_inspector: E::Inspector,
+    initializer: Init,
 }
 
 /// Container type for context exposed in [`TxTracer`].
 #[derive(Debug)]
-pub struct TracingCtx<'a, T, E: Evm> {
+pub struct TracingCtx<
+    'a,
+    T,
+    E: Evm,
+    Init: InspectorInitializer<E::Inspector, E::Error> = CloneInitializer<
+        <E as Evm>::Inspector,
+        <E as Evm>::Error,
+    >,
+> {
     /// The transaction that was just executed.
     pub tx: T,
     /// Result of transaction execution.
@@ -32,33 +87,48 @@ pub struct TracingCtx<'a, T, E: Evm> {
     fused_inspector: &'a E::Inspector,
     /// Whether the inspector was fused.
     was_fused: &'a mut bool,
+    /// Reference to the initializer.
+    initializer: &'a mut Init,
 }
 
-impl<'a, T, E: Evm<Inspector: Clone>> TracingCtx<'a, T, E> {
-    /// Fuses the inspector and returns the current inspector state.
-    pub fn take_inspector(&mut self) -> E::Inspector {
+impl<'a, T, E: Evm, Init: InspectorInitializer<E::Inspector, E::Error>> TracingCtx<'a, T, E, Init> {
+    /// Takes the current inspector and replaces it with a fresh one.
+    ///
+    /// This is useful when you want to keep the current inspector state but continue
+    /// tracing with a new one.
+    pub fn take_inspector(&mut self) -> Result<E::Inspector, E::Error>
+    where
+        E::Inspector: Clone,
+    {
         *self.was_fused = true;
-        core::mem::replace(self.inspector, self.fused_inspector.clone())
+        Ok(core::mem::replace(self.inspector, self.initializer.initialize()?))
     }
 }
 
-impl<E: Evm<Inspector: Clone, DB: DatabaseCommit>> TxTracer<E> {
-    /// Creates a new [`TxTracer`] instance.
-    pub fn new(mut evm: E) -> Self {
-        Self { fused_inspector: evm.inspector_mut().clone(), evm }
+impl<E: Evm<DB: DatabaseCommit>, Init: InspectorInitializer<E::Inspector, E::Error>>
+    TxTracer<E, Init>
+{
+    /// Creates a new [`TxTracer`] instance with a custom inspector initializer.
+    pub fn new_with_initializer(evm: E, mut initializer: Init) -> Result<Self, E::Error> {
+        let fused_inspector = initializer.initialize()?;
+        Ok(Self { fused_inspector, evm, initializer })
     }
 
-    fn fuse_inspector(&mut self) -> E::Inspector {
-        core::mem::replace(self.evm.inspector_mut(), self.fused_inspector.clone())
+    fn fuse_inspector(&mut self) -> Result<E::Inspector, E::Error> {
+        let fresh = self.initializer.initialize()?;
+        Ok(core::mem::replace(self.evm.inspector_mut(), fresh))
     }
 
     /// Executes a transaction, and returns its outcome along with the inspector state.
     pub fn trace(
         &mut self,
         tx: impl IntoTxEnv<E::Tx>,
-    ) -> Result<TraceOutput<E::HaltReason, E::Inspector>, E::Error> {
+    ) -> Result<TraceOutput<E::HaltReason, E::Inspector>, E::Error>
+    where
+        E::Inspector: Clone,
+    {
         let result = self.evm.transact_commit(tx);
-        let inspector = self.fuse_inspector();
+        let inspector = self.fuse_inspector()?;
         Ok(TraceOutput { result: result?, inspector })
     }
 
@@ -69,11 +139,17 @@ impl<E: Evm<Inspector: Clone, DB: DatabaseCommit>> TxTracer<E> {
         &mut self,
         txs: Txs,
         mut f: F,
-    ) -> TracerIter<'_, E, Txs::IntoIter, impl FnMut(TracingCtx<'_, T, E>) -> Result<O, E::Error>>
+    ) -> TracerIter<
+        '_,
+        E,
+        Init,
+        Txs::IntoIter,
+        impl FnMut(TracingCtx<'_, T, E, Init>) -> Result<O, E::Error>,
+    >
     where
         T: IntoTxEnv<E::Tx> + Clone,
         Txs: IntoIterator<Item = T>,
-        F: FnMut(TracingCtx<'_, Txs::Item, E>) -> O,
+        F: FnMut(TracingCtx<'_, Txs::Item, E, Init>) -> O,
     {
         self.try_trace_many(txs, move |ctx| Ok(f(ctx)))
     }
@@ -83,11 +159,11 @@ impl<E: Evm<Inspector: Clone, DB: DatabaseCommit>> TxTracer<E> {
         &mut self,
         txs: Txs,
         hook: F,
-    ) -> TracerIter<'_, E, Txs::IntoIter, F>
+    ) -> TracerIter<'_, E, Init, Txs::IntoIter, F>
     where
         T: IntoTxEnv<E::Tx> + Clone,
         Txs: IntoIterator<Item = T>,
-        F: FnMut(TracingCtx<'_, T, E>) -> Result<O, Err>,
+        F: FnMut(TracingCtx<'_, T, E, Init>) -> Result<O, Err>,
         Err: From<E::Error>,
     {
         TracerIter {
@@ -100,6 +176,17 @@ impl<E: Evm<Inspector: Clone, DB: DatabaseCommit>> TxTracer<E> {
     }
 }
 
+impl<E: Evm<DB: DatabaseCommit>> TxTracer<E, CloneInitializer<E::Inspector, E::Error>>
+where
+    E::Inspector: Clone,
+{
+    /// Creates a new [`TxTracer`] instance with a cloneable inspector.
+    pub fn new(mut evm: E) -> Result<Self, E::Error> {
+        let inspector = evm.inspector_mut().clone();
+        Self::new_with_initializer(evm, CloneInitializer::new(inspector))
+    }
+}
+
 /// Output of tracing a transaction.
 #[derive(Debug, Clone)]
 pub struct TraceOutput<H, I> {
@@ -112,15 +199,23 @@ pub struct TraceOutput<H, I> {
 /// Iterator used by tracer.
 #[derive(derive_more::Debug)]
 #[debug(bound(E::Inspector: Debug))]
-pub struct TracerIter<'a, E: Evm, Txs: Iterator, F> {
-    inner: &'a mut TxTracer<E>,
+pub struct TracerIter<
+    'a,
+    E: Evm,
+    Init: InspectorInitializer<E::Inspector, E::Error>,
+    Txs: Iterator,
+    F,
+> {
+    inner: &'a mut TxTracer<E, Init>,
     txs: Peekable<Txs>,
     hook: F,
     skip_last_commit: bool,
     fuse: bool,
 }
 
-impl<E: Evm, Txs: Iterator, F> TracerIter<'_, E, Txs, F> {
+impl<E: Evm, Init: InspectorInitializer<E::Inspector, E::Error>, Txs: Iterator, F>
+    TracerIter<'_, E, Init, Txs, F>
+{
     /// Flips the `skip_last_commit` flag thus making sure all transaction are committed.
     ///
     /// We are skipping last commit by default as it's expected that when tracing users are mostly
@@ -137,21 +232,22 @@ impl<E: Evm, Txs: Iterator, F> TracerIter<'_, E, Txs, F> {
     }
 }
 
-impl<E, T, Txs, F, O, Err> Iterator for TracerIter<'_, E, Txs, F>
+impl<E, Init, T, Txs, F, O, Err> Iterator for TracerIter<'_, E, Init, Txs, F>
 where
-    E: Evm<DB: DatabaseCommit, Inspector: Clone>,
+    E: Evm<DB: DatabaseCommit>,
+    Init: InspectorInitializer<E::Inspector, E::Error>,
     T: IntoTxEnv<E::Tx> + Clone,
     Txs: Iterator<Item = T>,
     Err: From<E::Error>,
-    F: FnMut(TracingCtx<'_, T, E>) -> Result<O, Err>,
+    F: FnMut(TracingCtx<'_, T, E, Init>) -> Result<O, Err>,
 {
     type Item = Result<O, Err>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let tx = self.txs.next()?;
         let result = self.inner.evm.transact(tx.clone());
 
-        let TxTracer { evm, fused_inspector } = self.inner;
+        let TxTracer { evm, fused_inspector, initializer } = self.inner;
         let (db, inspector, _) = evm.components_mut();
 
         let Ok(ResultAndState { result, state }) = result else {
@@ -166,6 +262,7 @@ where
             db,
             fused_inspector: &*fused_inspector,
             was_fused: &mut was_fused,
+            initializer,
         });
 
         // Only commit next transaction if `skip_last_commit` is disabled or there is a next