WEB3-322: feat: Add HostEvmEnv::extend method for merging environments

crates/steel/src/host/db/proof.rs

@@ -23,14 +23,15 @@ use alloy::{
     transports::Transport,
 };
 use alloy_primitives::{
-    map::{hash_map, AddressHashMap, B256HashMap, B256HashSet, HashSet},
+    map::{hash_map, AddressHashMap, B256HashMap, B256HashSet, HashMap, HashSet},
     Address, BlockNumber, Bytes, StorageKey, StorageValue, B256, U256,
 };
 use anyhow::{ensure, Context, Result};
 use revm::{
     primitives::{AccountInfo, Bytecode},
     Database,
 };
+use std::hash::{BuildHasher, Hash};
 
 /// A simple revm [Database] wrapper that records all DB queries.
 pub struct ProofDb<D> {
@@ -41,23 +42,28 @@ pub struct ProofDb<D> {
     inner: D,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
 struct AccountProof {
     /// The inclusion proof for this account.
     account_proof: Vec<Bytes>,
     /// The MPT inclusion proofs for several storage slots.
     storage_proofs: B256HashMap<StorageProof>,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
 struct StorageProof {
     /// The value that this key holds.
     value: StorageValue,
     /// In MPT inclusion proof for this particular slot.
     proof: Vec<Bytes>,
 }
 
-impl<D: Database> ProofDb<D> {
+impl<D> ProofDb<D> {
     /// Creates a new ProofDb instance, with a [Database].
-    pub fn new(db: D) -> Self {
+    pub fn new(db: D) -> Self
+    where
+        D: Database,
+    {
         Self {
             accounts: Default::default(),
             contracts: Default::default(),
@@ -83,6 +89,15 @@ impl<D: Database> ProofDb<D> {
     pub fn inner(&self) -> &D {
         &self.inner
     }
+
+    /// Extends the `ProofDb` with the contents of another `ProofDb`. It panics if they are not
+    /// consistent.
+    pub(crate) fn extend(&mut self, other: ProofDb<D>) {
+        extend_checked(&mut self.accounts, other.accounts);
+        extend_checked(&mut self.contracts, other.contracts);
+        extend_checked(&mut self.proofs, other.proofs);
+        self.block_hash_numbers.extend(other.block_hash_numbers);
+    }
 }
 
 impl<T: Transport + Clone, N: Network, P: Provider<T, N>> ProofDb<AlloyDb<T, N, P>> {
@@ -265,6 +280,32 @@ impl<DB: Database> Database for ProofDb<DB> {
     }
 }
 
+/// Extends a `HashMap` with the contents of an iterator.
+fn extend_checked<K, V, S, T>(map: &mut HashMap<K, V, S>, iter: T)
+where
+    K: Eq + Hash,
+    V: PartialEq,
+    S: BuildHasher,
+    T: IntoIterator<Item = (K, V)>,
+{
+    let iter = iter.into_iter();
+    let (lower_bound, _) = iter.size_hint();
+    map.reserve(lower_bound);
+
+    for (k, v) in iter {
+        match map.entry(k) {
+            hash_map::Entry::Vacant(entry) => {
+                entry.insert(v);
+            }
+            hash_map::Entry::Occupied(entry) => {
+                if entry.get() != &v {
+                    panic!("mismatching values for key")
+                }
+            }
+        }
+    }
+}
+
 fn filter_existing_keys(account_proof: Option<&AccountProof>) -> impl Fn(&StorageKey) -> bool + '_ {
     move |key| {
         !account_proof

crates/steel/src/host/mod.rs

@@ -174,6 +174,69 @@ impl<D, H: EvmBlockHeader, C> HostEvmEnv<D, H, C> {
 
         self
     }
+
+    /// Extends the environment with the contents of another compatible environment.
+    ///
+    /// ### Errors
+    ///
+    /// It returns an error if the environments are inconsistent, specifically if:
+    /// - The configurations don't match
+    /// - The headers don't match
+    ///
+    /// ### Panics
+    ///
+    /// It panics if the database states conflict.
+    ///
+    /// ### Use Cases
+    ///
+    /// This method is particularly useful for combining results from parallel preflights,
+    /// allowing you to execute multiple independent operations and merge their environments.
+    ///
+    /// ### Example
+    /// ```rust
+    /// # use risc0_steel::{ethereum::EthEvmEnv, Contract};
+    /// # use alloy_primitives::address;
+    /// # use alloy_sol_types::sol;
+    ///
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// # async fn main() -> anyhow::Result<()> {
+    /// # sol! {
+    /// #    interface IERC20 {
+    /// #        function balanceOf(address account) external view returns (uint);
+    /// #    }
+    /// # }
+    /// let call =
+    ///     IERC20::balanceOfCall { account: address!("F977814e90dA44bFA03b6295A0616a897441aceC") };
+    /// # let usdt_addr = address!("dAC17F958D2ee523a2206206994597C13D831ec7");
+    /// # let usdc_addr = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
+    ///
+    /// let url = "https://ethereum-rpc.publicnode.com".parse()?;
+    /// let builder = EthEvmEnv::builder().rpc(url);
+    ///
+    /// let mut env1 = builder.clone().build().await?;
+    /// let mut contract1 = Contract::preflight(usdt_addr, &mut env1);
+    /// let mut env2 = builder.clone().build().await?;
+    /// let mut contract2 = Contract::preflight(usdc_addr, &mut env2);
+    ///
+    /// tokio::join!(contract1.call_builder(&call).call(), contract2.call_builder(&call).call());
+    ///
+    /// env1.extend(env2)?;
+    /// let evm_input = env1.into_input().await?;
+    ///
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn extend(&mut self, other: Self) -> Result<()> {
+        ensure!(self.cfg_env == other.cfg_env, "configuration mismatch");
+        ensure!(
+            self.header.seal() == other.header.seal(),
+            "execution header mismatch"
+        );
+        // the commitments do not need to match as long as the cfg_env is consistent
+        self.db_mut().extend(other.db.unwrap());
+
+        Ok(())
+    }
 }
 
 impl<T, N, P, H> HostEvmEnv<AlloyDb<T, N, P>, H, ()>