feat(providers): add AtomicBlockchainProvider (#11705)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/chain-state/src/in_memory.rs

@@ -4,7 +4,7 @@ use crate::{
     CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications,
     ChainInfoTracker, MemoryOverlayStateProvider,
 };
-use alloy_eips::BlockNumHash;
+use alloy_eips::{BlockHashOrNumber, BlockNumHash};
 use alloy_primitives::{map::HashMap, Address, TxHash, B256};
 use parking_lot::RwLock;
 use reth_chainspec::ChainInfo;
@@ -514,7 +514,7 @@ impl CanonicalInMemoryState {
         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()
+            state.chain().map(|block_state| block_state.block()).collect()
         } else {
             Vec::new()
         };
@@ -692,10 +692,8 @@ impl BlockState {
     /// Returns a vector of `BlockStates` representing the entire in memory chain.
     /// The block state order in the output vector is newest to oldest (highest to lowest),
     /// including self as the first element.
-    pub fn chain(&self) -> Vec<&Self> {
-        let mut chain = vec![self];
-        self.append_parent_chain(&mut chain);
-        chain
+    pub fn chain(&self) -> impl Iterator<Item = &Self> {
+        std::iter::successors(Some(self), |state| state.parent.as_deref())
     }
 
     /// Appends the parent chain of this [`BlockState`] to the given vector.
@@ -715,10 +713,59 @@ impl BlockState {
     /// This merges the state of all blocks that are part of the chain that the this block is
     /// the head of. This includes all blocks that connect back to the canonical block on disk.
     pub fn state_provider(&self, historical: StateProviderBox) -> MemoryOverlayStateProvider {
-        let in_memory = self.chain().into_iter().map(|block_state| block_state.block()).collect();
+        let in_memory = self.chain().map(|block_state| block_state.block()).collect();
 
         MemoryOverlayStateProvider::new(historical, in_memory)
     }
+
+    /// Tries to find a block by [`BlockHashOrNumber`] in the chain ending at this block.
+    pub fn block_on_chain(&self, hash_or_num: BlockHashOrNumber) -> Option<&Self> {
+        self.chain().find(|block| match hash_or_num {
+            BlockHashOrNumber::Hash(hash) => block.hash() == hash,
+            BlockHashOrNumber::Number(number) => block.number() == number,
+        })
+    }
+
+    /// Tries to find a transaction by [`TxHash`] in the chain ending at this block.
+    pub fn transaction_on_chain(&self, hash: TxHash) -> Option<TransactionSigned> {
+        self.chain().find_map(|block_state| {
+            block_state
+                .block_ref()
+                .block()
+                .body
+                .transactions()
+                .find(|tx| tx.hash() == hash)
+                .cloned()
+        })
+    }
+
+    /// Tries to find a transaction with meta by [`TxHash`] in the chain ending at this block.
+    pub fn transaction_meta_on_chain(
+        &self,
+        tx_hash: TxHash,
+    ) -> Option<(TransactionSigned, TransactionMeta)> {
+        self.chain().find_map(|block_state| {
+            block_state
+                .block_ref()
+                .block()
+                .body
+                .transactions()
+                .enumerate()
+                .find(|(_, tx)| tx.hash() == tx_hash)
+                .map(|(index, tx)| {
+                    let meta = TransactionMeta {
+                        tx_hash,
+                        index: index as u64,
+                        block_hash: block_state.hash(),
+                        block_number: block_state.block_ref().block.number,
+                        base_fee: block_state.block_ref().block.header.base_fee_per_gas,
+                        timestamp: block_state.block_ref().block.timestamp,
+                        excess_blob_gas: block_state.block_ref().block.excess_blob_gas,
+                    };
+                    (tx.clone(), meta)
+                })
+        })
+    }
 }
 
 /// Represents an executed block stored in-memory.
@@ -1382,7 +1429,7 @@ mod tests {
         let parents = single_block.parent_state_chain();
         assert_eq!(parents.len(), 0);
 
-        let block_state_chain = single_block.chain();
+        let block_state_chain = single_block.chain().collect::<Vec<_>>();
         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);
@@ -1393,18 +1440,18 @@ mod tests {
         let mut test_block_builder = TestBlockBuilder::default();
         let chain = create_mock_state_chain(&mut test_block_builder, 3);
 
-        let block_state_chain = chain[2].chain();
+        let block_state_chain = chain[2].chain().collect::<Vec<_>>();
         assert_eq!(block_state_chain.len(), 3);
         assert_eq!(block_state_chain[0].block().block.number, 3);
         assert_eq!(block_state_chain[1].block().block.number, 2);
         assert_eq!(block_state_chain[2].block().block.number, 1);
 
-        let block_state_chain = chain[1].chain();
+        let block_state_chain = chain[1].chain().collect::<Vec<_>>();
         assert_eq!(block_state_chain.len(), 2);
         assert_eq!(block_state_chain[0].block().block.number, 2);
         assert_eq!(block_state_chain[1].block().block.number, 1);
 
-        let block_state_chain = chain[0].chain();
+        let block_state_chain = chain[0].chain().collect::<Vec<_>>();
         assert_eq!(block_state_chain.len(), 1);
         assert_eq!(block_state_chain[0].block().block.number, 1);
     }

crates/chain-state/src/lib.rs

@@ -22,7 +22,7 @@ pub use notifications::{
 };
 
 mod memory_overlay;
-pub use memory_overlay::MemoryOverlayStateProvider;
+pub use memory_overlay::{MemoryOverlayStateProvider, MemoryOverlayStateProviderRef};
 
 #[cfg(any(test, feature = "test-utils"))]
 /// Common test helpers

crates/chain-state/src/memory_overlay.rs

@@ -7,14 +7,26 @@ use alloy_primitives::{
 use reth_errors::ProviderResult;
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateProviderBox,
-    StateRootProvider, StorageRootProvider,
+    AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
+    StorageRootProvider,
 };
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
 };
 use std::sync::OnceLock;
 
+/// A state provider that stores references to in-memory blocks along with their state as well as a
+/// reference of the historical state provider for fallback lookups.
+#[allow(missing_debug_implementations)]
+pub struct MemoryOverlayStateProviderRef<'a> {
+    /// Historical state provider for state lookups that are not found in 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<ExecutedBlock>,
+    /// Lazy-loaded in-memory trie data.
+    pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
+}
+
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
 #[allow(missing_debug_implementations)]
@@ -27,193 +39,200 @@ pub struct MemoryOverlayStateProvider {
     pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
 }
 
-impl MemoryOverlayStateProvider {
-    /// Create new memory overlay state provider.
-    ///
-    /// ## Arguments
-    ///
-    /// - `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>, in_memory: Vec<ExecutedBlock>) -> Self {
-        Self { historical, in_memory, trie_state: OnceLock::new() }
-    }
-
-    /// Turn this state provider into a [`StateProviderBox`]
-    pub fn boxed(self) -> StateProviderBox {
-        Box::new(self)
-    }
-
-    /// Return lazy-loaded trie state aggregated from in-memory blocks.
-    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());
+macro_rules! impl_state_provider {
+    ([$($tokens:tt)*],$type:ty, $historical_type:ty) => {
+        impl $($tokens)*  $type {
+            /// Create new memory overlay state provider.
+            ///
+            /// ## Arguments
+            ///
+            /// - `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: $historical_type, in_memory: Vec<ExecutedBlock>) -> Self {
+                Self { historical, in_memory, trie_state: OnceLock::new() }
             }
-            trie_state
-        })
-    }
-}
 
-impl BlockHashReader for MemoryOverlayStateProvider {
-    fn block_hash(&self, number: BlockNumber) -> ProviderResult<Option<B256>> {
-        for block in &self.in_memory {
-            if block.block.number == number {
-                return Ok(Some(block.block.hash()))
+            /// Turn this state provider into a state provider
+            pub fn boxed(self) -> $historical_type {
+                Box::new(self)
+            }
+
+            /// Return lazy-loaded trie state aggregated from in-memory blocks.
+            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
+                })
             }
         }
 
-        self.historical.block_hash(number)
-    }
+        impl $($tokens)* BlockHashReader for $type {
+            fn block_hash(&self, number: BlockNumber) -> ProviderResult<Option<B256>> {
+                for block in &self.in_memory {
+                    if block.block.number == number {
+                        return Ok(Some(block.block.hash()))
+                    }
+                }
 
-    fn canonical_hashes_range(
-        &self,
-        start: BlockNumber,
-        end: BlockNumber,
-    ) -> ProviderResult<Vec<B256>> {
-        let range = start..end;
-        let mut earliest_block_number = None;
-        let mut in_memory_hashes = Vec::new();
-        for block in &self.in_memory {
-            if range.contains(&block.block.number) {
-                in_memory_hashes.insert(0, block.block.hash());
-                earliest_block_number = Some(block.block.number);
+                self.historical.block_hash(number)
+            }
+
+            fn canonical_hashes_range(
+                &self,
+                start: BlockNumber,
+                end: BlockNumber,
+            ) -> ProviderResult<Vec<B256>> {
+                let range = start..end;
+                let mut earliest_block_number = None;
+                let mut in_memory_hashes = Vec::new();
+                for block in &self.in_memory {
+                    if range.contains(&block.block.number) {
+                        in_memory_hashes.insert(0, block.block.hash());
+                        earliest_block_number = Some(block.block.number);
+                    }
+                }
+
+                let mut hashes =
+                    self.historical.canonical_hashes_range(start, earliest_block_number.unwrap_or(end))?;
+                hashes.append(&mut in_memory_hashes);
+                Ok(hashes)
             }
         }
 
-        let mut hashes =
-            self.historical.canonical_hashes_range(start, earliest_block_number.unwrap_or(end))?;
-        hashes.append(&mut in_memory_hashes);
-        Ok(hashes)
-    }
-}
+        impl $($tokens)* AccountReader for $type {
+            fn basic_account(&self, address: Address) -> ProviderResult<Option<Account>> {
+                for block in &self.in_memory {
+                    if let Some(account) = block.execution_output.account(&address) {
+                        return Ok(account)
+                    }
+                }
 
-impl AccountReader for MemoryOverlayStateProvider {
-    fn basic_account(&self, address: Address) -> ProviderResult<Option<Account>> {
-        for block in &self.in_memory {
-            if let Some(account) = block.execution_output.account(&address) {
-                return Ok(account)
+                self.historical.basic_account(address)
             }
         }
 
-        self.historical.basic_account(address)
-    }
-}
+        impl $($tokens)* StateRootProvider for $type {
+            fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
+                self.state_root_from_nodes(TrieInput::from_state(state))
+            }
 
-impl StateRootProvider for MemoryOverlayStateProvider {
-    fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
-        self.state_root_from_nodes(TrieInput::from_state(state))
-    }
+            fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
+                let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+                input.prepend_cached(nodes, state);
+                self.historical.state_root_from_nodes(input)
+            }
 
-    fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
-        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
-        input.prepend_cached(nodes, state);
-        self.historical.state_root_from_nodes(input)
-    }
+            fn state_root_with_updates(
+                &self,
+                state: HashedPostState,
+            ) -> ProviderResult<(B256, TrieUpdates)> {
+                self.state_root_from_nodes_with_updates(TrieInput::from_state(state))
+            }
 
-    fn state_root_with_updates(
-        &self,
-        state: HashedPostState,
-    ) -> ProviderResult<(B256, TrieUpdates)> {
-        self.state_root_from_nodes_with_updates(TrieInput::from_state(state))
-    }
-
-    fn state_root_from_nodes_with_updates(
-        &self,
-        mut input: TrieInput,
-    ) -> ProviderResult<(B256, TrieUpdates)> {
-        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
-        input.prepend_cached(nodes, state);
-        self.historical.state_root_from_nodes_with_updates(input)
-    }
-}
-
-impl StorageRootProvider for MemoryOverlayStateProvider {
-    // 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_state().state;
-        let mut hashed_storage =
-            state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
-        hashed_storage.extend(&storage);
-        self.historical.storage_root(address, hashed_storage)
-    }
-
-    // TODO: Currently this does not reuse available in-memory trie nodes.
-    fn storage_proof(
-        &self,
-        address: Address,
-        slot: B256,
-        storage: HashedStorage,
-    ) -> ProviderResult<reth_trie::StorageProof> {
-        let state = &self.trie_state().state;
-        let mut hashed_storage =
-            state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
-        hashed_storage.extend(&storage);
-        self.historical.storage_proof(address, slot, hashed_storage)
-    }
-}
-
-impl StateProofProvider for MemoryOverlayStateProvider {
-    fn proof(
-        &self,
-        mut input: TrieInput,
-        address: Address,
-        slots: &[B256],
-    ) -> ProviderResult<AccountProof> {
-        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
-        input.prepend_cached(nodes, state);
-        self.historical.proof(input, address, slots)
-    }
-
-    fn multiproof(
-        &self,
-        mut input: TrieInput,
-        targets: HashMap<B256, HashSet<B256>>,
-    ) -> ProviderResult<MultiProof> {
-        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<HashMap<B256, Bytes>> {
-        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
-        input.prepend_cached(nodes, state);
-        self.historical.witness(input, target)
-    }
-}
-
-impl StateProvider for MemoryOverlayStateProvider {
-    fn storage(
-        &self,
-        address: Address,
-        storage_key: StorageKey,
-    ) -> ProviderResult<Option<StorageValue>> {
-        for block in &self.in_memory {
-            if let Some(value) = block.execution_output.storage(&address, storage_key.into()) {
-                return Ok(Some(value))
+            fn state_root_from_nodes_with_updates(
+                &self,
+                mut input: TrieInput,
+            ) -> ProviderResult<(B256, TrieUpdates)> {
+                let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+                input.prepend_cached(nodes, state);
+                self.historical.state_root_from_nodes_with_updates(input)
             }
         }
 
-        self.historical.storage(address, storage_key)
-    }
+        impl $($tokens)* StorageRootProvider for $type {
+            // 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_state().state;
+                let mut hashed_storage =
+                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                hashed_storage.extend(&storage);
+                self.historical.storage_root(address, hashed_storage)
+            }
 
-    fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
-        for block in &self.in_memory {
-            if let Some(contract) = block.execution_output.bytecode(&code_hash) {
-                return Ok(Some(contract))
+            // TODO: Currently this does not reuse available in-memory trie nodes.
+            fn storage_proof(
+                &self,
+                address: Address,
+                slot: B256,
+                storage: HashedStorage,
+            ) -> ProviderResult<reth_trie::StorageProof> {
+                let state = &self.trie_state().state;
+                let mut hashed_storage =
+                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                hashed_storage.extend(&storage);
+                self.historical.storage_proof(address, slot, hashed_storage)
             }
         }
 
-        self.historical.bytecode_by_hash(code_hash)
-    }
+        impl $($tokens)* StateProofProvider for $type {
+            fn proof(
+                &self,
+                mut input: TrieInput,
+                address: Address,
+                slots: &[B256],
+            ) -> ProviderResult<AccountProof> {
+                let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+                input.prepend_cached(nodes, state);
+                self.historical.proof(input, address, slots)
+            }
+
+            fn multiproof(
+                &self,
+                mut input: TrieInput,
+                targets: HashMap<B256, HashSet<B256>>,
+            ) -> ProviderResult<MultiProof> {
+                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<HashMap<B256, Bytes>> {
+                let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+                input.prepend_cached(nodes, state);
+                self.historical.witness(input, target)
+            }
+        }
+
+        impl $($tokens)* StateProvider for $type {
+            fn storage(
+                &self,
+                address: Address,
+                storage_key: StorageKey,
+            ) -> ProviderResult<Option<StorageValue>> {
+                for block in &self.in_memory {
+                    if let Some(value) = block.execution_output.storage(&address, storage_key.into()) {
+                        return Ok(Some(value))
+                    }
+                }
+
+                self.historical.storage(address, storage_key)
+            }
+
+            fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
+                for block in &self.in_memory {
+                    if let Some(contract) = block.execution_output.bytecode(&code_hash) {
+                        return Ok(Some(contract))
+                    }
+                }
+
+                self.historical.bytecode_by_hash(code_hash)
+            }
+        }
+    };
 }
 
+impl_state_provider!([], MemoryOverlayStateProvider, Box<dyn StateProvider>);
+impl_state_provider!([<'a>], MemoryOverlayStateProviderRef<'a>, Box<dyn StateProvider + 'a>);
+
 /// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
 #[derive(Clone, Default, Debug)]
 pub(crate) struct MemoryOverlayTrieState {