refactor: use BundleBuilder as a state builder (#8861)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
Co-authored-by: Emilia Hane <elsaemiliaevahane@gmail.com>

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

@@ -4,7 +4,10 @@ use reth_primitives::{
 };
 use reth_trie::HashedPostState;
 use revm::{
-    db::{states::BundleState, BundleAccount},
+    db::{
+        states::{BundleBuilder, BundleState},
+        BundleAccount,
+    },
     primitives::AccountInfo,
 };
 
@@ -48,6 +51,19 @@ 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 from_state_builder(
+        state_builder: BundleBuilder,
+        receipts: Receipts,
+        first_block: BlockNumber,
+        requests: Vec<Requests>,
+    ) -> Self {
+        Self { bundle: state_builder.build(), receipts, first_block, requests }
+    }
+
     /// Return revm bundle state.
     pub const fn state(&self) -> &BundleState {
         &self.bundle
@@ -276,73 +292,9 @@ impl ExecutionOutcome {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use alloy_eips::{eip6110::DepositRequest, eip7002::WithdrawalRequest};
+    use alloy_eips::eip6110::DepositRequest;
     use alloy_primitives::{FixedBytes, LogData};
-    use reth_primitives::{Address, Receipts, Request, Requests, TxType, B256};
-    use std::collections::HashMap;
-
-    #[test]
-    fn test_initialisation() {
-        // Create a new BundleState object with initial data
-        let bundle = BundleState::new(
-            vec![(Address::new([2; 20]), None, Some(AccountInfo::default()), HashMap::default())],
-            vec![vec![(Address::new([2; 20]), None, vec![])]],
-            vec![],
-        );
-
-        // Create a Receipts object with a vector of receipt vectors
-        let receipts = Receipts {
-            receipt_vec: vec![vec![Some(Receipt {
-                tx_type: TxType::Legacy,
-                cumulative_gas_used: 46913,
-                logs: vec![],
-                success: true,
-                #[cfg(feature = "optimism")]
-                deposit_nonce: Some(18),
-                #[cfg(feature = "optimism")]
-                deposit_receipt_version: Some(34),
-            })]],
-        };
-
-        // Create a Requests object with a vector of requests, including DepositRequest and
-        // WithdrawalRequest
-        let requests = vec![Requests(vec![
-            Request::DepositRequest(DepositRequest {
-                pubkey: FixedBytes::<48>::from([1; 48]),
-                withdrawal_credentials: B256::from([0; 32]),
-                amount: 1111,
-                signature: FixedBytes::<96>::from([2; 96]),
-                index: 222,
-            }),
-            Request::DepositRequest(DepositRequest {
-                pubkey: FixedBytes::<48>::from([23; 48]),
-                withdrawal_credentials: B256::from([0; 32]),
-                amount: 34343,
-                signature: FixedBytes::<96>::from([43; 96]),
-                index: 1212,
-            }),
-            Request::WithdrawalRequest(WithdrawalRequest {
-                source_address: Address::from([1; 20]),
-                validator_pubkey: FixedBytes::<48>::from([10; 48]),
-                amount: 72,
-            }),
-        ])];
-
-        // Define the first block number
-        let first_block = 123;
-
-        // Create a ExecutionOutcome object with the created bundle, receipts, requests, and
-        // first_block
-        let exec_res = ExecutionOutcome {
-            bundle: bundle.clone(),
-            receipts: receipts.clone(),
-            requests: requests.clone(),
-            first_block,
-        };
-
-        // Assert that creating a new ExecutionOutcome using the constructor matches exec_res
-        assert_eq!(ExecutionOutcome::new(bundle, receipts, first_block, requests), exec_res);
-    }
+    use reth_primitives::{Receipts, Request, Requests, TxType, B256};
 
     #[test]
     fn test_block_number_to_index() {

crates/primitives-traits/src/account.rs

@@ -66,6 +66,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

@@ -19,7 +19,7 @@ reth-trie.workspace = true
 reth-etl.workspace = true
 reth-codecs.workspace = true
 reth-stages-types.workspace = true
-reth-revm.workspace = true
+revm.workspace = true
 reth-fs-util.workspace = true
 
 # eth

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

@@ -21,6 +21,7 @@ use reth_provider::{
 use reth_revm::db::states::BundleBuilder;
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_trie::{IntermediateStateRootState, StateRoot as StateRootComputer, StateRootProgress};
+use revm::db::states::BundleBuilder;
 use serde::{Deserialize, Serialize};
 use std::{
     collections::{BTreeMap, HashMap},
@@ -145,52 +146,61 @@ pub fn insert_state<'a, 'b, DB: Database>(
     alloc: impl Iterator<Item = (&'a Address, &'b GenesisAccount)>,
     block: u64,
 ) -> ProviderResult<()> {
+    // Create an empty bundle builder to accumulate the state changes
     let mut bundle_builder = BundleBuilder::default();
 
+    // Iterate over each address and its corresponding account
     for (address, account) in alloc {
+        // Check if the account has associated bytecode
         let bytecode_hash = if let Some(code) = &account.code {
+            // Create a new bytecode object from the code
             let bytecode = Bytecode::new_raw(code.clone());
+            // Compute the hash of the bytecode
             let hash = bytecode.hash_slow();
+            // Add the contract with its bytecode to the bundle builder
             bundle_builder = bundle_builder.contract(hash, bytecode.into());
+            // Store the hash for later use
             Some(hash)
         } else {
+            // No bytecode, so no hash
             None
         };
 
-        // get state
+        // Get the storage of the account, converting it to the required format
         let storage = account
             .storage
             .as_ref()
             .map(|m| {
                 m.iter()
-                    .map(|(key, value)| {
-                        let value = U256::from_be_bytes(value.0);
-                        ((*key).into(), (U256::ZERO, value))
-                    })
+                    .map(|(key, value)| ((*key).into(), (U256::ZERO, (*value).into())))
                     .collect::<HashMap<_, _>>()
             })
             .unwrap_or_default();
 
+        // Update the bundle builder with the account's state
         bundle_builder = bundle_builder
-            .revert_storage(
-                block,
-                *address,
-                storage.keys().map(|k| (*k, U256::ZERO)).collect::<Vec<(U256, U256)>>(),
-            )
+            .revert_storage(block, *address, storage.keys().map(|k| (*k, U256::ZERO)).collect())
             .state_present_account_info(
                 *address,
-                into_revm_acc(Account {
+                Account {
                     nonce: account.nonce.unwrap_or_default(),
                     balance: account.balance,
                     bytecode_hash,
-                }),
+                }
+                .into(),
             )
             .state_storage(*address, storage);
     }
 
-    let execution_outcome =
-        ExecutionOutcome::new(bundle_builder.build(), Receipts::default(), block, Vec::new());
+    // Create an execution outcome
+    let execution_outcome = ExecutionOutcome::from_state_builder(
+        bundle_builder,
+        Receipts::default(),
+        block,
+        Vec::new(),
+    );
 
+    // Write the execution outcome to the database storage
     execution_outcome.write_to_storage(tx, None, OriginalValuesKnown::Yes)?;
 
     trace!(target: "reth::cli", "Inserted state");

crates/storage/provider/src/providers/database/provider.rs

@@ -1,5 +1,5 @@
 use crate::{
-    bundle_state::{ExecutionOutcome, HashedStateChanges},
+    bundle_state::HashedStateChanges,
     providers::{database::metrics, static_file::StaticFileWriter, StaticFileProvider},
     to_range,
     traits::{
@@ -46,7 +46,10 @@ use reth_trie::{
     updates::TrieUpdates,
     HashedPostState, Nibbles, StateRoot,
 };
-use revm::{db::states::BundleBuilder, primitives::{BlockEnv, CfgEnvWithHandlerCfg}};
+use revm::{
+    db::states::BundleBuilder,
+    primitives::{BlockEnv, CfgEnvWithHandlerCfg},
+};
 use std::{
     cmp::Ordering,
     collections::{hash_map, BTreeMap, BTreeSet, HashMap, HashSet},
@@ -615,6 +618,7 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
         }
         let start_block_number = *range.start();
 
+        // Fetch the block bodies within the specified range.
         // We are not removing block meta as it is used to get block changesets.
         let block_bodies = self.get_or_take::<tables::BlockBodyIndices, false>(range.clone())?;
 
@@ -635,6 +639,7 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
         // Double option around Account represent if Account state is know (first option) and
         // account is removed (Second Option)
 
+        // Create an empty bundle builder to accumulate the state changes
         let mut bundle_builder = BundleBuilder::new(range);
 
         // This is not working for blocks that are not at tip. as plain state is not the last
@@ -644,48 +649,58 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
         let mut plain_accounts_cursor = self.tx.cursor_write::<tables::PlainAccountState>()?;
         let mut plain_storage_cursor = self.tx.cursor_dup_write::<tables::PlainStorageState>()?;
 
-        // add account changeset changes
+        // Iterate over the account changesets in reverse order.
         for (block_number, account_before) in account_changeset.into_iter().rev() {
+            // Destructure the account changeset to get the old info and address.
             let AccountBeforeTx { info: old_info, address } = account_before;
+
+            // If the address is not already in the bundle builder's state, add it.
             if !bundle_builder.get_states().contains(&address) {
-                let new_info = plain_accounts_cursor.seek_exact(address)?.map(|kv| kv.1);
-                if let Some(new_info) = new_info {
+                // Seek the address in the plain account state.
+                if let Some(new_info) = plain_accounts_cursor.seek_exact(address)?.map(|kv| kv.1) {
+                    // Add the new info to the bundle builder.
                     bundle_builder =
-                        bundle_builder.state_present_account_info(address, into_revm_acc(new_info));
+                        bundle_builder.state_present_account_info(address, new_info.into());
                 }
             }
+
+            // Add the old account info to the bundle builder if it exists.
             if let Some(old_info) = old_info {
                 bundle_builder =
-                    bundle_builder.state_original_account_info(address, into_revm_acc(old_info));
+                    bundle_builder.state_original_account_info(address, old_info.into());
             }
 
-            // insert old info into reverts.
+            // Insert the old account info into the reverts.
             bundle_builder = bundle_builder.revert_account_info(
                 block_number,
                 address,
-                Some(old_info.map(into_revm_acc)),
+                Some(old_info.map(Into::into)),
             );
         }
 
+        let mut state_storage: HashMap<Address, HashMap<U256, (U256, U256)>> = HashMap::new();
+
         // add storage changeset changes
         for (block_and_address, old_storage) in storage_changeset.into_iter().rev() {
             let BlockNumberAddress((block_number, address)) = block_and_address;
-            // get account state or insert from plain state.
+            // If the address is not already in the bundle builder's state, add it.
             if !bundle_builder.get_states().contains(&address) {
-                let present_info = plain_accounts_cursor.seek_exact(address)?.map(|kv| kv.1);
-                if let Some(present_info) = present_info {
+                if let Some(present_info) =
+                    plain_accounts_cursor.seek_exact(address)?.map(|kv| kv.1)
+                {
                     bundle_builder = bundle_builder
-                        .state_original_account_info(address, into_revm_acc(present_info))
-                        .state_present_account_info(address, into_revm_acc(present_info));
+                        .state_original_account_info(address, present_info.into())
+                        .state_present_account_info(address, present_info.into());
                 }
             }
 
-            let account_state = match bundle_builder.get_state_storage_mut().entry(address) {
+            // Get or insert the account state from plain state.
+            let account_state = match state_storage.entry(address) {
                 hash_map::Entry::Occupied(entry) => entry.into_mut(),
                 hash_map::Entry::Vacant(entry) => entry.insert(HashMap::new()),
             };
 
-            // match storage.
+            // Match storage changes and add to the account state.
             match account_state.entry(old_storage.key.into()) {
                 hash_map::Entry::Vacant(entry) => {
                     let new_storage = plain_storage_cursor
@@ -699,24 +714,27 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
                 }
             };
 
-            bundle_builder
-                .get_revert_storage_mut()
-                .entry((block_number, address))
-                .or_default()
-                .push((old_storage.key.into(), old_storage.value));
+            // Add the current state of storage to the bundle builder.
+            bundle_builder = bundle_builder
+                .state_storage(address, state_storage.entry(address).or_default().clone())
+                .revert_storage(
+                    block_number,
+                    address,
+                    // revert_storage.entry(block_and_address).or_default().clone(),
+                    vec![(old_storage.key.into(), old_storage.value)],
+                );
         }
 
         let bundle_state = bundle_builder.build();
 
         if TAKE {
             // iterate over local plain state remove all account and all storages.
-
             for (address, bundle_account) in &bundle_state.state {
                 // revert account if needed.
                 if bundle_account.is_info_changed() {
                     let existing_entry = plain_accounts_cursor.seek_exact(*address)?;
                     if let Some(account) = &bundle_account.original_info {
-                        plain_accounts_cursor.upsert(*address, into_reth_acc(account.clone()))?;
+                        plain_accounts_cursor.upsert(*address, account.clone().into())?;
                     } else if existing_entry.is_some() {
                         plain_accounts_cursor.delete_current()?;
                     }
@@ -738,7 +756,7 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
                     }
 
                     // insert value if needed
-                    if storage_slot.original_value() != U256::ZERO {
+                    if !storage_slot.original_value().is_zero() {
                         plain_storage_cursor.upsert(*address, storage_entry)?;
                     }
                 }
@@ -762,12 +780,7 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
             receipts.push(block_receipts);
         }
 
-        Ok(ExecutionOutcome::new(
-            bundle_state,
-            reth_primitives::Receipts::from(receipts),
-            start_block_number,
-            Vec::new(),
-        ))
+        Ok(ExecutionOutcome::new(bundle_state, receipts.into(), start_block_number, Vec::new()))
     }
 
     /// Return list of entries from table