feat: Introduced Initializers to trace inspector

crates/evm/src/evm.rs

@@ -2,7 +2,7 @@
 
 use crate::{tracing::TxTracer, EvmEnv, EvmError, IntoTxEnv};
 use alloy_primitives::{Address, Bytes};
-use core::{error::Error, fmt::Debug, hash::Hash};
+use core::{convert::Infallible, error::Error, fmt::Debug, hash::Hash};
 use revm::{
     context::{result::ExecutionResult, BlockEnv},
     context_interface::{
@@ -244,7 +244,7 @@ pub trait EvmFactoryExt: EvmFactory {
         db: DB,
         input: EvmEnv<Self::Spec>,
         fused_inspector: I,
-    ) -> TxTracer<Self::Evm<DB, I>>
+    ) -> Result<TxTracer<Self::Evm<DB, I>>, Infallible>
     where
         DB: Database + DatabaseCommit,
         I: Inspector<Self::Context<DB>> + Clone,

crates/evm/src/tracing.rs

@@ -1,23 +1,68 @@
 //! Helpers for tracing.
 
 use crate::{Evm, IntoTxEnv};
-use core::{fmt::Debug, iter::Peekable};
+use core::{convert::Infallible, fmt::Debug, iter::Peekable};
 use revm::{
     context::result::{ExecutionResult, ResultAndState},
     state::EvmState,
     DatabaseCommit,
 };
 
+/// Trait for inspector initialization.
+pub trait InspectorInitializer<T> {
+    /// Associated inspector initialization error type.
+    type Error;
+
+    /// Create a new inspector instance.
+    fn initialize(&mut self) -> Result<T, Self::Error>;
+}
+
+/// A simple initializer that clones the inspector.
+#[derive(Debug, Clone)]
+pub struct CloneInitializer<T: Clone>(T);
+
+impl<T: Clone> CloneInitializer<T> {
+    /// Create a new clone initializer.
+    pub fn new(value: T) -> Self {
+        Self(value)
+    }
+}
+
+impl<T: Clone> InspectorInitializer<T> for CloneInitializer<T> {
+    type Error = Infallible;
+
+    fn initialize(&mut self) -> Result<T, Self::Error> {
+        Ok(self.0.clone())
+    }
+}
+
+/// Implementation for closures.
+impl<T, E, F: FnMut() -> Result<T, E>> InspectorInitializer<T> for F {
+    type Error = E;
+
+    fn initialize(&mut self) -> Result<T, E> {
+        self()
+    }
+}
+
 /// A helper type for tracing transactions.
 #[derive(Debug, Clone)]
-pub struct TxTracer<E: Evm> {
+pub struct TxTracer<
+    E: Evm,
+    Init: InspectorInitializer<E::Inspector> = CloneInitializer<<E as Evm>::Inspector>,
+> {
     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> = CloneInitializer<<E as Evm>::Inspector>,
+> {
     /// The transaction that was just executed.
     pub tx: T,
     /// Result of transaction execution.
@@ -28,37 +73,47 @@ pub struct TracingCtx<'a, T, E: Evm> {
     pub inspector: &'a mut E::Inspector,
     /// Database used when executing the transaction, _before_ committing the state changes.
     pub db: &'a mut E::DB,
-    /// Fused inspector.
-    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>> 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, Init::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>> TxTracer<E, Init> {
+    /// Creates a new [`TxTracer`] instance with a custom inspector initializer.
+    pub fn new_with_initializer(evm: E, initializer: Init) -> Result<Self, Init::Error> {
+        Ok(Self { 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, Init::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>, Init::Error>
+    where
+        <Init as InspectorInitializer<E::Inspector>>::Error: From<E::Error>,
+    {
         let result = self.evm.transact_commit(tx);
-        let inspector = self.fuse_inspector();
+        let inspector = self.fuse_inspector()?;
         Ok(TraceOutput { result: result?, inspector })
     }
 
@@ -69,11 +124,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 +144,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 +161,17 @@ impl<E: Evm<Inspector: Clone, DB: DatabaseCommit>> TxTracer<E> {
     }
 }
 
+impl<E: Evm<DB: DatabaseCommit>> TxTracer<E, CloneInitializer<E::Inspector>>
+where
+    E::Inspector: Clone,
+{
+    /// Creates a new [`TxTracer`] instance with a cloneable inspector.
+    pub fn new(mut evm: E) -> Result<Self, Infallible> {
+        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 +184,17 @@ 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>, 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>, 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 +211,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>,
     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, initializer } = self.inner;
         let (db, inspector, _) = evm.components_mut();
 
         let Ok(ResultAndState { result, state }) = result else {
@@ -164,8 +239,8 @@ where
             state: &state,
             inspector,
             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
@@ -175,7 +250,7 @@ where
         }
 
         if self.fuse && !was_fused {
-            self.inner.fuse_inspector();
+            let _ = self.inner.fuse_inspector();
         }
 
         Some(output)