revert: fix(engine): recompute trie updates for forked blocks (#16500) (#16565)

Author: shekhirin

crates/chain-state/src/in_memory.rs

@@ -798,83 +798,28 @@ impl<N: NodePrimitives> ExecutedBlock<N> {
     }
 }
 
-/// 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)]
-    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,
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    Default,
+    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,
+    /// Trie updates that result of applying the block.
+    pub trie: Arc<TrieUpdates>,
 }
 
 impl<N: NodePrimitives> ExecutedBlockWithTrieUpdates<N> {
@@ -883,15 +828,15 @@ impl<N: NodePrimitives> ExecutedBlockWithTrieUpdates<N> {
         recovered_block: Arc<RecoveredBlock<N::Block>>,
         execution_output: Arc<ExecutionOutcome<N::Receipt>>,
         hashed_state: Arc<HashedPostState>,
-        trie: ExecutedTrieUpdates,
+        trie: Arc<TrieUpdates>,
     ) -> 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 reference to the trie updates for the block
     #[inline]
-    pub fn trie_updates(&self) -> Option<&TrieUpdates> {
-        self.trie.as_ref()
+    pub fn trie_updates(&self) -> &TrieUpdates {
+        &self.trie
     }
 
     /// Converts the value into [`SealedBlock`].

crates/chain-state/src/memory_overlay.rs

@@ -26,7 +26,7 @@ pub struct MemoryOverlayStateProviderRef<
     /// The collection of executed parent blocks. Expected order is newest to oldest.
     pub(crate) in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
     /// Lazy-loaded in-memory trie data.
-    pub(crate) trie_input: OnceLock<TrieInput>,
+    pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
 }
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
@@ -45,7 +45,7 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
         historical: Box<dyn StateProvider + 'a>,
         in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
     ) -> Self {
-        Self { historical, in_memory, trie_input: OnceLock::new() }
+        Self { historical, in_memory, trie_state: OnceLock::new() }
     }
 
     /// Turn this state provider into a state provider
@@ -54,14 +54,14 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
     }
 
     /// Return lazy-loaded trie state aggregated from in-memory blocks.
-    fn trie_input(&self) -> &TrieInput {
-        self.trie_input.get_or_init(|| {
-            TrieInput::from_blocks(
-                self.in_memory
-                    .iter()
-                    .rev()
-                    .map(|block| (block.hashed_state.as_ref(), block.trie.as_ref())),
-            )
+    fn trie_state(&self) -> &MemoryOverlayTrieState {
+        self.trie_state.get_or_init(|| {
+            let mut trie_state = MemoryOverlayTrieState::default();
+            for block in self.in_memory.iter().rev() {
+                trie_state.state.extend_ref(block.hashed_state.as_ref());
+                trie_state.nodes.extend_ref(block.trie.as_ref());
+            }
+            trie_state
         })
     }
 }
@@ -117,7 +117,8 @@ impl<N: NodePrimitives> StateRootProvider for MemoryOverlayStateProviderRef<'_,
     }
 
     fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
-        input.prepend_self(self.trie_input().clone());
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
         self.historical.state_root_from_nodes(input)
     }
 
@@ -132,15 +133,16 @@ impl<N: NodePrimitives> StateRootProvider for MemoryOverlayStateProviderRef<'_,
         &self,
         mut input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        input.prepend_self(self.trie_input().clone());
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
         self.historical.state_root_from_nodes_with_updates(input)
     }
 }
 
 impl<N: NodePrimitives> StorageRootProvider for MemoryOverlayStateProviderRef<'_, N> {
     // TODO: Currently this does not reuse available in-memory trie nodes.
     fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
-        let state = &self.trie_input().state;
+        let state = &self.trie_state().state;
         let mut hashed_storage =
             state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
         hashed_storage.extend(&storage);
@@ -154,7 +156,7 @@ impl<N: NodePrimitives> StorageRootProvider for MemoryOverlayStateProviderRef<'_
         slot: B256,
         storage: HashedStorage,
     ) -> ProviderResult<reth_trie::StorageProof> {
-        let state = &self.trie_input().state;
+        let state = &self.trie_state().state;
         let mut hashed_storage =
             state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
         hashed_storage.extend(&storage);
@@ -168,7 +170,7 @@ impl<N: NodePrimitives> StorageRootProvider for MemoryOverlayStateProviderRef<'_
         slots: &[B256],
         storage: HashedStorage,
     ) -> ProviderResult<StorageMultiProof> {
-        let state = &self.trie_input().state;
+        let state = &self.trie_state().state;
         let mut hashed_storage =
             state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
         hashed_storage.extend(&storage);
@@ -183,7 +185,8 @@ impl<N: NodePrimitives> StateProofProvider for MemoryOverlayStateProviderRef<'_,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        input.prepend_self(self.trie_input().clone());
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
         self.historical.proof(input, address, slots)
     }
 
@@ -192,12 +195,14 @@ impl<N: NodePrimitives> StateProofProvider for MemoryOverlayStateProviderRef<'_,
         mut input: TrieInput,
         targets: MultiProofTargets,
     ) -> ProviderResult<MultiProof> {
-        input.prepend_self(self.trie_input().clone());
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
         self.historical.multiproof(input, targets)
     }
 
     fn witness(&self, mut input: TrieInput, target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
-        input.prepend_self(self.trie_input().clone());
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
         self.historical.witness(input, target)
     }
 }
@@ -233,3 +238,12 @@ impl<N: NodePrimitives> StateProvider for MemoryOverlayStateProviderRef<'_, N> {
         self.historical.bytecode_by_hash(code_hash)
     }
 }
+
+/// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
+#[derive(Clone, Default, Debug)]
+pub(crate) struct MemoryOverlayTrieState {
+    /// The collection of aggregated in-memory trie updates.
+    pub(crate) nodes: TrieUpdates,
+    /// The collection of hashed state from in-memory blocks.
+    pub(crate) state: HashedPostState,
+}

crates/chain-state/src/test_utils.rs

@@ -1,6 +1,6 @@
 use crate::{
     in_memory::ExecutedBlockWithTrieUpdates, CanonStateNotification, CanonStateNotifications,
-    CanonStateSubscriptions, ExecutedTrieUpdates,
+    CanonStateSubscriptions,
 };
 use alloy_consensus::{
     Header, SignableTransaction, Transaction as _, TxEip1559, TxReceipt, EMPTY_ROOT_HASH,
@@ -25,7 +25,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::{
@@ -222,7 +222,7 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
                 vec![Requests::default()],
             )),
             Arc::new(HashedPostState::default()),
-            ExecutedTrieUpdates::empty(),
+            Arc::new(TrieUpdates::default()),
         )
     }
 

crates/engine/tree/src/tree/error.rs

@@ -1,7 +1,6 @@
 //! Internal errors for the tree module.
 
 use alloy_consensus::BlockHeader;
-use alloy_primitives::B256;
 use reth_consensus::ConsensusError;
 use reth_errors::{BlockExecutionError, BlockValidationError, ProviderError};
 use reth_evm::execute::InternalBlockExecutionError;
@@ -18,20 +17,6 @@ pub enum AdvancePersistenceError {
     /// A provider error
     #[error(transparent)]
     Provider(#[from] ProviderError),
-    /// Missing ancestor.
-    ///
-    /// This error occurs when we need to compute the state root for a block with missing trie
-    /// updates, but the ancestor block is not available. State root computation requires the state
-    /// from the parent block as a starting point.
-    ///
-    /// A block may be missing the trie updates when it's a fork chain block 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 until the moment when database is unwound to the canonical
-    /// chain.
-    ///
-    /// Also see [`reth_chain_state::ExecutedTrieUpdates::Missing`].
-    #[error("Missing ancestor with hash {0}")]
-    MissingAncestor(B256),
 }
 
 #[derive(thiserror::Error)]