breaking up simulator code

Cargo.toml

@@ -24,6 +24,7 @@ path = "bin/submit_transaction.rs"
 [dependencies]
 init4-bin-base = { git = "https://github.com/init4tech/bin-base.git" }
 
+signet-evm = { git = "https://github.com/init4tech/signet-sdk", branch = "main" }
 signet-bundle = { git = "https://github.com/init4tech/signet-sdk", branch = "main" }
 signet-types = { git = "https://github.com/init4tech/signet-sdk", branch = "main" }
 signet-zenith = { git = "https://github.com/init4tech/signet-sdk", branch = "main" }
@@ -60,3 +61,4 @@ tokio = { version = "1.36.0", features = ["full", "macros", "rt-multi-thread"] }
 
 async-trait = "0.1.80"
 oauth2 = "4.4.2"
+futures-util = "0.3.31"

src/tasks/mod.rs

@@ -16,5 +16,5 @@ pub mod submit;
 /// Tx polling task
 pub mod tx_poller;
 
-/// Tx and bundle simulation task
-pub mod simulator;
+/// Simulation logic
+pub mod sim;

src/tasks/sim/env.rs

@@ -0,0 +1,223 @@
+use crate::tasks::sim::SimOutcome;
+use signet_bundle::{SignetEthBundle, SignetEthBundleDriver, SignetEthBundleError};
+use signet_evm::SignetLayered;
+use signet_types::config::SignetSystemConstants;
+use std::{convert::Infallible, marker::PhantomData};
+use trevm::{
+    Block, BundleDriver, Cfg, DbConnect, EvmFactory, Tx,
+    db::{TryCachingDb, cow::CacheOnWrite},
+    helpers::Ctx,
+    inspectors::{Layered, TimeLimit},
+    revm::{
+        DatabaseRef, Inspector, context::result::EVMError, database::Cache,
+        inspector::NoOpInspector,
+    },
+};
+
+use super::BundleOrTx;
+
+/// Factory for creating simulation tasks.
+#[derive(Debug, Clone)]
+pub struct SimEnv<Db, C, B, Insp = NoOpInspector> {
+    /// The database to use for the simulation.
+    db: Db,
+
+    /// The system constants for the Signet network.
+    constants: SignetSystemConstants,
+
+    /// Chain cfg to use for the simulation.
+    cfg: C,
+
+    /// Block to use for the simulation.
+    block: B,
+
+    /// The max time to spend on any simulation.
+    execution_timeout: std::time::Duration,
+
+    _pd: PhantomData<fn() -> Insp>,
+}
+
+impl<Db, C, B, Insp> SimEnv<Db, C, B, Insp> {
+    /// Creates a new `SimFactory` instance.
+    pub fn new(
+        db: Db,
+        constants: SignetSystemConstants,
+        cfg: C,
+        block: B,
+        execution_timeout: std::time::Duration,
+    ) -> Self {
+        Self { db, constants, cfg, block, execution_timeout, _pd: PhantomData }
+    }
+
+    /// Get a reference to the database.
+    pub fn db_mut(&mut self) -> &mut Db {
+        &mut self.db
+    }
+
+    /// Get a reference to the system constants.
+    pub fn constants(&self) -> &SignetSystemConstants {
+        &self.constants
+    }
+
+    /// Get a reference to the chain cfg.
+    pub fn cfg(&self) -> &C {
+        &self.cfg
+    }
+
+    /// Get a mutable reference to the chain cfg.
+    pub fn cfg_mut(&mut self) -> &mut C {
+        &mut self.cfg
+    }
+
+    /// Get a reference to the block.
+    pub fn block(&self) -> &B {
+        &self.block
+    }
+
+    /// Get a mutable reference to the block.
+    pub fn block_mut(&mut self) -> &mut B {
+        &mut self.block
+    }
+
+    /// Get the exectuion timeout.
+    pub fn execution_timeout(&self) -> std::time::Duration {
+        self.execution_timeout
+    }
+
+    /// Set the execution timeout.
+    pub fn set_execution_timeout(&mut self, timeout: std::time::Duration) {
+        self.execution_timeout = timeout;
+    }
+}
+
+impl<Db, C, B, Insp> DbConnect for SimEnv<Db, C, B, Insp>
+where
+    Db: DatabaseRef + Clone + Sync,
+    C: Sync,
+    B: Sync,
+    Insp: Sync,
+{
+    type Database = CacheOnWrite<Db>;
+
+    type Error = Infallible;
+
+    fn connect(&self) -> Result<Self::Database, Self::Error> {
+        Ok(CacheOnWrite::new(self.db.clone()))
+    }
+}
+
+impl<Db, C, B, Insp> EvmFactory for SimEnv<Db, C, B, Insp>
+where
+    Db: DatabaseRef + Clone + Sync,
+    C: Sync,
+    B: Sync,
+    Insp: Inspector<Ctx<CacheOnWrite<Db>>> + Default + Sync,
+{
+    type Insp = SignetLayered<Layered<TimeLimit, Insp>>;
+
+    fn create(&self) -> Result<trevm::EvmNeedsCfg<Self::Database, Self::Insp>, Self::Error> {
+        let db = self.connect().unwrap();
+
+        let inspector = Layered::new(TimeLimit::new(self.execution_timeout), Insp::default());
+
+        Ok(signet_evm::signet_evm_with_inspector(db, inspector, self.constants.clone()))
+    }
+}
+
+impl<Db, C, B, Insp> SimEnv<Db, C, B, Insp>
+where
+    Db: DatabaseRef + Clone + Sync,
+    C: Cfg + Sync,
+    B: Block + Sync,
+    Insp: Inspector<Ctx<CacheOnWrite<Db>>> + Default + Sync,
+{
+    /// Simulates a transaction in the context of a block.
+    ///
+    /// This function runs the simulation in a separate thread and waits for
+    /// the result or the deadline to expire.
+    fn simulate_tx<'a, T>(
+        &self,
+        transaction: &'a T,
+    ) -> Result<SimOutcome<&'a T, Cache>, SignetEthBundleError<CacheOnWrite<Db>>>
+    where
+        T: Tx,
+    {
+        let trevm = self.create_with_block(&self.cfg, &self.block).unwrap();
+
+        // Get the initial beneficiary balance
+        let beneificiary = trevm.beneficiary();
+        let initial_beneficiary_balance =
+            trevm.try_read_balance_ref(beneificiary).map_err(EVMError::Database)?;
+
+        // If succesful, take the cache. If failed, return the error.
+        match trevm.run_tx(transaction) {
+            Ok(trevm) => {
+                // Get the beneficiary balance after the transaction and calculate the
+                // increase
+                let beneficiary_balance =
+                    trevm.try_read_balance_ref(beneificiary).map_err(EVMError::Database)?;
+                let increase = beneficiary_balance.saturating_sub(initial_beneficiary_balance);
+
+                let cache = trevm.accept_state().into_db().into_cache();
+
+                // Create the outcome
+                Ok(SimOutcome::new_unchecked(transaction, cache, increase))
+            }
+            Err(e) => Err(SignetEthBundleError::from(e.into_error())),
+        }
+    }
+
+    /// Simulates a bundle in the context of a block.
+    fn simulate_bundle<'a>(
+        &self,
+        bundle: &'a SignetEthBundle,
+    ) -> Result<SimOutcome<&'a SignetEthBundle, Cache>, SignetEthBundleError<CacheOnWrite<Db>>>
+    where
+        Insp: Inspector<Ctx<CacheOnWrite<Db>>> + Default + Sync,
+    {
+        let mut driver =
+            SignetEthBundleDriver::new(&bundle, std::time::Instant::now() + self.execution_timeout);
+        let trevm = self.create_with_block(&self.cfg, &self.block).unwrap();
+
+        // run the bundle
+        let trevm = match driver.run_bundle(trevm) {
+            Ok(result) => result,
+            Err(e) => return Err(e.into_error()),
+        };
+
+        // evaluate the result
+        let score = driver.beneficiary_balance_increase();
+
+        let db = trevm.into_db();
+        let cache = db.into_cache();
+
+        Ok(SimOutcome::new_unchecked(bundle, cache, score))
+    }
+
+    /// Simulate a [`BundleOrTx`], containing either a [`SignetEthBundle`] or a
+    /// [`TxEnvelope`]
+    pub fn simulate<'a, 'b: 'a>(
+        &self,
+        item: &'a BundleOrTx<'b>,
+    ) -> Result<SimOutcome<&'a BundleOrTx<'b>>, SignetEthBundleError<CacheOnWrite<Db>>> {
+        match item {
+            BundleOrTx::Bundle(bundle) => {
+                Ok(self.simulate_bundle(bundle.as_ref())?.map_in(|_| item))
+            }
+            BundleOrTx::Tx(tx) => Ok(self.simulate_tx(tx.as_ref())?.map_in(|_| item)),
+        }
+    }
+}
+
+impl<Db, C, B, Insp> SimEnv<Db, C, B, Insp>
+where
+    Db: TryCachingDb + DatabaseRef + Clone + Sync,
+    C: Cfg + Sync,
+    B: Block + Sync,
+    Insp: Inspector<Ctx<CacheOnWrite<Db>>> + Default + Sync,
+{
+    /// Accepts a cache from the simulation and extends the database with it.
+    pub fn accept_cache(&mut self, cache: Cache) -> Result<(), <Db as TryCachingDb>::Error> {
+        self.db_mut().try_extend(cache)
+    }
+}

src/tasks/sim/mod.rs

@@ -0,0 +1,11 @@
+mod outcome;
+pub use outcome::SimOutcome;
+
+mod env;
+pub use env::SimEnv;
+
+mod task;
+pub use task::SimTask;
+
+mod item;
+pub use item::BundleOrTx;

src/tasks/sim/outcome.rs

@@ -0,0 +1,69 @@
+use alloy::primitives::U256;
+use trevm::revm::database::Cache;
+
+/// An evaluated EVM transaction with a particular score.
+#[derive(Debug, Clone)]
+pub struct SimOutcome<In, Out = Cache, S = U256> {
+    /// The transaction or bundle being executed.
+    input: In,
+    /// The result of the tx/bundle execution.
+    output: Out,
+    /// The score calculated by the evaluation function.
+    score: S,
+}
+
+impl<In, Out, S> SimOutcome<In, Out, S> {
+    /// Creates a new `Best` instance.
+    pub fn new(input: In, output: Out, evaluator: impl FnOnce(&Out) -> S) -> Self {
+        let score = evaluator(&output);
+        Self::new_unchecked(input, output, score)
+    }
+
+    /// Transform the input type to a new type, while keeping the output and
+    /// score the same.
+    pub fn map_in_into<In2>(self) -> SimOutcome<In2, Out, S>
+    where
+        In2: From<In>,
+    {
+        self.map_in(Into::into)
+    }
+
+    /// Transform the input type to a new type using a closure.
+    pub fn map_in<In2, F>(self, f: F) -> SimOutcome<In2, Out, S>
+    where
+        F: FnOnce(In) -> In2,
+    {
+        let input = f(self.input);
+        SimOutcome { input, output: self.output, score: self.score }
+    }
+
+    /// Creates a new `Best` instance without evaluating the score.
+    pub fn new_unchecked(input: In, output: Out, score: S) -> Self {
+        Self { input, output, score }
+    }
+
+    /// Get a reference to the input, usually a transaction or bundle.
+    pub fn input(&self) -> &In {
+        &self.input
+    }
+
+    /// Get a reference to the output.
+    pub fn output(&self) -> &Out {
+        &self.output
+    }
+
+    /// Get a reference to the score.
+    pub fn score(&self) -> &S {
+        &self.score
+    }
+
+    /// Deconstruct the `SimOutcome` into its parts.
+    pub fn into_parts(self) -> (In, Out, S) {
+        (self.input, self.output, self.score)
+    }
+
+    /// Deconstruct the `SimOutcome` into its output.
+    pub fn into_output(self) -> Out {
+        self.output
+    }
+}