feat(trie): Merge trie changesets changes into main

crates/chain-state/benches/canonical_hashes_range.rs

@@ -2,7 +2,7 @@
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use reth_chain_state::{
-    test_utils::TestBlockBuilder, ExecutedBlockWithTrieUpdates, MemoryOverlayStateProviderRef,
+    test_utils::TestBlockBuilder, ExecutedBlock, MemoryOverlayStateProviderRef,
 };
 use reth_ethereum_primitives::EthPrimitives;
 use reth_storage_api::{noop::NoopProvider, BlockHashReader};
@@ -84,10 +84,7 @@ fn bench_canonical_hashes_range(c: &mut Criterion) {
 
 fn setup_provider_with_blocks(
     num_blocks: usize,
-) -> (
-    MemoryOverlayStateProviderRef<'static, EthPrimitives>,
-    Vec<ExecutedBlockWithTrieUpdates<EthPrimitives>>,
-) {
+) -> (MemoryOverlayStateProviderRef<'static, EthPrimitives>, Vec<ExecutedBlock<EthPrimitives>>) {
     let mut builder = TestBlockBuilder::<EthPrimitives>::default();
 
     let blocks: Vec<_> = builder.get_executed_blocks(1000..1000 + num_blocks as u64).collect();

crates/chain-state/src/in_memory.rs

@@ -242,7 +242,7 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
     /// Updates the pending block with the given block.
     ///
     /// Note: This assumes that the parent block of the pending block is canonical.
-    pub fn set_pending_block(&self, pending: ExecutedBlockWithTrieUpdates<N>) {
+    pub fn set_pending_block(&self, pending: ExecutedBlock<N>) {
         // fetch the state of the pending block's parent block
         let parent = self.state_by_hash(pending.recovered_block().parent_hash());
         let pending = BlockState::with_parent(pending, parent);
@@ -258,7 +258,7 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
     /// them to their parent blocks.
     fn update_blocks<I, R>(&self, new_blocks: I, reorged: R)
     where
-        I: IntoIterator<Item = ExecutedBlockWithTrieUpdates<N>>,
+        I: IntoIterator<Item = ExecutedBlock<N>>,
         R: IntoIterator<Item = ExecutedBlock<N>>,
     {
         {
@@ -568,22 +568,19 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct BlockState<N: NodePrimitives = EthPrimitives> {
     /// The executed block that determines the state after this block has been executed.
-    block: ExecutedBlockWithTrieUpdates<N>,
+    block: ExecutedBlock<N>,
     /// The block's parent block if it exists.
     parent: Option<Arc<BlockState<N>>>,
 }
 
 impl<N: NodePrimitives> BlockState<N> {
     /// [`BlockState`] constructor.
-    pub const fn new(block: ExecutedBlockWithTrieUpdates<N>) -> Self {
+    pub const fn new(block: ExecutedBlock<N>) -> Self {
         Self { block, parent: None }
     }
 
     /// [`BlockState`] constructor with parent.
-    pub const fn with_parent(
-        block: ExecutedBlockWithTrieUpdates<N>,
-        parent: Option<Arc<Self>>,
-    ) -> Self {
+    pub const fn with_parent(block: ExecutedBlock<N>, parent: Option<Arc<Self>>) -> Self {
         Self { block, parent }
     }
 
@@ -597,12 +594,12 @@ impl<N: NodePrimitives> BlockState<N> {
     }
 
     /// Returns the executed block that determines the state.
-    pub fn block(&self) -> ExecutedBlockWithTrieUpdates<N> {
+    pub fn block(&self) -> ExecutedBlock<N> {
         self.block.clone()
     }
 
     /// Returns a reference to the executed block that determines the state.
-    pub const fn block_ref(&self) -> &ExecutedBlockWithTrieUpdates<N> {
+    pub const fn block_ref(&self) -> &ExecutedBlock<N> {
         &self.block
     }
 
@@ -730,6 +727,8 @@ pub struct ExecutedBlock<N: NodePrimitives = EthPrimitives> {
     pub execution_output: Arc<ExecutionOutcome<N::Receipt>>,
     /// Block's hashed state.
     pub hashed_state: Arc<HashedPostState>,
+    /// Trie updates that result from calculating the state root for the block.
+    pub trie_updates: Arc<TrieUpdates>,
 }
 
 impl<N: NodePrimitives> Default for ExecutedBlock<N> {
@@ -738,6 +737,7 @@ impl<N: NodePrimitives> Default for ExecutedBlock<N> {
             recovered_block: Default::default(),
             execution_output: Default::default(),
             hashed_state: Default::default(),
+            trie_updates: Default::default(),
         }
     }
 }
@@ -767,113 +767,16 @@ impl<N: NodePrimitives> ExecutedBlock<N> {
         &self.hashed_state
     }
 
-    /// Returns a [`BlockNumber`] of the block.
+    /// Returns a reference to the trie updates resulting from the execution outcome
     #[inline]
-    pub fn block_number(&self) -> BlockNumber {
-        self.recovered_block.header().number()
+    pub fn trie_updates(&self) -> &TrieUpdates {
+        &self.trie_updates
     }
-}
 
-/// Trie updates that result from calculating the state root for the block.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum ExecutedTrieUpdates {
-    /// Trie updates present. State root was calculated, and the trie updates can be applied to the
-    /// database.
-    Present(Arc<TrieUpdates>),
-    /// Trie updates missing. State root was calculated, but the trie updates cannot be applied to
-    /// the current database state. To apply the updates, the state root must be recalculated, and
-    /// new trie updates must be generated.
-    ///
-    /// This can happen when processing fork chain blocks that are building on top of the
-    /// historical database state. Since we don't store the historical trie state, we cannot
-    /// generate the trie updates for it.
-    Missing,
-}
-
-impl ExecutedTrieUpdates {
-    /// Creates a [`ExecutedTrieUpdates`] with present but empty trie updates.
-    pub fn empty() -> Self {
-        Self::Present(Arc::default())
-    }
-
-    /// Sets the trie updates to the provided value as present.
-    pub fn set_present(&mut self, updates: Arc<TrieUpdates>) {
-        *self = Self::Present(updates);
-    }
-
-    /// Takes the present trie updates, leaving the state as missing.
-    pub fn take_present(&mut self) -> Option<Arc<TrieUpdates>> {
-        match self {
-            Self::Present(updates) => {
-                let updates = core::mem::take(updates);
-                *self = Self::Missing;
-                Some(updates)
-            }
-            Self::Missing => None,
-        }
-    }
-
-    /// Returns a reference to the trie updates if present.
-    #[allow(clippy::missing_const_for_fn)] // false positive
-    pub fn as_ref(&self) -> Option<&TrieUpdates> {
-        match self {
-            Self::Present(updates) => Some(updates),
-            Self::Missing => None,
-        }
-    }
-
-    /// Returns `true` if the trie updates are present.
-    pub const fn is_present(&self) -> bool {
-        matches!(self, Self::Present(_))
-    }
-
-    /// Returns `true` if the trie updates are missing.
-    pub const fn is_missing(&self) -> bool {
-        matches!(self, Self::Missing)
-    }
-}
-
-/// An [`ExecutedBlock`] with its [`TrieUpdates`].
-///
-/// We store it as separate type because [`TrieUpdates`] are only available for blocks stored in
-/// memory and can't be obtained for canonical persisted blocks.
-#[derive(
-    Clone, Debug, PartialEq, Eq, derive_more::Deref, derive_more::DerefMut, derive_more::Into,
-)]
-pub struct ExecutedBlockWithTrieUpdates<N: NodePrimitives = EthPrimitives> {
-    /// Inner [`ExecutedBlock`].
-    #[deref]
-    #[deref_mut]
-    #[into]
-    pub block: ExecutedBlock<N>,
-    /// Trie updates that result from calculating the state root for the block.
-    ///
-    /// If [`ExecutedTrieUpdates::Missing`], the trie updates should be computed when persisting
-    /// the block **on top of the canonical parent**.
-    pub trie: ExecutedTrieUpdates,
-}
-
-impl<N: NodePrimitives> ExecutedBlockWithTrieUpdates<N> {
-    /// [`ExecutedBlock`] constructor.
-    pub const fn new(
-        recovered_block: Arc<RecoveredBlock<N::Block>>,
-        execution_output: Arc<ExecutionOutcome<N::Receipt>>,
-        hashed_state: Arc<HashedPostState>,
-        trie: ExecutedTrieUpdates,
-    ) -> Self {
-        Self { block: ExecutedBlock { recovered_block, execution_output, hashed_state }, trie }
-    }
-
-    /// Returns a reference to the trie updates for the block, if present.
+    /// Returns a [`BlockNumber`] of the block.
     #[inline]
-    pub fn trie_updates(&self) -> Option<&TrieUpdates> {
-        self.trie.as_ref()
-    }
-
-    /// Converts the value into [`SealedBlock`].
-    pub fn into_sealed_block(self) -> SealedBlock<N::Block> {
-        let block = Arc::unwrap_or_clone(self.block.recovered_block);
-        block.into_sealed_block()
+    pub fn block_number(&self) -> BlockNumber {
+        self.recovered_block.header().number()
     }
 }
 
@@ -883,18 +786,14 @@ pub enum NewCanonicalChain<N: NodePrimitives = EthPrimitives> {
     /// A simple append to the current canonical head
     Commit {
         /// all blocks that lead back to the canonical head
-        new: Vec<ExecutedBlockWithTrieUpdates<N>>,
+        new: Vec<ExecutedBlock<N>>,
     },
     /// A reorged chain consists of two chains that trace back to a shared ancestor block at which
     /// point they diverge.
     Reorg {
         /// All blocks of the _new_ chain
-        new: Vec<ExecutedBlockWithTrieUpdates<N>>,
+        new: Vec<ExecutedBlock<N>>,
         /// All blocks of the _old_ chain
-        ///
-        /// These are not [`ExecutedBlockWithTrieUpdates`] because we don't always have the trie
-        /// updates for the old canonical chain. For example, in case of node being restarted right
-        /// before the reorg [`TrieUpdates`] can't be fetched from database.
         old: Vec<ExecutedBlock<N>>,
     },
 }
@@ -1257,7 +1156,7 @@ mod tests {
             block1.recovered_block().hash()
         );
 
-        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1.block] };
+        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1] };
         state.update_chain(chain);
         assert_eq!(
             state.head_state().unwrap().block_ref().recovered_block().hash(),
@@ -1539,7 +1438,7 @@ mod tests {
         // Test reorg notification
         let chain_reorg = NewCanonicalChain::Reorg {
             new: vec![block1a.clone(), block2a.clone()],
-            old: vec![block1.block.clone(), block2.block.clone()],
+            old: vec![block1.clone(), block2.clone()],
         };
 
         assert_eq!(

crates/chain-state/src/memory_overlay.rs

@@ -1,4 +1,4 @@
-use super::ExecutedBlockWithTrieUpdates;
+use super::ExecutedBlock;
 use alloy_consensus::BlockHeader;
 use alloy_primitives::{keccak256, Address, BlockNumber, Bytes, StorageKey, StorageValue, B256};
 use reth_errors::ProviderResult;
@@ -24,7 +24,7 @@ pub struct MemoryOverlayStateProviderRef<
     /// Historical state provider for state lookups that are not found in memory blocks.
     pub(crate) historical: Box<dyn StateProvider + 'a>,
     /// The collection of executed parent blocks. Expected order is newest to oldest.
-    pub(crate) in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
+    pub(crate) in_memory: Vec<ExecutedBlock<N>>,
     /// Lazy-loaded in-memory trie data.
     pub(crate) trie_input: OnceLock<TrieInput>,
 }
@@ -41,10 +41,7 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
     /// - `in_memory` - the collection of executed ancestor blocks in reverse.
     /// - `historical` - a historical state provider for the latest ancestor block stored in the
     ///   database.
-    pub fn new(
-        historical: Box<dyn StateProvider + 'a>,
-        in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
-    ) -> Self {
+    pub fn new(historical: Box<dyn StateProvider + 'a>, in_memory: Vec<ExecutedBlock<N>>) -> Self {
         Self { historical, in_memory, trie_input: OnceLock::new() }
     }
 
@@ -60,7 +57,7 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
                 self.in_memory
                     .iter()
                     .rev()
-                    .map(|block| (block.hashed_state.as_ref(), block.trie.as_ref())),
+                    .map(|block| (block.hashed_state.as_ref(), block.trie_updates.as_ref())),
             )
         })
     }

crates/chain-state/src/test_utils.rs

@@ -1,6 +1,6 @@
 use crate::{
-    in_memory::ExecutedBlockWithTrieUpdates, CanonStateNotification, CanonStateNotifications,
-    CanonStateSubscriptions, ExecutedTrieUpdates,
+    in_memory::ExecutedBlock, CanonStateNotification, CanonStateNotifications,
+    CanonStateSubscriptions,
 };
 use alloy_consensus::{Header, SignableTransaction, TxEip1559, TxReceipt, EMPTY_ROOT_HASH};
 use alloy_eips::{
@@ -23,7 +23,7 @@ use reth_primitives_traits::{
     SignedTransaction,
 };
 use reth_storage_api::NodePrimitivesProvider;
-use reth_trie::{root::state_root_unhashed, HashedPostState};
+use reth_trie::{root::state_root_unhashed, updates::TrieUpdates, HashedPostState};
 use revm_database::BundleState;
 use revm_state::AccountInfo;
 use std::{
@@ -198,53 +198,53 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
         fork
     }
 
-    /// Gets an [`ExecutedBlockWithTrieUpdates`] with [`BlockNumber`], receipts and parent hash.
+    /// Gets an [`ExecutedBlock`] with [`BlockNumber`], receipts and parent hash.
     fn get_executed_block(
         &mut self,
         block_number: BlockNumber,
         receipts: Vec<Vec<Receipt>>,
         parent_hash: B256,
-    ) -> ExecutedBlockWithTrieUpdates {
+    ) -> ExecutedBlock {
         let block_with_senders = self.generate_random_block(block_number, parent_hash);
 
         let (block, senders) = block_with_senders.split_sealed();
-        ExecutedBlockWithTrieUpdates::new(
-            Arc::new(RecoveredBlock::new_sealed(block, senders)),
-            Arc::new(ExecutionOutcome::new(
+        ExecutedBlock {
+            recovered_block: Arc::new(RecoveredBlock::new_sealed(block, senders)),
+            execution_output: Arc::new(ExecutionOutcome::new(
                 BundleState::default(),
                 receipts,
                 block_number,
                 vec![Requests::default()],
             )),
-            Arc::new(HashedPostState::default()),
-            ExecutedTrieUpdates::empty(),
-        )
+            hashed_state: Arc::new(HashedPostState::default()),
+            trie_updates: Arc::new(TrieUpdates::default()),
+        }
     }
 
-    /// Generates an [`ExecutedBlockWithTrieUpdates`] that includes the given receipts.
+    /// Generates an [`ExecutedBlock`] that includes the given receipts.
     pub fn get_executed_block_with_receipts(
         &mut self,
         receipts: Vec<Vec<Receipt>>,
         parent_hash: B256,
-    ) -> ExecutedBlockWithTrieUpdates {
+    ) -> ExecutedBlock {
         let number = rand::rng().random::<u64>();
         self.get_executed_block(number, receipts, parent_hash)
     }
 
-    /// Generates an [`ExecutedBlockWithTrieUpdates`] with the given [`BlockNumber`].
+    /// Generates an [`ExecutedBlock`] with the given [`BlockNumber`].
     pub fn get_executed_block_with_number(
         &mut self,
         block_number: BlockNumber,
         parent_hash: B256,
-    ) -> ExecutedBlockWithTrieUpdates {
+    ) -> ExecutedBlock {
         self.get_executed_block(block_number, vec![vec![]], parent_hash)
     }
 
     /// Generates a range of executed blocks with ascending block numbers.
     pub fn get_executed_blocks(
         &mut self,
         range: Range<u64>,
-    ) -> impl Iterator<Item = ExecutedBlockWithTrieUpdates> + '_ {
+    ) -> impl Iterator<Item = ExecutedBlock> + '_ {
         let mut parent_hash = B256::default();
         range.map(move |number| {
             let current_parent_hash = parent_hash;

crates/cli/commands/src/stage/drop.rs

@@ -15,7 +15,7 @@ use reth_db_common::{
 };
 use reth_node_api::{HeaderTy, ReceiptTy, TxTy};
 use reth_node_core::args::StageEnum;
-use reth_provider::{DBProvider, DatabaseProviderFactory, StaticFileProviderFactory};
+use reth_provider::{DBProvider, DatabaseProviderFactory, StaticFileProviderFactory, TrieWriter};
 use reth_prune::PruneSegment;
 use reth_stages::StageId;
 use reth_static_file_types::StaticFileSegment;
@@ -138,6 +138,10 @@ impl<C: ChainSpecParser> Command<C> {
                     None,
                 )?;
             }
+            StageEnum::MerkleChangeSets => {
+                provider_rw.clear_trie_changesets()?;
+                reset_stage_checkpoint(tx, StageId::MerkleChangeSets)?;
+            }
             StageEnum::AccountHistory | StageEnum::StorageHistory => {
                 tx.clear::<tables::AccountsHistory>()?;
                 tx.clear::<tables::StoragesHistory>()?;