refactor: Use BundleBuilder instead of hashmaps

bin/reth/src/commands/stage/drop.rs

@@ -103,7 +103,7 @@ impl Command {
                     Default::default(),
                 )?;
                 let alloc = &self.env.chain.genesis().alloc;
-                insert_genesis_state::<DatabaseEnv>(tx, alloc.len(), alloc.iter())?;
+                insert_genesis_state::<DatabaseEnv>(tx, alloc.iter())?;
             }
             StageEnum::AccountHashing => {
                 tx.clear::<tables::HashedAccounts>()?;

crates/evm/execution-types/src/execution_outcome.rs

@@ -1,15 +1,12 @@
 use reth_primitives::{
-    logs_bloom,
-    revm::compat::{into_reth_acc, into_revm_acc},
-    Account, Address, BlockNumber, Bloom, Bytecode, Log, Receipt, Receipts, Requests, StorageEntry,
-    B256, U256,
+    logs_bloom, revm::compat::into_reth_acc, Account, Address, BlockNumber, Bloom, Bytecode, Log,
+    Receipt, Receipts, Requests, B256, U256,
 };
 use reth_trie::HashedPostState;
 use revm::{
     db::{states::BundleState, BundleAccount},
     primitives::AccountInfo,
 };
-use std::collections::HashMap;
 
 /// Represents the outcome of block execution, including post-execution changes and reverts.
 ///
@@ -36,16 +33,6 @@ pub struct ExecutionOutcome {
     pub requests: Vec<Requests>,
 }
 
-/// Type used to initialize revms bundle state.
-pub type BundleStateInit =
-    HashMap<Address, (Option<Account>, Option<Account>, HashMap<B256, (U256, U256)>)>;
-
-/// Types used inside `RevertsInit` to initialize revms reverts.
-pub type AccountRevertInit = (Option<Option<Account>>, Vec<StorageEntry>);
-
-/// Type used to initialize revms reverts.
-pub type RevertsInit = HashMap<BlockNumber, HashMap<Address, AccountRevertInit>>;
-
 impl ExecutionOutcome {
     /// Creates a new `ExecutionOutcome`.
     ///
@@ -60,48 +47,6 @@ impl ExecutionOutcome {
         Self { bundle, receipts, first_block, requests }
     }
 
-    /// Creates a new `ExecutionOutcome` from initialization parameters.
-    ///
-    /// This constructor initializes a new `ExecutionOutcome` instance using detailed
-    /// initialization parameters.
-    pub fn new_init(
-        state_init: BundleStateInit,
-        revert_init: RevertsInit,
-        contracts_init: Vec<(B256, Bytecode)>,
-        receipts: Receipts,
-        first_block: BlockNumber,
-        requests: Vec<Requests>,
-    ) -> Self {
-        // sort reverts by block number
-        let mut reverts = revert_init.into_iter().collect::<Vec<_>>();
-        reverts.sort_unstable_by_key(|a| a.0);
-
-        // initialize revm bundle
-        let bundle = BundleState::new(
-            state_init.into_iter().map(|(address, (original, present, storage))| {
-                (
-                    address,
-                    original.map(into_revm_acc),
-                    present.map(into_revm_acc),
-                    storage.into_iter().map(|(k, s)| (k.into(), s)).collect(),
-                )
-            }),
-            reverts.into_iter().map(|(_, reverts)| {
-                // does not needs to be sorted, it is done when taking reverts.
-                reverts.into_iter().map(|(address, (original, storage))| {
-                    (
-                        address,
-                        original.map(|i| i.map(into_revm_acc)),
-                        storage.into_iter().map(|entry| (entry.key.into(), entry.value)),
-                    )
-                })
-            }),
-            contracts_init.into_iter().map(|(code_hash, bytecode)| (code_hash, bytecode.0)),
-        );
-
-        Self { bundle, receipts, first_block, requests }
-    }
-
     /// Return revm bundle state.
     pub const fn state(&self) -> &BundleState {
         &self.bundle
@@ -395,37 +340,7 @@ mod tests {
         };
 
         // Assert that creating a new ExecutionOutcome using the constructor matches exec_res
-        assert_eq!(
-            ExecutionOutcome::new(bundle, receipts.clone(), first_block, requests.clone()),
-            exec_res
-        );
-
-        // Create a BundleStateInit object and insert initial data
-        let mut state_init: BundleStateInit = HashMap::new();
-        state_init
-            .insert(Address::new([2; 20]), (None, Some(Account::default()), HashMap::default()));
-
-        // Create a HashMap for account reverts and insert initial data
-        let mut revert_inner: HashMap<Address, AccountRevertInit> = HashMap::new();
-        revert_inner.insert(Address::new([2; 20]), (None, vec![]));
-
-        // Create a RevertsInit object and insert the revert_inner data
-        let mut revert_init: RevertsInit = HashMap::new();
-        revert_init.insert(123, revert_inner);
-
-        // Assert that creating a new ExecutionOutcome using the new_init method matches
-        // exec_res
-        assert_eq!(
-            ExecutionOutcome::new_init(
-                state_init,
-                revert_init,
-                vec![],
-                receipts,
-                first_block,
-                requests,
-            ),
-            exec_res
-        );
+        assert_eq!(ExecutionOutcome::new(bundle, receipts, first_block, requests), exec_res);
     }
 
     #[test]

crates/primitives/src/account.rs

@@ -23,6 +23,12 @@ impl Bytecode {
     }
 }
 
+impl From<Bytecode> for RevmBytecode {
+    fn from(value: Bytecode) -> Self {
+        value.0
+    }
+}
+
 impl Compact for Bytecode {
     fn to_compact<B>(self, buf: &mut B) -> usize
     where

crates/storage/db-common/Cargo.toml

@@ -18,6 +18,7 @@ reth-trie.workspace = true
 reth-etl.workspace = true
 reth-codecs.workspace = true
 reth-stages-types.workspace = true
+reth-revm.workspace = true
 
 # misc
 eyre.workspace = true

crates/storage/db-common/src/init.rs

@@ -6,17 +6,17 @@ use reth_db::tables;
 use reth_db_api::{database::Database, transaction::DbTxMut, DatabaseError};
 use reth_etl::Collector;
 use reth_primitives::{
-    Account, Address, Bytecode, ChainSpec, GenesisAccount, Receipts, StaticFileSegment,
-    StorageEntry, B256, U256,
+    revm::compat::into_revm_acc, Account, Address, Bytecode, ChainSpec, GenesisAccount, Receipts,
+    StaticFileSegment, StorageEntry, B256, U256,
 };
 use reth_provider::{
-    bundle_state::{BundleStateInit, RevertsInit},
     errors::provider::ProviderResult,
     providers::{StaticFileProvider, StaticFileWriter},
     BlockHashReader, BlockNumReader, ChainSpecProvider, DatabaseProviderRW, ExecutionOutcome,
     HashingWriter, HistoryWriter, OriginalValuesKnown, ProviderError, ProviderFactory,
     StageCheckpointWriter, StateWriter, StaticFileProviderFactory,
 };
+use reth_revm::db::states::BundleBuilder;
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_trie::{IntermediateStateRootState, StateRoot as StateRootComputer, StateRootProgress};
 use serde::{Deserialize, Serialize};
@@ -116,7 +116,7 @@ pub fn init_genesis<DB: Database>(factory: ProviderFactory<DB>) -> Result<B256,
     let static_file_provider = factory.static_file_provider();
     insert_genesis_header::<DB>(tx, &static_file_provider, chain.clone())?;
 
-    insert_genesis_state::<DB>(tx, alloc.len(), alloc.iter())?;
+    insert_genesis_state::<DB>(tx, alloc.iter())?;
 
     // insert sync stage
     for stage in StageId::ALL {
@@ -132,28 +132,24 @@ pub fn init_genesis<DB: Database>(factory: ProviderFactory<DB>) -> Result<B256,
 /// Inserts the genesis state into the database.
 pub fn insert_genesis_state<'a, 'b, DB: Database>(
     tx: &<DB as Database>::TXMut,
-    capacity: usize,
     alloc: impl Iterator<Item = (&'a Address, &'b GenesisAccount)>,
 ) -> ProviderResult<()> {
-    insert_state::<DB>(tx, capacity, alloc, 0)
+    insert_state::<DB>(tx, alloc, 0)
 }
 
 /// Inserts state at given block into database.
 pub fn insert_state<'a, 'b, DB: Database>(
     tx: &<DB as Database>::TXMut,
-    capacity: usize,
     alloc: impl Iterator<Item = (&'a Address, &'b GenesisAccount)>,
     block: u64,
 ) -> ProviderResult<()> {
-    let mut state_init: BundleStateInit = HashMap::with_capacity(capacity);
-    let mut reverts_init = HashMap::with_capacity(capacity);
-    let mut contracts: HashMap<B256, Bytecode> = HashMap::with_capacity(capacity);
+    let mut bundle_builder = BundleBuilder::default();
 
     for (address, account) in alloc {
         let bytecode_hash = if let Some(code) = &account.code {
             let bytecode = Bytecode::new_raw(code.clone());
             let hash = bytecode.hash_slow();
-            contracts.insert(hash, bytecode);
+            bundle_builder = bundle_builder.contract(hash, bytecode.into());
             Some(hash)
         } else {
             None
@@ -167,40 +163,31 @@ pub fn insert_state<'a, 'b, DB: Database>(
                 m.iter()
                     .map(|(key, value)| {
                         let value = U256::from_be_bytes(value.0);
-                        (*key, (U256::ZERO, value))
+                        ((*key).into(), (U256::ZERO, value))
                     })
                     .collect::<HashMap<_, _>>()
             })
             .unwrap_or_default();
 
-        reverts_init.insert(
-            *address,
-            (Some(None), storage.keys().map(|k| StorageEntry::new(*k, U256::ZERO)).collect()),
-        );
-
-        state_init.insert(
-            *address,
-            (
-                None,
-                Some(Account {
+        bundle_builder = bundle_builder
+            .revert_storage(
+                block,
+                *address,
+                storage.keys().map(|k| (*k, U256::ZERO)).collect::<Vec<(U256, U256)>>(),
+            )
+            .state_present_account_info(
+                *address,
+                into_revm_acc(Account {
                     nonce: account.nonce.unwrap_or_default(),
                     balance: account.balance,
                     bytecode_hash,
                 }),
-                storage,
-            ),
-        );
+            )
+            .state_storage(*address, storage);
     }
-    let all_reverts_init: RevertsInit = HashMap::from([(block, reverts_init)]);
 
-    let execution_outcome = ExecutionOutcome::new_init(
-        state_init,
-        all_reverts_init,
-        contracts.into_iter().collect(),
-        Receipts::default(),
-        block,
-        Vec::new(),
-    );
+    let execution_outcome =
+        ExecutionOutcome::new(bundle_builder.build(), Receipts::default(), block, Vec::new());
 
     execution_outcome.write_to_storage(tx, None, OriginalValuesKnown::Yes)?;
 
@@ -437,7 +424,6 @@ fn dump_state<DB: Database>(
             let tx = provider_rw.deref_mut().tx_mut();
             insert_state::<DB>(
                 tx,
-                accounts.len(),
                 accounts.iter().map(|(address, account)| (address, account)),
                 block,
             )?;

crates/storage/provider/src/bundle_state/mod.rs

@@ -5,9 +5,7 @@ mod hashed_state_changes;
 mod state_changes;
 mod state_reverts;
 
-pub use bundle_state_with_receipts::{
-    AccountRevertInit, BundleStateInit, ExecutionOutcome, OriginalValuesKnown, RevertsInit,
-};
+pub use bundle_state_with_receipts::{ExecutionOutcome, OriginalValuesKnown};
 pub use hashed_state_changes::HashedStateChanges;
 pub use state_changes::StateChanges;
 pub use state_reverts::{StateReverts, StorageRevertsIter};