fix: move block env creation for host to env task (#169)

* fix: move block env creation for host to sim task

* rename vars to clearly distinguish between rollup and host (#170)

* comment

---------

Co-authored-by: Anna Carroll <anna-carroll@users.noreply.github.com>
Co-authored-by: Anna Carroll <anna.s.carroll@gmail.com>

Author: prestwich

src/tasks/block/sim.rs

@@ -5,7 +5,7 @@ use crate::{
     config::{BuilderConfig, HostProvider, RuProvider},
     tasks::env::SimEnv,
 };
-use alloy::{eips::BlockId, network::Ethereum, primitives::BlockNumber};
+use alloy::{consensus::Header, eips::BlockId, network::Ethereum};
 use init4_bin_base::{
     deps::metrics::{counter, histogram},
     utils::calc::SlotCalculator,
@@ -21,13 +21,9 @@ use tokio::{
     task::JoinHandle,
 };
 use tracing::{Instrument, Span, debug, instrument};
-use trevm::{
-    Block,
-    revm::{
-        context::BlockEnv,
-        database::{AlloyDB, WrapDatabaseAsync},
-        inspector::NoOpInspector,
-    },
+use trevm::revm::{
+    database::{AlloyDB, WrapDatabaseAsync},
+    inspector::NoOpInspector,
 };
 
 type HostAlloyDatabaseProvider = WrapDatabaseAsync<AlloyDB<Ethereum, HostProvider>>;
@@ -58,8 +54,18 @@ pub struct SimResult {
 }
 
 impl SimResult {
+    /// Get a reference to the previous host header.
+    pub const fn prev_host(&self) -> &Header {
+        self.sim_env.prev_host()
+    }
+
+    /// Get a reference to the previous rollup header.
+    pub const fn prev_rollup(&self) -> &Header {
+        self.sim_env.prev_rollup()
+    }
+
     /// Returns the block number of the built block.
-    pub const fn block_number(&self) -> u64 {
+    pub const fn rollup_block_number(&self) -> u64 {
         self.block.block_number()
     }
 
@@ -134,7 +140,7 @@ impl Simulator {
         sim_env: SimEnv,
     ) -> eyre::Result<BuiltBlock> {
         let concurrency_limit = self.config.concurrency_limit();
-        let max_host_gas = self.config.max_host_gas(sim_env.prev_host.gas_limit);
+        let max_host_gas = self.config.max_host_gas(sim_env.prev_host().gas_limit);
 
         let (rollup_env, host_env) = self.create_envs(constants, &sim_env).await;
 
@@ -170,29 +176,25 @@ impl Simulator {
         HostEnv<HostAlloyDatabaseProvider, NoOpInspector>,
     ) {
         // Host DB and Env
-        let host_block_number = BlockNumber::from(sim_env.prev_host.number);
+        let host_block_number = sim_env.host_block_number();
         let host_db = self.create_host_db(host_block_number).await;
-        let mut host_block_env = BlockEnv::default();
-        sim_env.prev_host.fill_block_env(&mut host_block_env);
 
         let host_env = HostEnv::<HostAlloyDatabaseProvider, NoOpInspector>::new(
             host_db,
             constants.clone(),
             &self.config.host_cfg_env(),
-            &host_block_env,
+            sim_env.host_env(),
         );
 
         // Rollup DB and Env
-        let rollup_block_number = BlockNumber::from(sim_env.prev_header.number);
+        let rollup_block_number = sim_env.rollup_block_number();
         let rollup_db = self.create_rollup_db(rollup_block_number);
-        let mut rollup_block_env = BlockEnv::default();
-        sim_env.prev_header.fill_block_env(&mut rollup_block_env);
 
         let rollup_env = RollupEnv::<RollupAlloyDatabaseProvider, NoOpInspector>::new(
             rollup_db,
             constants,
             &self.config.ru_cfg_env(),
-            &rollup_block_env,
+            sim_env.rollup_env(),
         );
 
         (rollup_env, host_env)

src/tasks/cache/task.rs

@@ -16,7 +16,7 @@ use tracing::{debug, info};
 #[derive(Debug)]
 pub struct CacheTask {
     /// The channel to receive the block environment.
-    env: watch::Receiver<Option<SimEnv>>,
+    envs: watch::Receiver<Option<SimEnv>>,
     /// The channel to receive the transaction bundles.
     bundles: mpsc::UnboundedReceiver<TxCacheBundle>,
     /// The channel to receive the transactions.
@@ -30,24 +30,27 @@ impl CacheTask {
         bundles: mpsc::UnboundedReceiver<TxCacheBundle>,
         txns: mpsc::UnboundedReceiver<TxEnvelope>,
     ) -> Self {
-        Self { env, bundles, txns }
+        Self { envs: env, bundles, txns }
     }
 
     async fn task_future(mut self, cache: SimCache) {
         loop {
             let mut basefee = 0;
             tokio::select! {
                 biased;
-                res = self.env.changed() => {
+                res = self.envs.changed() => {
                     if res.is_err() {
                         debug!("Cache task: env channel closed, exiting");
                         break;
                     }
-                    if let Some(env) = self.env.borrow_and_update().as_ref() {
-                        basefee = env.block_env.basefee;
-                        info!(basefee, block_env_number = env.block_env.number.to::<u64>(), block_env_timestamp = env.block_env.timestamp.to::<u64>(), "rollup block env changed, clearing cache");
+
+                    if let Some(env) = self.envs.borrow_and_update().as_ref() {
+                        let sim_env = env.rollup_env();
+
+                        basefee = sim_env.basefee;
+                        info!(basefee, block_env_number = sim_env.number.to::<u64>(), block_env_timestamp = sim_env.timestamp.to::<u64>(), "rollup block env changed, clearing cache");
                         cache.clean(
-                            env.block_env.number.to(), env.block_env.timestamp.to()
+                            sim_env.number.to(), sim_env.timestamp.to()
                         );
                     }
                 }

src/tasks/env.rs

@@ -25,31 +25,82 @@ pub struct EnvTask {
     ru_provider: RuProvider,
 }
 
-/// Contains a signet BlockEnv and its corresponding host Header.
+/// An environment for simulating a block.
+#[derive(Debug, Clone)]
+pub struct Environment {
+    block_env: BlockEnv,
+    prev_header: Header,
+}
+
+impl Environment {
+    /// Create a new `Environment` with the given block environment and
+    /// previous header.
+    pub const fn new(block_env: BlockEnv, prev_header: Header) -> Self {
+        Self { block_env, prev_header }
+    }
+
+    /// Get a reference to the block environment.
+    pub const fn block_env(&self) -> &BlockEnv {
+        &self.block_env
+    }
+
+    /// Get a reference to the previous block header.
+    pub const fn prev_header(&self) -> &Header {
+        &self.prev_header
+    }
+
+    /// Create a new empty `Environment` for testing purposes.
+    #[doc(hidden)]
+    pub fn for_testing() -> Self {
+        Self { block_env: Default::default(), prev_header: Header::default() }
+    }
+}
+
+/// Contains environments to simulate both host and rollup blocks.
 #[derive(Debug, Clone)]
 pub struct SimEnv {
-    /// The signet block environment, for rollup block simulation.
-    pub block_env: BlockEnv,
-    /// The header of the previous rollup block.
-    pub prev_header: Header,
-    /// The header of the previous host block.
-    pub prev_host: Header,
-    /// A tracing span associated with this block
+    /// The host environment, for host block simulation.
+    pub host: Environment,
+
+    /// The rollup environment, for rollup block simulation.
+    pub rollup: Environment,
+
+    /// A tracing span associated with this block simulation.
     pub span: Span,
 }
 
 impl SimEnv {
-    /// Returns the block number of the signet block environment.
-    pub const fn block_number(&self) -> u64 {
-        self.prev_header.number.saturating_add(1)
+    /// Get a reference to previous rollup header.
+    pub const fn prev_rollup(&self) -> &Header {
+        &self.rollup.prev_header
+    }
+
+    /// Get a reference to the previous host header.
+    pub const fn prev_host(&self) -> &Header {
+        &self.host.prev_header
+    }
+
+    /// Get the block number of the rollup block environment.
+    pub const fn rollup_block_number(&self) -> u64 {
+        self.prev_rollup().number.saturating_add(1)
     }
 
-    /// Returns the host block number for the signet block environment.
+    /// Get the block number for the host block environment.
     pub const fn host_block_number(&self) -> u64 {
-        self.prev_host.number.saturating_add(1)
+        self.prev_host().number.saturating_add(1)
     }
 
-    /// Returns a reference to the tracing span associated with this block env.
+    /// Get a reference to the rollup block environment.
+    pub const fn rollup_env(&self) -> &BlockEnv {
+        &self.rollup.block_env
+    }
+
+    /// Get a reference to the host block environment.
+    pub const fn host_env(&self) -> &BlockEnv {
+        &self.host.block_env
+    }
+
+    /// Get a reference to the tracing span associated with this block env.
     pub const fn span(&self) -> &Span {
         &self.span
     }
@@ -71,8 +122,8 @@ impl EnvTask {
     }
 
     /// Construct a [`BlockEnv`] by from the previous block header.
-    fn construct_block_env(&self, previous: &Header) -> BlockEnv {
-        BlockEnv {
+    fn construct_block_env(&self, previous: Header) -> Environment {
+        let env = BlockEnv {
             number: U256::from(previous.number + 1),
             beneficiary: self.config.builder_rewards_address,
             // NB: EXACTLY the same as the previous block
@@ -87,14 +138,15 @@ impl EnvTask {
                 excess_blob_gas: 0,
                 blob_gasprice: 0,
             }),
-        }
+        };
+        Environment::new(env, previous)
     }
 
     /// Returns a sender that sends [`SimEnv`] for communicating the next block environment.
     async fn task_fut(self, sender: watch::Sender<Option<SimEnv>>) {
         let span = info_span!("EnvTask::task_fut::init");
 
-        let mut headers = match self.ru_provider.subscribe_blocks().await {
+        let mut rollup_headers = match self.ru_provider.subscribe_blocks().await {
             Ok(poller) => poller,
             Err(err) => {
                 span_error!(span, %err, "Failed to subscribe to blocks");
@@ -106,7 +158,7 @@ impl EnvTask {
         drop(span);
 
         while let Some(rollup_header) =
-            headers.next().instrument(info_span!("EnvTask::task_fut::stream")).await
+            rollup_headers.next().instrument(info_span!("EnvTask::task_fut::stream")).await
         {
             let host_block_number =
                 self.config.constants.rollup_block_to_host_block_num(rollup_header.number);
@@ -118,7 +170,7 @@ impl EnvTask {
                 span,
                 error!("error fetching previous host block - skipping block submission")
             );
-            let prev_host = opt_unwrap_or_continue!(
+            let host_header = opt_unwrap_or_continue!(
                 host_block_opt,
                 span,
                 warn!("previous host block not found - skipping block submission")
@@ -127,23 +179,17 @@ impl EnvTask {
             .inner;
 
             // Construct the block env using the previous block header
-            let signet_env = self.construct_block_env(&rollup_header);
+            let rollup_env = self.construct_block_env(rollup_header.into());
+            let host_env = self.construct_block_env(host_header);
+
             span_debug!(
                 span,
-                signet_env_number = signet_env.number.to::<u64>(),
-                signet_env_basefee = signet_env.basefee,
-                "constructed signet block env"
+                rollup_env_number = rollup_env.block_env.number.to::<u64>(),
+                rollup_env_basefee = rollup_env.block_env.basefee,
+                "constructed block env"
             );
 
-            if sender
-                .send(Some(SimEnv {
-                    span,
-                    block_env: signet_env,
-                    prev_header: rollup_header.inner,
-                    prev_host,
-                }))
-                .is_err()
-            {
+            if sender.send(Some(SimEnv { span, rollup: rollup_env, host: host_env })).is_err() {
                 // The receiver has been dropped, so we can stop the task.
                 tracing::debug!("receiver dropped, stopping task");
                 break;

src/tasks/submit/flashbots.rs

@@ -79,7 +79,7 @@ impl FlashbotsTask {
             self.config.clone(),
         );
 
-        let tx = prep.prep_transaction(&sim_result.sim_env.prev_host).await?;
+        let tx = prep.prep_transaction(sim_result.prev_host()).await?;
 
         let sendable = self.host_provider().fill(tx.into_request()).await?;
 

tests/block_builder_test.rs

@@ -1,6 +1,7 @@
 //! Tests for the block building task.
 
 use alloy::{
+    eips::BlockId,
     network::Ethereum,
     node_bindings::Anvil,
     primitives::U256,
@@ -10,7 +11,7 @@ use alloy::{
 use builder::{
     tasks::{
         block::sim::Simulator,
-        env::{EnvTask, SimEnv},
+        env::{EnvTask, Environment, SimEnv},
     },
     test_utils::{new_signed_tx, setup_logging, setup_test_config, test_block_env},
 };
@@ -26,8 +27,6 @@ use std::time::{Duration, Instant};
 #[ignore = "integration test"]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_handle_build() {
-    use alloy::eips::BlockId;
-
     setup_logging();
 
     // Make a test config
@@ -69,17 +68,15 @@ async fn test_handle_build() {
 
     // Setup the block envs
     let finish_by = Instant::now() + Duration::from_secs(2);
-    let host_header = host_provider.get_block(BlockId::latest()).await.unwrap().unwrap().header;
     let ru_header = ru_provider.get_block(BlockId::latest()).await.unwrap().unwrap().header.inner;
     let number = ru_header.number + 1;
     let timestamp = ru_header.timestamp + config.slot_calculator.slot_duration();
     let block_env = test_block_env(config, number, 7, timestamp);
 
     // Spawn the block builder task
     let sim_env = SimEnv {
-        block_env: block_env.clone(),
-        prev_header: ru_header.clone(),
-        prev_host: host_header.into(),
+        host: Environment::for_testing(),
+        rollup: Environment::new(block_env, ru_header),
         span: tracing::Span::none(),
     };
     let got = block_builder.handle_build(constants, sim_items, finish_by, sim_env).await;