feat(provider): StateRootProvider::state_root_with_updates (#5485)

crates/blockchain-tree/src/blockchain_tree.rs

@@ -104,8 +104,8 @@ impl<DB: Database, EF: ExecutorFactory> BlockchainTree<DB, EF> {
             externals.fetch_latest_canonical_hashes(config.num_of_canonical_hashes() as usize)?;
 
         // TODO(rakita) save last finalized block inside database but for now just take
-        // tip-max_reorg_depth
-        // task: https://github.com/paradigmxyz/reth/issues/1712
+        // `tip - max_reorg_depth`
+        // https://github.com/paradigmxyz/reth/issues/1712
         let last_finalized_block_number = if last_canonical_hashes.len() > max_reorg_depth {
             // we pick `Highest - max_reorg_depth` block as last finalized block.
             last_canonical_hashes.keys().nth_back(max_reorg_depth)

crates/blockchain-tree/src/chain.rs

@@ -158,15 +158,19 @@ impl AppendableChain {
         state.revert_to(parent.number);
 
         // Revert changesets to get the state of the parent that we need to apply the change.
-        let post_state_data = BundleStateDataRef {
+        let bundle_state_data = BundleStateDataRef {
             state: &state,
             sidechain_block_hashes: &side_chain_block_hashes,
             canonical_block_hashes,
             canonical_fork,
         };
-        let block_state =
-            Self::validate_and_execute_sidechain(block.clone(), parent, post_state_data, externals)
-                .map_err(|err| InsertBlockError::new(block.block.clone(), err.into()))?;
+        let block_state = Self::validate_and_execute_sidechain(
+            block.clone(),
+            parent,
+            bundle_state_data,
+            externals,
+        )
+        .map_err(|err| InsertBlockError::new(block.block.clone(), err.into()))?;
         state.extend(block_state);
 
         // If all is okay, return new chain back. Present chain is not modified.
@@ -185,7 +189,7 @@ impl AppendableChain {
     fn validate_and_execute<BSDP, DB, EF>(
         block: SealedBlockWithSenders,
         parent_block: &SealedHeader,
-        post_state_data_provider: BSDP,
+        bundle_state_data_provider: BSDP,
         externals: &TreeExternals<DB, EF>,
         block_kind: BlockKind,
         block_validation_kind: BlockValidationKind,
@@ -203,10 +207,10 @@ impl AppendableChain {
 
         // get the state provider.
         let db = externals.database();
-        let canonical_fork = post_state_data_provider.canonical_fork();
+        let canonical_fork = bundle_state_data_provider.canonical_fork();
         let state_provider = db.history_by_block_number(canonical_fork.number)?;
 
-        let provider = BundleStateProvider::new(state_provider, post_state_data_provider);
+        let provider = BundleStateProvider::new(state_provider, bundle_state_data_provider);
 
         let mut executor = externals.executor_factory.with_state(&provider);
         executor.execute_and_verify_receipt(&block, U256::MAX, Some(senders))?;
@@ -232,7 +236,7 @@ impl AppendableChain {
     fn validate_and_execute_sidechain<BSDP, DB, EF>(
         block: SealedBlockWithSenders,
         parent_block: &SealedHeader,
-        post_state_data_provider: BSDP,
+        bundle_state_data_provider: BSDP,
         externals: &TreeExternals<DB, EF>,
     ) -> RethResult<BundleStateWithReceipts>
     where
@@ -243,7 +247,7 @@ impl AppendableChain {
         Self::validate_and_execute(
             block,
             parent_block,
-            post_state_data_provider,
+            bundle_state_data_provider,
             externals,
             BlockKind::ForksHistoricalBlock,
             BlockValidationKind::SkipStateRootValidation,
@@ -279,7 +283,7 @@ impl AppendableChain {
     {
         let parent_block = self.chain.tip();
 
-        let post_state_data = BundleStateDataRef {
+        let bundle_state_data = BundleStateDataRef {
             state: self.state(),
             sidechain_block_hashes: &side_chain_block_hashes,
             canonical_block_hashes,
@@ -289,7 +293,7 @@ impl AppendableChain {
         let block_state = Self::validate_and_execute(
             block.clone(),
             parent_block,
-            post_state_data,
+            bundle_state_data,
             externals,
             block_kind,
             block_validation_kind,

crates/revm/Cargo.toml

@@ -22,6 +22,9 @@ revm.workspace = true
 # common
 tracing.workspace = true
 
+[dev-dependencies]
+reth-trie.workspace = true
+
 [features]
 optimism = [
   "revm/optimism",

crates/revm/src/processor.rs

@@ -568,6 +568,7 @@ mod tests {
     use reth_provider::{
         AccountReader, BlockHashReader, BundleStateWithReceipts, StateRootProvider,
     };
+    use reth_trie::updates::TrieUpdates;
     use revm::{Database, TransitionState};
     use std::collections::HashMap;
 
@@ -627,6 +628,13 @@ mod tests {
         fn state_root(&self, _bundle_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
             unimplemented!("state root computation is not supported")
         }
+
+        fn state_root_with_updates(
+            &self,
+            _bundle_state: &BundleStateWithReceipts,
+        ) -> ProviderResult<(B256, TrieUpdates)> {
+            unimplemented!("state root computation is not supported")
+        }
     }
 
     impl StateProvider for StateProviderTest {

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

@@ -13,6 +13,7 @@ use reth_primitives::{
 };
 use reth_trie::{
     hashed_cursor::{HashedPostState, HashedPostStateCursorFactory, HashedStorage},
+    updates::TrieUpdates,
     StateRoot, StateRootError,
 };
 use revm::{db::states::BundleState, primitives::AccountInfo};
@@ -154,6 +155,20 @@ impl BundleStateWithReceipts {
         hashed_state.sorted()
     }
 
+    /// Returns [StateRoot] calculator.
+    fn state_root_calculator<'a, 'b, TX: DbTx>(
+        &self,
+        tx: &'a TX,
+        hashed_post_state: &'b HashedPostState,
+    ) -> StateRoot<'a, TX, HashedPostStateCursorFactory<'a, 'b, TX>> {
+        let (account_prefix_set, storage_prefix_set) = hashed_post_state.construct_prefix_sets();
+        let hashed_cursor_factory = HashedPostStateCursorFactory::new(tx, hashed_post_state);
+        StateRoot::new(tx)
+            .with_hashed_cursor_factory(hashed_cursor_factory)
+            .with_changed_account_prefixes(account_prefix_set)
+            .with_changed_storage_prefixes(storage_prefix_set)
+    }
+
     /// Calculate the state root for this [BundleState].
     /// Internally, function calls [Self::hash_state_slow] to obtain the [HashedPostState].
     /// Afterwards, it retrieves the prefixsets from the [HashedPostState] and uses them to
@@ -196,13 +211,17 @@ impl BundleStateWithReceipts {
     /// The state root for this [BundleState].
     pub fn state_root_slow<TX: DbTx>(&self, tx: &TX) -> Result<B256, StateRootError> {
         let hashed_post_state = self.hash_state_slow();
-        let (account_prefix_set, storage_prefix_set) = hashed_post_state.construct_prefix_sets();
-        let hashed_cursor_factory = HashedPostStateCursorFactory::new(tx, &hashed_post_state);
-        StateRoot::new(tx)
-            .with_hashed_cursor_factory(hashed_cursor_factory)
-            .with_changed_account_prefixes(account_prefix_set)
-            .with_changed_storage_prefixes(storage_prefix_set)
-            .root()
+        self.state_root_calculator(tx, &hashed_post_state).root()
+    }
+
+    /// Calculates the state root for this [BundleState] and returns it alongside trie updates.
+    /// See [Self::state_root_slow] for more info.
+    pub fn state_root_slow_with_updates<TX: DbTx>(
+        &self,
+        tx: &TX,
+    ) -> Result<(B256, TrieUpdates), StateRootError> {
+        let hashed_post_state = self.hash_state_slow();
+        self.state_root_calculator(tx, &hashed_post_state).root_with_updates()
     }
 
     /// Transform block number to the index of block.

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

@@ -19,7 +19,7 @@ impl From<StateChangeset> for StateChanges {
 }
 
 impl StateChanges {
-    /// Write the post state to the database.
+    /// Write the bundle state to the database.
     pub fn write_to_db<TX: DbTxMut + DbTx>(mut self, tx: &TX) -> Result<(), DatabaseError> {
         // sort all entries so they can be written to database in more performant way.
         // and take smaller memory footprint.
@@ -28,28 +28,28 @@ impl StateChanges {
         self.0.contracts.par_sort_by_key(|a| a.0);
 
         // Write new account state
-        tracing::trace!(target: "provider::post_state", len = self.0.accounts.len(), "Writing new account state");
+        tracing::trace!(target: "provider::bundle_state", len = self.0.accounts.len(), "Writing new account state");
         let mut accounts_cursor = tx.cursor_write::<tables::PlainAccountState>()?;
         // write account to database.
         for (address, account) in self.0.accounts.into_iter() {
             if let Some(account) = account {
-                tracing::trace!(target: "provider::post_state", ?address, "Updating plain state account");
+                tracing::trace!(target: "provider::bundle_state", ?address, "Updating plain state account");
                 accounts_cursor.upsert(address, into_reth_acc(account))?;
             } else if accounts_cursor.seek_exact(address)?.is_some() {
-                tracing::trace!(target: "provider::post_state", ?address, "Deleting plain state account");
+                tracing::trace!(target: "provider::bundle_state", ?address, "Deleting plain state account");
                 accounts_cursor.delete_current()?;
             }
         }
 
         // Write bytecode
-        tracing::trace!(target: "provider::post_state", len = self.0.contracts.len(), "Writing bytecodes");
+        tracing::trace!(target: "provider::bundle_state", len = self.0.contracts.len(), "Writing bytecodes");
         let mut bytecodes_cursor = tx.cursor_write::<tables::Bytecodes>()?;
         for (hash, bytecode) in self.0.contracts.into_iter() {
             bytecodes_cursor.upsert(hash, Bytecode(bytecode))?;
         }
 
         // Write new storage state and wipe storage if needed.
-        tracing::trace!(target: "provider::post_state", len = self.0.storage.len(), "Writing new storage state");
+        tracing::trace!(target: "provider::bundle_state", len = self.0.storage.len(), "Writing new storage state");
         let mut storages_cursor = tx.cursor_dup_write::<tables::PlainStorageState>()?;
         for PlainStorageChangeset { address, wipe_storage, storage } in self.0.storage.into_iter() {
             // Wiping of storage.
@@ -65,7 +65,7 @@ impl StateChanges {
             storage.par_sort_unstable_by_key(|a| a.key);
 
             for entry in storage.into_iter() {
-                tracing::trace!(target: "provider::post_state", ?address, ?entry.key, "Updating plain state storage");
+                tracing::trace!(target: "provider::bundle_state", ?address, ?entry.key, "Updating plain state storage");
                 if let Some(db_entry) = storages_cursor.seek_by_key_subkey(address, entry.key)? {
                     if db_entry.key == entry.key {
                         storages_cursor.delete_current()?;

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

@@ -4,21 +4,22 @@ use crate::{
 };
 use reth_interfaces::provider::{ProviderError, ProviderResult};
 use reth_primitives::{trie::AccountProof, Account, Address, BlockNumber, Bytecode, B256};
+use reth_trie::updates::TrieUpdates;
 
 /// A state provider that either resolves to data in a wrapped [`crate::BundleStateWithReceipts`],
 /// or an underlying state provider.
 #[derive(Debug)]
 pub struct BundleStateProvider<SP: StateProvider, BSDP: BundleStateDataProvider> {
     /// The inner state provider.
     pub(crate) state_provider: SP,
-    /// Post state data,
-    pub(crate) post_state_data_provider: BSDP,
+    /// Bundle state data,
+    pub(crate) bundle_state_data_provider: BSDP,
 }
 
 impl<SP: StateProvider, BSDP: BundleStateDataProvider> BundleStateProvider<SP, BSDP> {
-    /// Create new post-state provider
-    pub fn new(state_provider: SP, post_state_data_provider: BSDP) -> Self {
-        Self { state_provider, post_state_data_provider }
+    /// Create new bundle state provider
+    pub fn new(state_provider: SP, bundle_state_data_provider: BSDP) -> Self {
+        Self { state_provider, bundle_state_data_provider }
     }
 }
 
@@ -28,7 +29,7 @@ impl<SP: StateProvider, BSDP: BundleStateDataProvider> BlockHashReader
     for BundleStateProvider<SP, BSDP>
 {
     fn block_hash(&self, block_number: BlockNumber) -> ProviderResult<Option<B256>> {
-        let block_hash = self.post_state_data_provider.block_hash(block_number);
+        let block_hash = self.bundle_state_data_provider.block_hash(block_number);
         if block_hash.is_some() {
             return Ok(block_hash)
         }
@@ -48,7 +49,7 @@ impl<SP: StateProvider, BSDP: BundleStateDataProvider> AccountReader
     for BundleStateProvider<SP, BSDP>
 {
     fn basic_account(&self, address: Address) -> ProviderResult<Option<Account>> {
-        if let Some(account) = self.post_state_data_provider.state().account(&address) {
+        if let Some(account) = self.bundle_state_data_provider.state().account(&address) {
             Ok(account)
         } else {
             self.state_provider.basic_account(address)
@@ -59,11 +60,20 @@ impl<SP: StateProvider, BSDP: BundleStateDataProvider> AccountReader
 impl<SP: StateProvider, BSDP: BundleStateDataProvider> StateRootProvider
     for BundleStateProvider<SP, BSDP>
 {
-    fn state_root(&self, post_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
-        let mut state = self.post_state_data_provider.state().clone();
-        state.extend(post_state.clone());
+    fn state_root(&self, bundle_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
+        let mut state = self.bundle_state_data_provider.state().clone();
+        state.extend(bundle_state.clone());
         self.state_provider.state_root(&state)
     }
+
+    fn state_root_with_updates(
+        &self,
+        bundle_state: &BundleStateWithReceipts,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let mut state = self.bundle_state_data_provider.state().clone();
+        state.extend(bundle_state.clone());
+        self.state_provider.state_root_with_updates(&state)
+    }
 }
 
 impl<SP: StateProvider, BSDP: BundleStateDataProvider> StateProvider
@@ -76,7 +86,7 @@ impl<SP: StateProvider, BSDP: BundleStateDataProvider> StateProvider
     ) -> ProviderResult<Option<reth_primitives::StorageValue>> {
         let u256_storage_key = storage_key.into();
         if let Some(value) =
-            self.post_state_data_provider.state().storage(&account, u256_storage_key)
+            self.bundle_state_data_provider.state().storage(&account, u256_storage_key)
         {
             return Ok(Some(value))
         }
@@ -85,7 +95,7 @@ impl<SP: StateProvider, BSDP: BundleStateDataProvider> StateProvider
     }
 
     fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
-        if let Some(bytecode) = self.post_state_data_provider.state().bytecode(&code_hash) {
+        if let Some(bytecode) = self.bundle_state_data_provider.state().bytecode(&code_hash) {
             return Ok(Some(bytecode))
         }
 

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

@@ -571,14 +571,14 @@ where
 
     fn pending_with_provider(
         &self,
-        post_state_data: Box<dyn BundleStateDataProvider>,
+        bundle_state_data: Box<dyn BundleStateDataProvider>,
     ) -> ProviderResult<StateProviderBox<'_>> {
-        let canonical_fork = post_state_data.canonical_fork();
+        let canonical_fork = bundle_state_data.canonical_fork();
         trace!(target: "providers::blockchain", ?canonical_fork, "Returning post state provider");
 
         let state_provider = self.history_by_block_hash(canonical_fork.hash)?;
-        let post_state_provider = BundleStateProvider::new(state_provider, post_state_data);
-        Ok(Box::new(post_state_provider))
+        let bundle_state_provider = BundleStateProvider::new(state_provider, bundle_state_data);
+        Ok(Box::new(bundle_state_provider))
     }
 }
 

crates/storage/provider/src/providers/state/historical.rs

@@ -14,6 +14,7 @@ use reth_interfaces::provider::ProviderResult;
 use reth_primitives::{
     trie::AccountProof, Account, Address, BlockNumber, Bytecode, StorageKey, StorageValue, B256,
 };
+use reth_trie::updates::TrieUpdates;
 
 /// State provider for a given block number which takes a tx reference.
 ///
@@ -198,7 +199,14 @@ impl<'b, TX: DbTx> BlockHashReader for HistoricalStateProviderRef<'b, TX> {
 }
 
 impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
-    fn state_root(&self, _post_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
+    fn state_root(&self, _bundle_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
+        Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
+    }
+
+    fn state_root_with_updates(
+        &self,
+        _bundle_state: &BundleStateWithReceipts,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
         Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
     }
 }

crates/storage/provider/src/providers/state/latest.rs

@@ -12,6 +12,7 @@ use reth_primitives::{
     keccak256, trie::AccountProof, Account, Address, BlockNumber, Bytecode, StorageKey,
     StorageValue, B256,
 };
+use reth_trie::updates::TrieUpdates;
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -62,6 +63,15 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
     fn state_root(&self, bundle_state: &BundleStateWithReceipts) -> ProviderResult<B256> {
         bundle_state.state_root_slow(self.db).map_err(|err| ProviderError::Database(err.into()))
     }
+
+    fn state_root_with_updates(
+        &self,
+        bundle_state: &BundleStateWithReceipts,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        bundle_state
+            .state_root_slow_with_updates(self.db)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
 }
 
 impl<'b, TX: DbTx> StateProvider for LatestStateProviderRef<'b, TX> {