feat: integrate memory overlay in canonical state (#9817)

2026-02-18 02:41:31 -05:00 · 2024-07-26 11:51:11 +02:00
parent f734e614d7
commit ac2853884a
3 changed files with 265 additions and 28 deletions
--- a/crates/chain-state/src/in_memory.rs
+++ b/crates/chain-state/src/in_memory.rs
@@ -1,7 +1,8 @@
 //! Types for tracking the canonical chain state in memory.

 use crate::{
-    CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications, ChainInfoTracker,
+    CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications,
+    ChainInfoTracker, MemoryOverlayStateProvider,
 };
 use parking_lot::RwLock;
 use reth_chainspec::ChainInfo;
@@ -10,6 +11,7 @@ use reth_primitives::{
    Address, BlockNumHash, Header, Receipt, Receipts, SealedBlock, SealedBlockWithSenders,
    SealedHeader, B256,
 };
+use reth_storage_api::StateProviderBox;
 use reth_trie::{updates::TrieUpdates, HashedPostState};
 use std::{collections::HashMap, ops::Deref, sync::Arc, time::Instant};
 use tokio::sync::broadcast;
@@ -354,6 +356,24 @@ impl CanonicalInMemoryState {
    pub fn notify_canon_state(&self, event: CanonStateNotification) {
        self.inner.canon_state_notification_sender.send(event).ok();
    }
+
+    /// Return state provider with reference to in-memory blocks that overlay database state.
+    ///
+    /// This merges the state of all blocks that are part of the chain that the requested block is
+    /// the head of. This includes all blocks that connect back to the canonical block on disk.
+    pub fn state_provider(
+        &self,
+        hash: B256,
+        historical: StateProviderBox,
+    ) -> MemoryOverlayStateProvider {
+        let in_memory = if let Some(state) = self.state_by_hash(hash) {
+            state.chain().into_iter().map(|block_state| block_state.block()).collect()
+        } else {
+            Vec::new()
+        };
+
+        MemoryOverlayStateProvider::new(in_memory, historical)
+    }
 }

 /// State after applying the given block, this block is part of the canonical chain that partially
@@ -434,6 +454,27 @@ impl BlockState {
            })
            .unwrap_or_default()
    }
+
+    /// Returns a vector of parent `BlockStates` starting from the oldest one.
+    pub fn parent_state_chain(&self) -> Vec<&Self> {
+        let mut parents = Vec::new();
+        let mut current = self.parent.as_deref();
+
+        while let Some(parent) = current {
+            parents.insert(0, parent);
+            current = parent.parent.as_deref();
+        }
+
+        parents
+    }
+
+    /// Returns a vector of `BlockStates` representing the entire in memory chain,
+    /// including self as the last element.
+    pub fn chain(&self) -> Vec<&Self> {
+        let mut chain = self.parent_state_chain();
+        chain.push(self);
+        chain
+    }
 }

 /// Represents an executed block stored in-memory.
@@ -575,17 +616,104 @@ mod tests {
    use super::*;
    use crate::test_utils::{get_executed_block_with_number, get_executed_block_with_receipts};
    use rand::Rng;
-    use reth_primitives::Receipt;
+    use reth_errors::ProviderResult;
+    use reth_primitives::{Account, BlockNumber, Bytecode, Receipt, StorageKey, StorageValue};
+    use reth_storage_api::{
+        AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
+    };
+    use reth_trie::AccountProof;

-    fn create_mock_state(block_number: u64) -> BlockState {
-        BlockState::new(get_executed_block_with_number(block_number, B256::random()))
+    fn create_mock_state(block_number: u64, parent_hash: B256) -> BlockState {
+        BlockState::new(get_executed_block_with_number(block_number, parent_hash))
+    }
+
+    fn create_mock_state_chain(num_blocks: u64) -> Vec<BlockState> {
+        let mut chain = Vec::with_capacity(num_blocks as usize);
+        let mut parent_hash = B256::random();
+        let mut parent_state: Option<BlockState> = None;
+
+        for i in 1..=num_blocks {
+            let mut state = create_mock_state(i, parent_hash);
+            if let Some(parent) = parent_state {
+                state.parent = Some(Box::new(parent));
+            }
+            parent_hash = state.hash();
+            parent_state = Some(state.clone());
+            chain.push(state);
+        }
+
+        chain
+    }
+
+    struct MockStateProvider;
+
+    impl StateProvider for MockStateProvider {
+        fn storage(
+            &self,
+            _address: Address,
+            _storage_key: StorageKey,
+        ) -> ProviderResult<Option<StorageValue>> {
+            Ok(None)
+        }
+
+        fn bytecode_by_hash(&self, _code_hash: B256) -> ProviderResult<Option<Bytecode>> {
+            Ok(None)
+        }
+    }
+
+    impl BlockHashReader for MockStateProvider {
+        fn block_hash(&self, _number: BlockNumber) -> ProviderResult<Option<B256>> {
+            Ok(None)
+        }
+
+        fn canonical_hashes_range(
+            &self,
+            _start: BlockNumber,
+            _end: BlockNumber,
+        ) -> ProviderResult<Vec<B256>> {
+            Ok(vec![])
+        }
+    }
+
+    impl AccountReader for MockStateProvider {
+        fn basic_account(&self, _address: Address) -> ProviderResult<Option<Account>> {
+            Ok(None)
+        }
+    }
+
+    impl StateRootProvider for MockStateProvider {
+        fn hashed_state_root(&self, _hashed_state: &HashedPostState) -> ProviderResult<B256> {
+            Ok(B256::random())
+        }
+
+        fn hashed_state_root_with_updates(
+            &self,
+            _hashed_state: &HashedPostState,
+        ) -> ProviderResult<(B256, TrieUpdates)> {
+            Ok((B256::random(), TrieUpdates::default()))
+        }
+
+        fn state_root(&self, _bundle_state: &revm::db::BundleState) -> ProviderResult<B256> {
+            Ok(B256::random())
+        }
+    }
+
+    impl StateProofProvider for MockStateProvider {
+        fn hashed_proof(
+            &self,
+            _hashed_state: &HashedPostState,
+            _address: Address,
+            _slots: &[B256],
+        ) -> ProviderResult<AccountProof> {
+            Ok(AccountProof::new(Address::random()))
+        }
    }

    #[test]
    fn test_in_memory_state_impl_state_by_hash() {
        let mut state_by_hash = HashMap::new();
        let number = rand::thread_rng().gen::<u64>();
-        let state = Arc::new(create_mock_state(number));
+        let state = Arc::new(create_mock_state(number, B256::random()));
        state_by_hash.insert(state.hash(), state.clone());

        let in_memory_state = InMemoryState::new(state_by_hash, HashMap::new(), None);
@@ -600,7 +728,7 @@ mod tests {
        let mut hash_by_number = HashMap::new();

        let number = rand::thread_rng().gen::<u64>();
-        let state = Arc::new(create_mock_state(number));
+        let state = Arc::new(create_mock_state(number, B256::random()));
        let hash = state.hash();

        state_by_hash.insert(hash, state.clone());
@@ -616,9 +744,9 @@ mod tests {
    fn test_in_memory_state_impl_head_state() {
        let mut state_by_hash = HashMap::new();
        let mut hash_by_number = HashMap::new();
-        let state1 = Arc::new(create_mock_state(1));
-        let state2 = Arc::new(create_mock_state(2));
+        let state1 = Arc::new(create_mock_state(1, B256::random()));
        let hash1 = state1.hash();
+        let state2 = Arc::new(create_mock_state(2, hash1));
        let hash2 = state2.hash();
        hash_by_number.insert(1, hash1);
        hash_by_number.insert(2, hash2);
@@ -635,7 +763,7 @@ mod tests {
    #[test]
    fn test_in_memory_state_impl_pending_state() {
        let pending_number = rand::thread_rng().gen::<u64>();
-        let pending_state = create_mock_state(pending_number);
+        let pending_state = create_mock_state(pending_number, B256::random());
        let pending_hash = pending_state.hash();

        let in_memory_state =
@@ -733,4 +861,104 @@ mod tests {

        assert_eq!(state.inner.in_memory_state.block_count(), 1);
    }
+
+    #[test]
+    fn test_canonical_in_memory_state_state_provider() {
+        let block1 = get_executed_block_with_number(1, B256::random());
+        let block2 = get_executed_block_with_number(2, block1.block().hash());
+        let block3 = get_executed_block_with_number(3, block2.block().hash());
+
+        let state1 = BlockState::new(block1.clone());
+        let state2 = BlockState::with_parent(block2.clone(), Some(state1.clone()));
+        let state3 = BlockState::with_parent(block3.clone(), Some(state2.clone()));
+
+        let mut blocks = HashMap::new();
+        blocks.insert(block1.block().hash(), Arc::new(state1));
+        blocks.insert(block2.block().hash(), Arc::new(state2));
+        blocks.insert(block3.block().hash(), Arc::new(state3));
+
+        let mut numbers = HashMap::new();
+        numbers.insert(1, block1.block().hash());
+        numbers.insert(2, block2.block().hash());
+        numbers.insert(3, block3.block().hash());
+
+        let canonical_state = CanonicalInMemoryState::new(blocks, numbers, None);
+
+        let historical: StateProviderBox = Box::new(MockStateProvider);
+
+        let overlay_provider = canonical_state.state_provider(block3.block().hash(), historical);
+
+        assert_eq!(overlay_provider.in_memory.len(), 3);
+        assert_eq!(overlay_provider.in_memory[0].block().number, 1);
+        assert_eq!(overlay_provider.in_memory[1].block().number, 2);
+        assert_eq!(overlay_provider.in_memory[2].block().number, 3);
+
+        assert_eq!(
+            overlay_provider.in_memory[1].block().parent_hash,
+            overlay_provider.in_memory[0].block().hash()
+        );
+        assert_eq!(
+            overlay_provider.in_memory[2].block().parent_hash,
+            overlay_provider.in_memory[1].block().hash()
+        );
+
+        let unknown_hash = B256::random();
+        let empty_overlay_provider =
+            canonical_state.state_provider(unknown_hash, Box::new(MockStateProvider));
+        assert_eq!(empty_overlay_provider.in_memory.len(), 0);
+    }
+
+    #[test]
+    fn test_block_state_parent_blocks() {
+        let chain = create_mock_state_chain(4);
+
+        let parents = chain[3].parent_state_chain();
+        assert_eq!(parents.len(), 3);
+        assert_eq!(parents[0].block().block.number, 1);
+        assert_eq!(parents[1].block().block.number, 2);
+        assert_eq!(parents[2].block().block.number, 3);
+
+        let parents = chain[2].parent_state_chain();
+        assert_eq!(parents.len(), 2);
+        assert_eq!(parents[0].block().block.number, 1);
+        assert_eq!(parents[1].block().block.number, 2);
+
+        let parents = chain[0].parent_state_chain();
+        assert_eq!(parents.len(), 0);
+    }
+
+    #[test]
+    fn test_block_state_single_block_state_chain() {
+        let single_block_number = 1;
+        let single_block = create_mock_state(single_block_number, B256::random());
+        let single_block_hash = single_block.block().block.hash();
+
+        let parents = single_block.parent_state_chain();
+        assert_eq!(parents.len(), 0);
+
+        let block_state_chain = single_block.chain();
+        assert_eq!(block_state_chain.len(), 1);
+        assert_eq!(block_state_chain[0].block().block.number, single_block_number);
+        assert_eq!(block_state_chain[0].block().block.hash(), single_block_hash);
+    }
+
+    #[test]
+    fn test_block_state_chain() {
+        let chain = create_mock_state_chain(3);
+
+        let block_state_chain = chain[2].chain();
+        assert_eq!(block_state_chain.len(), 3);
+        assert_eq!(block_state_chain[0].block().block.number, 1);
+        assert_eq!(block_state_chain[1].block().block.number, 2);
+        assert_eq!(block_state_chain[2].block().block.number, 3);
+
+        let block_state_chain = chain[1].chain();
+        assert_eq!(block_state_chain.len(), 2);
+        assert_eq!(block_state_chain[0].block().block.number, 1);
+        assert_eq!(block_state_chain[1].block().block.number, 2);
+
+        let block_state_chain = chain[0].chain();
+        assert_eq!(block_state_chain.len(), 1);
+        assert_eq!(block_state_chain[0].block().block.number, 1);
+    }
 }
--- a/crates/chain-state/src/memory_overlay.rs
+++ b/crates/chain-state/src/memory_overlay.rs
@@ -11,11 +11,11 @@ use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 #[allow(missing_debug_implementations)]
 pub struct MemoryOverlayStateProvider {
    /// The collection of executed parent blocks.
-    in_memory: Vec<ExecutedBlock>,
+    pub(crate) in_memory: Vec<ExecutedBlock>,
    /// The collection of hashed state from in-memory blocks.
-    hashed_post_state: HashedPostState,
+    pub(crate) hashed_post_state: HashedPostState,
    /// Historical state provider for state lookups that are not found in in-memory blocks.
-    historical: Box<dyn StateProvider>,
+    pub(crate) historical: Box<dyn StateProvider>,
 }

 impl MemoryOverlayStateProvider {
--- a/crates/storage/provider/src/providers/blockchain_provider.rs
+++ b/crates/storage/provider/src/providers/blockchain_provider.rs
@@ -1,12 +1,12 @@
 use crate::{
-    providers::BundleStateProvider, AccountReader, BlockHashReader, BlockIdReader, BlockNumReader,
-    BlockReader, BlockReaderIdExt, BlockSource, BlockchainTreePendingStateProvider,
-    CanonChainTracker, CanonStateNotifications, CanonStateSubscriptions, ChainSpecProvider,
-    ChangeSetReader, DatabaseProviderFactory, DatabaseProviderRO, EvmEnvProvider,
-    FullExecutionDataProvider, HeaderProvider, ProviderError, ProviderFactory,
-    PruneCheckpointReader, ReceiptProvider, ReceiptProviderIdExt, RequestsProvider,
-    StageCheckpointReader, StateProviderBox, StateProviderFactory, StaticFileProviderFactory,
-    TransactionVariant, TransactionsProvider, WithdrawalsProvider,
+    providers::{BundleStateProvider, StaticFileProvider},
+    AccountReader, BlockHashReader, BlockIdReader, BlockNumReader, BlockReader, BlockReaderIdExt,
+    BlockSource, BlockchainTreePendingStateProvider, CanonChainTracker, CanonStateNotifications,
+    CanonStateSubscriptions, ChainSpecProvider, ChangeSetReader, DatabaseProviderFactory,
+    DatabaseProviderRO, EvmEnvProvider, FullExecutionDataProvider, HeaderProvider, ProviderError,
+    ProviderFactory, PruneCheckpointReader, ReceiptProvider, ReceiptProviderIdExt,
+    RequestsProvider, StageCheckpointReader, StateProviderBox, StateProviderFactory,
+    StaticFileProviderFactory, TransactionVariant, TransactionsProvider, WithdrawalsProvider,
 };
 use alloy_rpc_types_engine::ForkchoiceState;
 use reth_chain_state::CanonicalInMemoryState;
@@ -33,8 +33,6 @@ use std::{
 };
 use tracing::trace;

-use super::StaticFileProvider;
-
 /// The main type for interacting with the blockchain.
 ///
 /// This type serves as the main entry point for interacting with the blockchain and provides data
@@ -624,15 +622,28 @@ where
    fn pending(&self) -> ProviderResult<StateProviderBox> {
        trace!(target: "providers::blockchain", "Getting provider for pending state");

-        // TODO: check in memory overlay https://github.com/paradigmxyz/reth/issues/9614
+        if let Some(block) = self.canonical_in_memory_state.pending_block_num_hash() {
+            let historical = self.database.history_by_block_hash(block.hash)?;
+            let pending_provider =
+                self.canonical_in_memory_state.state_provider(block.hash, historical);
+
+            return Ok(Box::new(pending_provider));
+        }

        // fallback to latest state if the pending block is not available
        self.latest()
    }

-    fn pending_state_by_hash(&self, _block_hash: B256) -> ProviderResult<Option<StateProviderBox>> {
-        // TODO: check in memory overlay https://github.com/paradigmxyz/reth/issues/9614
+    fn pending_state_by_hash(&self, block_hash: B256) -> ProviderResult<Option<StateProviderBox>> {
+        let historical = self.database.history_by_block_hash(block_hash)?;
+        if let Some(block) = self.canonical_in_memory_state.pending_block_num_hash() {
+            if block.hash == block_hash {
+                let pending_provider =
+                    self.canonical_in_memory_state.state_provider(block_hash, historical);

+                return Ok(Some(Box::new(pending_provider)))
+            }
+        }
        Ok(None)
    }

@@ -640,10 +651,8 @@ where
        &self,
        bundle_state_data: Box<dyn FullExecutionDataProvider>,
    ) -> ProviderResult<StateProviderBox> {
-        let canonical_fork = bundle_state_data.canonical_fork();
-        trace!(target: "providers::blockchain", ?canonical_fork, "Returning post state provider");
+        let state_provider = self.pending()?;

-        let state_provider = self.history_by_block_hash(canonical_fork.hash)?;
        let bundle_state_provider = BundleStateProvider::new(state_provider, bundle_state_data);
        Ok(Box::new(bundle_state_provider))
    }