fix: add state source type support

Author: dylanlott

src/test_utils/block.rs

@@ -3,7 +3,7 @@
 //! for testing block simulation.
 
 use super::{
-    db::TestDb,
+    db::{TestDb, TestStateSource},
     env::{TestHostEnv, TestRollupEnv, TestSimEnvBuilder},
 };
 use signet_sim::{BlockBuild, BuiltBlock, SimCache};
@@ -12,7 +12,8 @@ use tokio::time::Instant;
 use trevm::revm::inspector::NoOpInspector;
 
 /// Test block builder type using in-memory databases.
-pub type TestBlockBuild = BlockBuild<TestDb, TestDb, NoOpInspector, NoOpInspector>;
+pub type TestBlockBuild =
+    BlockBuild<TestDb, TestDb, TestStateSource, TestStateSource, NoOpInspector, NoOpInspector>;
 
 /// Builder for creating test `BlockBuild` instances.
 /// Configures all the parameters needed for block simulation
@@ -110,12 +111,15 @@ impl TestBlockBuildBuilder {
     /// This creates a `BlockBuild` ready for simulation.
     /// Call `.build().await` on the result to execute the simulation and get a `BuiltBlock`.
     pub fn build(self) -> TestBlockBuild {
-        let (rollup_env, host_env) = match (self.rollup_env, self.host_env) {
-            (Some(rollup), Some(host)) => (rollup, host),
-            _ => {
-                let builder = self.sim_env_builder.unwrap_or_default();
-                builder.build()
+        let sim_env_builder = self.sim_env_builder.unwrap_or_default();
+
+        let (rollup_env, host_env, ru_source, host_source) = match (self.rollup_env, self.host_env)
+        {
+            (Some(rollup), Some(host)) => {
+                let (ru_source, host_source) = sim_env_builder.build_state_sources();
+                (rollup, host, ru_source, host_source)
             }
+            _ => sim_env_builder.build_with_sources(),
         };
 
         let finish_by = Instant::now() + self.deadline_duration;
@@ -128,6 +132,8 @@ impl TestBlockBuildBuilder {
             self.sim_cache,
             self.max_gas,
             self.max_host_gas,
+            ru_source,
+            host_source,
         )
     }
 }

src/test_utils/db.rs

@@ -3,8 +3,11 @@
 //! for testing block simulation without requiring network access.
 
 use alloy::primitives::{Address, B256, U256};
+use signet_sim::AcctInfo;
 use trevm::revm::{
     database::{CacheDB, EmptyDB},
+    database_interface::DatabaseRef,
+    primitives::KECCAK_EMPTY,
     state::AccountInfo,
 };
 
@@ -14,6 +17,35 @@ use trevm::revm::{
 /// with `RollupEnv` and `HostEnv` for offline simulation testing.
 pub type TestDb = CacheDB<EmptyDB>;
 
+/// A [`StateSource`] backed by a [`TestDb`] for offline testing.
+///
+/// This wraps an in-memory database and implements [`signet_sim::StateSource`]
+/// so it can be used as the async state source parameter in [`BlockBuild::new`].
+///
+/// [`StateSource`]: signet_sim::StateSource
+/// [`BlockBuild::new`]: signet_sim::BlockBuild::new
+#[derive(Debug, Clone)]
+pub struct TestStateSource {
+    db: TestDb,
+}
+
+impl TestStateSource {
+    /// Create a new [`TestStateSource`] from a [`TestDb`].
+    pub const fn new(db: TestDb) -> Self {
+        Self { db }
+    }
+}
+
+impl signet_sim::StateSource for TestStateSource {
+    type Error = <TestDb as DatabaseRef>::Error;
+
+    async fn account_details(&self, address: &Address) -> Result<AcctInfo, Self::Error> {
+        let info = self.db.basic_ref(*address)?.unwrap_or_default();
+        let has_code = info.code_hash() != KECCAK_EMPTY;
+        Ok(AcctInfo { nonce: info.nonce, balance: info.balance, has_code })
+    }
+}
+
 /// Builder for creating pre-populated test databases.
 /// Use this builder to set up blockchain state (accounts, contracts, storage)
 /// before running simulations.

src/test_utils/env.rs

@@ -2,7 +2,7 @@
 //! This module provides builders for creating `RollupEnv` and `HostEnv`
 //! instances with in-memory databases for offline testing.
 
-use super::db::{TestDb, TestDbBuilder};
+use super::db::{TestDb, TestDbBuilder, TestStateSource};
 use crate::tasks::block::cfg::SignetCfgEnv;
 use alloy::primitives::{Address, B256, U256};
 use signet_constants::SignetSystemConstants;
@@ -117,12 +117,27 @@ impl TestSimEnvBuilder {
         HostEnv::new(self.host_db.clone(), self.constants.clone(), &cfg, &self.host_block_env)
     }
 
+    /// Build [`TestStateSource`] instances from the current databases.
+    pub fn build_state_sources(&self) -> (TestStateSource, TestStateSource) {
+        (TestStateSource::new(self.rollup_db.clone()), TestStateSource::new(self.host_db.clone()))
+    }
+
     /// Build both environments as a tuple.
     pub fn build(self) -> (TestRollupEnv, TestHostEnv) {
         let rollup = self.build_rollup_env();
         let host = self.build_host_env();
         (rollup, host)
     }
+
+    /// Build environments and state sources together.
+    pub fn build_with_sources(
+        &self,
+    ) -> (TestRollupEnv, TestHostEnv, TestStateSource, TestStateSource) {
+        let rollup = self.build_rollup_env();
+        let host = self.build_host_env();
+        let (ru_source, host_source) = self.build_state_sources();
+        (rollup, host, ru_source, host_source)
+    }
 }
 
 #[cfg(test)]

src/test_utils/mod.rs

@@ -11,7 +11,7 @@ mod tx;
 
 // Re-export test harness components
 pub use block::{TestBlockBuild, TestBlockBuildBuilder, quick_build_block};
-pub use db::{TestDb, TestDbBuilder};
+pub use db::{TestDb, TestDbBuilder, TestStateSource};
 pub use env::{TestHostEnv, TestRollupEnv, TestSimEnvBuilder};
 pub use scenarios::{
     DEFAULT_BALANCE, DEFAULT_BASEFEE, basic_scenario, custom_funded_scenario, funded_test_db,