feat: start host env breakout (#125)

Author: prestwich

crates/sim/src/env/host.rs

@@ -0,0 +1,69 @@
+use crate::{InnerDb, SimDb, TimeLimited};
+use std::{marker::PhantomData, sync::Arc, time::Instant};
+use trevm::{
+    helpers::Ctx,
+    inspectors::{Layered, TimeLimit},
+    revm::{
+        database::CacheDB,
+        inspector::{Inspector, NoOpInspector},
+        DatabaseRef,
+    },
+    TrevmBuilder,
+};
+
+/// A host simulation environment.
+#[derive(Debug)]
+pub struct HostEnv<Db, Insp = NoOpInspector> {
+    db: InnerDb<Db>,
+    _pd: PhantomData<fn() -> Insp>,
+}
+
+impl<Db, Insp> Clone for HostEnv<Db, Insp> {
+    fn clone(&self) -> Self {
+        Self { db: self.db.clone(), _pd: PhantomData }
+    }
+}
+
+impl<Db> From<Db> for HostEnv<Db, NoOpInspector>
+where
+    Db: DatabaseRef + Send + Sync,
+{
+    fn from(db: Db) -> Self {
+        Self::new(db)
+    }
+}
+
+impl<Db, Insp> HostEnv<Db, Insp> {
+    /// Create a new host environment.
+    pub fn new(db: Db) -> Self {
+        Self { db: Arc::new(CacheDB::new(db)), _pd: PhantomData }
+    }
+
+    /// Get a mutable reference to the inner database.
+    pub const fn db_mut(&mut self) -> &mut InnerDb<Db> {
+        &mut self.db
+    }
+}
+
+impl<Db, Insp> HostEnv<Db, Insp>
+where
+    Db: DatabaseRef + Send + Sync,
+    Insp: Inspector<Ctx<SimDb<Db>>> + Default + Sync,
+{
+    /// Connect a fresh database for the simulation.
+    pub fn sim_db(&self) -> crate::SimDb<Db> {
+        crate::SimDb::new(self.db.clone())
+    }
+
+    /// Create a new EVM for the host environment that will finish by the
+    /// given instant.
+    pub fn create_evm(
+        &self,
+        finish_by: Instant,
+    ) -> trevm::EvmNeedsCfg<crate::SimDb<Db>, TimeLimited<Insp>> {
+        let db = self.sim_db();
+        let inspector = Layered::new(TimeLimit::new(finish_by - Instant::now()), Insp::default());
+
+        TrevmBuilder::new().with_insp(inspector).with_db(db).build_trevm().unwrap()
+    }
+}

crates/sim/src/env/mod.rs

@@ -0,0 +1,11 @@
+mod host;
+pub use host::HostEnv;
+
+mod rollup;
+pub use rollup::RollupEnv;
+
+mod shared;
+pub use shared::SharedSimEnv;
+
+mod sim_env;
+pub use sim_env::SimEnv;

crates/sim/src/env/rollup.rs

@@ -0,0 +1,88 @@
+use crate::{InnerDb, SimDb, TimeLimited};
+use signet_evm::EvmNeedsCfg;
+use signet_types::constants::SignetSystemConstants;
+use std::{marker::PhantomData, sync::Arc, time::Instant};
+use trevm::{
+    db::{cow::CacheOnWrite, TryCachingDb},
+    helpers::Ctx,
+    inspectors::{Layered, TimeLimit},
+    revm::{
+        database::{Cache, CacheDB},
+        inspector::NoOpInspector,
+        DatabaseRef, Inspector,
+    },
+};
+
+/// A rollup simulation environment.
+#[derive(Debug)]
+pub struct RollupEnv<Db, Insp = NoOpInspector> {
+    db: InnerDb<Db>,
+    constants: SignetSystemConstants,
+    _pd: PhantomData<fn() -> Insp>,
+}
+
+impl<Db, Insp> Clone for RollupEnv<Db, Insp> {
+    fn clone(&self) -> Self {
+        Self { db: self.db.clone(), constants: self.constants.clone(), _pd: PhantomData }
+    }
+}
+
+impl<Db, Insp> RollupEnv<Db, Insp> {
+    /// Create a new rollup environment.
+    pub fn new(db: Db, constants: SignetSystemConstants) -> Self {
+        Self { db: Arc::new(CacheDB::new(db)), constants, _pd: PhantomData }
+    }
+
+    /// Get a mutable reference to the inner database.
+    pub const fn db_mut(&mut self) -> &mut InnerDb<Db> {
+        &mut self.db
+    }
+
+    /// Get the constants for this environment.
+    pub const fn constants(&self) -> &SignetSystemConstants {
+        &self.constants
+    }
+}
+
+impl<Db, Insp> RollupEnv<Db, Insp>
+where
+    Db: DatabaseRef + Send + Sync,
+    Insp: Inspector<Ctx<SimDb<Db>>> + Default + Sync,
+{
+    /// Connect a fresh database for the simulation.
+    pub fn sim_db(&self) -> SimDb<Db> {
+        CacheOnWrite::new(self.db.clone())
+    }
+
+    /// Create a new EVM for the rollup environment that will finish by the
+    /// given instant.
+    pub fn create_evm(&self, finish_by: Instant) -> EvmNeedsCfg<SimDb<Db>, TimeLimited<Insp>> {
+        let db = self.sim_db();
+
+        let inspector = Layered::new(TimeLimit::new(finish_by - Instant::now()), Insp::default());
+
+        signet_evm::signet_evm_with_inspector(db, inspector, self.constants.clone())
+    }
+}
+
+impl<Db, Insp> RollupEnv<Db, Insp>
+where
+    Db: DatabaseRef,
+    Insp: Inspector<Ctx<SimDb<Db>>> + Default + Sync,
+{
+    /// Accepts a cache from the simulation and extends the database with it.
+    pub fn accept_cache(
+        &mut self,
+        cache: Cache,
+    ) -> Result<(), <InnerDb<Db> as TryCachingDb>::Error> {
+        self.db_mut().try_extend(cache)
+    }
+
+    /// Accepts a cache from the simulation and extends the database with it.
+    pub fn accept_cache_ref(
+        &mut self,
+        cache: &Cache,
+    ) -> Result<(), <InnerDb<Db> as TryCachingDb>::Error> {
+        self.db_mut().try_extend_ref(cache)
+    }
+}

crates/sim/src/env/shared.rs

@@ -0,0 +1,103 @@
+use crate::{env::RollupEnv, outcome::SimulatedItem, SimCache, SimDb, SimEnv};
+use core::fmt;
+use std::{ops::Deref, sync::Arc};
+use tokio::{select, sync::watch};
+use tracing::{instrument, trace};
+use trevm::{
+    helpers::Ctx,
+    revm::{inspector::NoOpInspector, DatabaseRef, Inspector},
+    Block, Cfg,
+};
+
+/// A simulation environment.
+///
+/// Contains enough information to run a simulation.
+pub struct SharedSimEnv<Db, Insp = NoOpInspector> {
+    inner: Arc<SimEnv<Db, Insp>>,
+}
+
+impl<Db, Insp> fmt::Debug for SharedSimEnv<Db, Insp> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("SharedSimEnv")
+            .field("finish_by", &self.inner.finish_by())
+            .field("concurrency_limit", &self.inner.concurrency_limit())
+            .finish_non_exhaustive()
+    }
+}
+
+impl<Db, Insp> Deref for SharedSimEnv<Db, Insp> {
+    type Target = SimEnv<Db, Insp>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<Db, Insp> From<SimEnv<Db, Insp>> for SharedSimEnv<Db, Insp>
+where
+    Db: DatabaseRef + Send + Sync + 'static,
+    Insp: Inspector<Ctx<SimDb<Db>>> + Default + Sync + 'static,
+{
+    fn from(inner: SimEnv<Db, Insp>) -> Self {
+        Self { inner: Arc::new(inner) }
+    }
+}
+
+impl<Db, Insp> SharedSimEnv<Db, Insp>
+where
+    Db: DatabaseRef + Send + Sync + 'static,
+    Insp: Inspector<Ctx<SimDb<Db>>> + Default + Sync + 'static,
+{
+    /// Creates a new `SimEnv` instance.
+    pub fn new<C, B>(
+        rollup: RollupEnv<Db, Insp>,
+        cfg: C,
+        block: B,
+        finish_by: std::time::Instant,
+        concurrency_limit: usize,
+        sim_items: SimCache,
+    ) -> Self
+    where
+        C: Cfg,
+        B: Block,
+    {
+        SimEnv::new(rollup, cfg, block, finish_by, concurrency_limit, sim_items).into()
+    }
+
+    /// Run a simulation round, returning the best item.
+    #[instrument(skip(self))]
+    pub async fn sim_round(&mut self, max_gas: u64) -> Option<SimulatedItem> {
+        let (best_tx, mut best_watcher) = watch::channel(None);
+
+        let this = self.inner.clone();
+
+        // Spawn a blocking task to run the simulations.
+        let sim_task = tokio::task::spawn_blocking(move || this.sim_round(max_gas, best_tx));
+
+        // Either simulation is done, or we time out
+        select! {
+            _ = tokio::time::sleep_until(self.finish_by().into()) => {
+                trace!("Sim round timed out");
+            },
+            _ = sim_task => {
+                trace!("Sim round done");
+            },
+        }
+
+        // Check what the current best outcome is.
+        let best = best_watcher.borrow_and_update();
+        trace!(score = %best.as_ref().map(|candidate| candidate.score).unwrap_or_default(), "Read outcome from channel");
+        let outcome = best.as_ref()?;
+
+        // Remove the item from the cache.
+        let item = self.sim_items().remove(outcome.cache_rank)?;
+        // Accept the cache from the simulation.
+        Arc::get_mut(&mut self.inner)
+            .expect("sims dropped already")
+            .rollup_mut()
+            .accept_cache_ref(&outcome.cache)
+            .ok()?;
+
+        Some(SimulatedItem { gas_used: outcome.gas_used, score: outcome.score, item })
+    }
+}