use crate::{ traits::{BlockSource, ReceiptProvider}, AccountReader, BlockHashReader, BlockIdReader, BlockNumReader, BlockReader, BlockReaderIdExt, ChainSpecProvider, ChangeSetReader, EvmEnvProvider, HeaderProvider, ReceiptProviderIdExt, RequestsProvider, StateProvider, StateProviderBox, StateProviderFactory, StateRootProvider, TransactionVariant, TransactionsProvider, WithdrawalsProvider, }; use parking_lot::Mutex; use reth_chainspec::{ChainInfo, ChainSpec}; use reth_db_api::models::{AccountBeforeTx, StoredBlockBodyIndices}; use reth_evm::ConfigureEvmEnv; use reth_primitives::{ keccak256, Account, Address, Block, BlockHash, BlockHashOrNumber, BlockId, BlockNumber, BlockNumberOrTag, BlockWithSenders, Bytecode, Bytes, Header, Receipt, SealedBlock, SealedBlockWithSenders, SealedHeader, StorageKey, StorageValue, TransactionMeta, TransactionSigned, TransactionSignedNoHash, TxHash, TxNumber, Withdrawal, Withdrawals, B256, U256, }; use reth_stages_types::{StageCheckpoint, StageId}; use reth_storage_api::{StageCheckpointReader, StateProofProvider}; use reth_storage_errors::provider::{ProviderError, ProviderResult}; use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, }; use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg}; use std::{ collections::{BTreeMap, HashMap}, ops::{RangeBounds, RangeInclusive}, sync::Arc, }; /// A mock implementation for Provider interfaces. #[derive(Debug, Clone)] pub struct MockEthProvider { /// Local block store pub blocks: Arc>>, /// Local header store pub headers: Arc>>, /// Local account store pub accounts: Arc>>, /// Local chain spec pub chain_spec: Arc, /// Local state roots pub state_roots: Arc>>, } impl Default for MockEthProvider { fn default() -> Self { Self { blocks: Default::default(), headers: Default::default(), accounts: Default::default(), chain_spec: Arc::new(reth_chainspec::ChainSpecBuilder::mainnet().build()), state_roots: Default::default(), } } } /// An extended account for local store #[derive(Debug, Clone)] pub struct ExtendedAccount { account: Account, bytecode: Option, storage: HashMap, } impl ExtendedAccount { /// Create new instance of extended account pub fn new(nonce: u64, balance: U256) -> Self { Self { account: Account { nonce, balance, bytecode_hash: None }, bytecode: None, storage: Default::default(), } } /// Set bytecode and bytecode hash on the extended account pub fn with_bytecode(mut self, bytecode: Bytes) -> Self { let hash = keccak256(&bytecode); self.account.bytecode_hash = Some(hash); self.bytecode = Some(Bytecode::new_raw(bytecode)); self } /// Add storage to the extended account. If the storage key is already present, /// the value is updated. pub fn extend_storage( mut self, storage: impl IntoIterator, ) -> Self { self.storage.extend(storage); self } } impl MockEthProvider { /// Add block to local block store pub fn add_block(&self, hash: B256, block: Block) { self.add_header(hash, block.header.clone()); self.blocks.lock().insert(hash, block); } /// Add multiple blocks to local block store pub fn extend_blocks(&self, iter: impl IntoIterator) { for (hash, block) in iter { self.add_header(hash, block.header.clone()); self.add_block(hash, block) } } /// Add header to local header store pub fn add_header(&self, hash: B256, header: Header) { self.headers.lock().insert(hash, header); } /// Add multiple headers to local header store pub fn extend_headers(&self, iter: impl IntoIterator) { for (hash, header) in iter { self.add_header(hash, header) } } /// Add account to local account store pub fn add_account(&self, address: Address, account: ExtendedAccount) { self.accounts.lock().insert(address, account); } /// Add account to local account store pub fn extend_accounts(&self, iter: impl IntoIterator) { for (address, account) in iter { self.add_account(address, account) } } /// Add state root to local state root store pub fn add_state_root(&self, state_root: B256) { self.state_roots.lock().push(state_root); } } impl HeaderProvider for MockEthProvider { fn header(&self, block_hash: &BlockHash) -> ProviderResult> { let lock = self.headers.lock(); Ok(lock.get(block_hash).cloned()) } fn header_by_number(&self, num: u64) -> ProviderResult> { let lock = self.headers.lock(); Ok(lock.values().find(|h| h.number == num).cloned()) } fn header_td(&self, hash: &BlockHash) -> ProviderResult> { let lock = self.headers.lock(); Ok(lock.get(hash).map(|target| { lock.values() .filter(|h| h.number < target.number) .fold(target.difficulty, |td, h| td + h.difficulty) })) } fn header_td_by_number(&self, number: BlockNumber) -> ProviderResult> { let lock = self.headers.lock(); let sum = lock .values() .filter(|h| h.number <= number) .fold(U256::ZERO, |td, h| td + h.difficulty); Ok(Some(sum)) } fn headers_range(&self, range: impl RangeBounds) -> ProviderResult> { let lock = self.headers.lock(); let mut headers: Vec<_> = lock.values().filter(|header| range.contains(&header.number)).cloned().collect(); headers.sort_by_key(|header| header.number); Ok(headers) } fn sealed_header(&self, number: BlockNumber) -> ProviderResult> { Ok(self.header_by_number(number)?.map(|h| h.seal_slow())) } fn sealed_headers_while( &self, range: impl RangeBounds, mut predicate: impl FnMut(&SealedHeader) -> bool, ) -> ProviderResult> { Ok(self .headers_range(range)? .into_iter() .map(|h| h.seal_slow()) .take_while(|h| predicate(h)) .collect()) } } impl ChainSpecProvider for MockEthProvider { fn chain_spec(&self) -> Arc { self.chain_spec.clone() } } impl TransactionsProvider for MockEthProvider { fn transaction_id(&self, tx_hash: TxHash) -> ProviderResult> { let lock = self.blocks.lock(); let tx_number = lock .values() .flat_map(|block| &block.body) .position(|tx| tx.hash() == tx_hash) .map(|pos| pos as TxNumber); Ok(tx_number) } fn transaction_by_id(&self, id: TxNumber) -> ProviderResult> { let lock = self.blocks.lock(); let transaction = lock.values().flat_map(|block| &block.body).nth(id as usize).cloned(); Ok(transaction) } fn transaction_by_id_no_hash( &self, id: TxNumber, ) -> ProviderResult> { let lock = self.blocks.lock(); let transaction = lock .values() .flat_map(|block| &block.body) .nth(id as usize) .map(|tx| Into::::into(tx.clone())); Ok(transaction) } fn transaction_by_hash(&self, hash: TxHash) -> ProviderResult> { Ok(self .blocks .lock() .iter() .find_map(|(_, block)| block.body.iter().find(|tx| tx.hash() == hash).cloned())) } fn transaction_by_hash_with_meta( &self, hash: TxHash, ) -> ProviderResult> { let lock = self.blocks.lock(); for (block_hash, block) in lock.iter() { for (index, tx) in block.body.iter().enumerate() { if tx.hash() == hash { let meta = TransactionMeta { tx_hash: hash, index: index as u64, block_hash: *block_hash, block_number: block.header.number, base_fee: block.header.base_fee_per_gas, excess_blob_gas: block.header.excess_blob_gas, timestamp: block.header.timestamp, }; return Ok(Some((tx.clone(), meta))) } } } Ok(None) } fn transaction_block(&self, id: TxNumber) -> ProviderResult> { let lock = self.blocks.lock(); let mut current_tx_number: TxNumber = 0; for block in lock.values() { if current_tx_number + (block.body.len() as TxNumber) > id { return Ok(Some(block.header.number)) } current_tx_number += block.body.len() as TxNumber; } Ok(None) } fn transactions_by_block( &self, id: BlockHashOrNumber, ) -> ProviderResult>> { Ok(self.block(id)?.map(|b| b.body)) } fn transactions_by_block_range( &self, range: impl RangeBounds, ) -> ProviderResult>> { // init btreemap so we can return in order let mut map = BTreeMap::new(); for (_, block) in self.blocks.lock().iter() { if range.contains(&block.number) { map.insert(block.number, block.body.clone()); } } Ok(map.into_values().collect()) } fn transactions_by_tx_range( &self, range: impl RangeBounds, ) -> ProviderResult> { let lock = self.blocks.lock(); let transactions = lock .values() .flat_map(|block| &block.body) .enumerate() .filter_map(|(tx_number, tx)| { if range.contains(&(tx_number as TxNumber)) { Some(tx.clone().into()) } else { None } }) .collect(); Ok(transactions) } fn senders_by_tx_range( &self, range: impl RangeBounds, ) -> ProviderResult> { let lock = self.blocks.lock(); let transactions = lock .values() .flat_map(|block| &block.body) .enumerate() .filter_map(|(tx_number, tx)| { if range.contains(&(tx_number as TxNumber)) { Some(tx.recover_signer()?) } else { None } }) .collect(); Ok(transactions) } fn transaction_sender(&self, id: TxNumber) -> ProviderResult> { self.transaction_by_id(id).map(|tx_option| tx_option.map(|tx| tx.recover_signer().unwrap())) } } impl ReceiptProvider for MockEthProvider { fn receipt(&self, _id: TxNumber) -> ProviderResult> { Ok(None) } fn receipt_by_hash(&self, _hash: TxHash) -> ProviderResult> { Ok(None) } fn receipts_by_block(&self, _block: BlockHashOrNumber) -> ProviderResult>> { Ok(None) } fn receipts_by_tx_range( &self, _range: impl RangeBounds, ) -> ProviderResult> { Ok(vec![]) } } impl ReceiptProviderIdExt for MockEthProvider {} impl BlockHashReader for MockEthProvider { fn block_hash(&self, number: u64) -> ProviderResult> { let lock = self.blocks.lock(); let hash = lock.iter().find_map(|(hash, b)| (b.number == number).then_some(*hash)); Ok(hash) } fn canonical_hashes_range( &self, start: BlockNumber, end: BlockNumber, ) -> ProviderResult> { let range = start..end; let lock = self.blocks.lock(); let mut hashes: Vec<_> = lock.iter().filter(|(_, block)| range.contains(&block.number)).collect(); hashes.sort_by_key(|(_, block)| block.number); Ok(hashes.into_iter().map(|(hash, _)| *hash).collect()) } } impl BlockNumReader for MockEthProvider { fn chain_info(&self) -> ProviderResult { let best_block_number = self.best_block_number()?; let lock = self.headers.lock(); Ok(lock .iter() .find(|(_, header)| header.number == best_block_number) .map(|(hash, header)| ChainInfo { best_hash: *hash, best_number: header.number }) .unwrap_or_default()) } fn best_block_number(&self) -> ProviderResult { let lock = self.headers.lock(); lock.iter() .max_by_key(|h| h.1.number) .map(|(_, header)| header.number) .ok_or(ProviderError::BestBlockNotFound) } fn last_block_number(&self) -> ProviderResult { self.best_block_number() } fn block_number(&self, hash: B256) -> ProviderResult> { let lock = self.blocks.lock(); let num = lock.iter().find_map(|(h, b)| (*h == hash).then_some(b.number)); Ok(num) } } impl BlockIdReader for MockEthProvider { fn pending_block_num_hash(&self) -> ProviderResult> { Ok(None) } fn safe_block_num_hash(&self) -> ProviderResult> { Ok(None) } fn finalized_block_num_hash(&self) -> ProviderResult> { Ok(None) } } impl BlockReader for MockEthProvider { fn find_block_by_hash( &self, hash: B256, _source: BlockSource, ) -> ProviderResult> { self.block(hash.into()) } fn block(&self, id: BlockHashOrNumber) -> ProviderResult> { let lock = self.blocks.lock(); match id { BlockHashOrNumber::Hash(hash) => Ok(lock.get(&hash).cloned()), BlockHashOrNumber::Number(num) => Ok(lock.values().find(|b| b.number == num).cloned()), } } fn pending_block(&self) -> ProviderResult> { Ok(None) } fn pending_block_with_senders(&self) -> ProviderResult> { Ok(None) } fn pending_block_and_receipts(&self) -> ProviderResult)>> { Ok(None) } fn ommers(&self, _id: BlockHashOrNumber) -> ProviderResult>> { Ok(None) } fn block_body_indices(&self, _num: u64) -> ProviderResult> { Ok(None) } fn block_with_senders( &self, _id: BlockHashOrNumber, _transaction_kind: TransactionVariant, ) -> ProviderResult> { Ok(None) } fn sealed_block_with_senders( &self, _id: BlockHashOrNumber, _transaction_kind: TransactionVariant, ) -> ProviderResult> { Ok(None) } fn block_range(&self, range: RangeInclusive) -> ProviderResult> { let lock = self.blocks.lock(); let mut blocks: Vec<_> = lock.values().filter(|block| range.contains(&block.number)).cloned().collect(); blocks.sort_by_key(|block| block.number); Ok(blocks) } fn block_with_senders_range( &self, _range: RangeInclusive, ) -> ProviderResult> { Ok(vec![]) } fn sealed_block_with_senders_range( &self, _range: RangeInclusive, ) -> ProviderResult> { Ok(vec![]) } } impl BlockReaderIdExt for MockEthProvider { fn block_by_id(&self, id: BlockId) -> ProviderResult> { match id { BlockId::Number(num) => self.block_by_number_or_tag(num), BlockId::Hash(hash) => self.block_by_hash(hash.block_hash), } } fn sealed_header_by_id(&self, id: BlockId) -> ProviderResult> { self.header_by_id(id)?.map_or_else(|| Ok(None), |h| Ok(Some(h.seal_slow()))) } fn header_by_id(&self, id: BlockId) -> ProviderResult> { match self.block_by_id(id)? { None => Ok(None), Some(block) => Ok(Some(block.header)), } } fn ommers_by_id(&self, id: BlockId) -> ProviderResult>> { match id { BlockId::Number(num) => self.ommers_by_number_or_tag(num), BlockId::Hash(hash) => self.ommers(BlockHashOrNumber::Hash(hash.block_hash)), } } } impl AccountReader for MockEthProvider { fn basic_account(&self, address: Address) -> ProviderResult> { Ok(self.accounts.lock().get(&address).cloned().map(|a| a.account)) } } impl StageCheckpointReader for MockEthProvider { fn get_stage_checkpoint(&self, _id: StageId) -> ProviderResult> { Ok(None) } fn get_stage_checkpoint_progress(&self, _id: StageId) -> ProviderResult>> { Ok(None) } fn get_all_checkpoints(&self) -> ProviderResult> { Ok(vec![]) } } impl StateRootProvider for MockEthProvider { fn hashed_state_root(&self, _state: HashedPostState) -> ProviderResult { Ok(self.state_roots.lock().pop().unwrap_or_default()) } fn hashed_state_root_from_nodes( &self, _nodes: TrieUpdates, _hashed_state: HashedPostState, _prefix_sets: TriePrefixSetsMut, ) -> ProviderResult { Ok(self.state_roots.lock().pop().unwrap_or_default()) } fn hashed_state_root_with_updates( &self, _state: HashedPostState, ) -> ProviderResult<(B256, TrieUpdates)> { let state_root = self.state_roots.lock().pop().unwrap_or_default(); Ok((state_root, Default::default())) } fn hashed_storage_root( &self, _address: Address, _hashed_storage: HashedStorage, ) -> ProviderResult { Ok(B256::default()) } fn hashed_state_root_from_nodes_with_updates( &self, _nodes: TrieUpdates, _hashed_state: HashedPostState, _prefix_sets: TriePrefixSetsMut, ) -> ProviderResult<(B256, TrieUpdates)> { let state_root = self.state_roots.lock().pop().unwrap_or_default(); Ok((state_root, Default::default())) } } impl StateProofProvider for MockEthProvider { fn hashed_proof( &self, _hashed_state: HashedPostState, address: Address, _slots: &[B256], ) -> ProviderResult { Ok(AccountProof::new(address)) } fn witness( &self, _overlay: HashedPostState, _target: HashedPostState, ) -> ProviderResult> { Ok(HashMap::default()) } } impl StateProvider for MockEthProvider { fn storage( &self, account: Address, storage_key: StorageKey, ) -> ProviderResult> { let lock = self.accounts.lock(); Ok(lock.get(&account).and_then(|account| account.storage.get(&storage_key)).cloned()) } fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult> { let lock = self.accounts.lock(); Ok(lock.values().find_map(|account| { match (account.account.bytecode_hash.as_ref(), account.bytecode.as_ref()) { (Some(bytecode_hash), Some(bytecode)) if *bytecode_hash == code_hash => { Some(bytecode.clone()) } _ => None, } })) } } impl EvmEnvProvider for MockEthProvider { fn fill_env_at( &self, _cfg: &mut CfgEnvWithHandlerCfg, _block_env: &mut BlockEnv, _at: BlockHashOrNumber, _evm_config: EvmConfig, ) -> ProviderResult<()> where EvmConfig: ConfigureEvmEnv, { Ok(()) } fn fill_env_with_header( &self, _cfg: &mut CfgEnvWithHandlerCfg, _block_env: &mut BlockEnv, _header: &Header, _evm_config: EvmConfig, ) -> ProviderResult<()> where EvmConfig: ConfigureEvmEnv, { Ok(()) } fn fill_cfg_env_at( &self, _cfg: &mut CfgEnvWithHandlerCfg, _at: BlockHashOrNumber, _evm_config: EvmConfig, ) -> ProviderResult<()> where EvmConfig: ConfigureEvmEnv, { Ok(()) } fn fill_cfg_env_with_header( &self, _cfg: &mut CfgEnvWithHandlerCfg, _header: &Header, _evm_config: EvmConfig, ) -> ProviderResult<()> where EvmConfig: ConfigureEvmEnv, { Ok(()) } } impl StateProviderFactory for MockEthProvider { fn latest(&self) -> ProviderResult { Ok(Box::new(self.clone())) } fn history_by_block_number(&self, _block: BlockNumber) -> ProviderResult { Ok(Box::new(self.clone())) } fn history_by_block_hash(&self, _block: BlockHash) -> ProviderResult { Ok(Box::new(self.clone())) } fn state_by_block_hash(&self, _block: BlockHash) -> ProviderResult { Ok(Box::new(self.clone())) } fn state_by_block_number_or_tag( &self, number_or_tag: BlockNumberOrTag, ) -> ProviderResult { match number_or_tag { BlockNumberOrTag::Latest => self.latest(), BlockNumberOrTag::Finalized => { // we can only get the finalized state by hash, not by num let hash = self.finalized_block_hash()?.ok_or(ProviderError::FinalizedBlockNotFound)?; // only look at historical state self.history_by_block_hash(hash) } BlockNumberOrTag::Safe => { // we can only get the safe state by hash, not by num let hash = self.safe_block_hash()?.ok_or(ProviderError::SafeBlockNotFound)?; self.history_by_block_hash(hash) } BlockNumberOrTag::Earliest => self.history_by_block_number(0), BlockNumberOrTag::Pending => self.pending(), BlockNumberOrTag::Number(num) => self.history_by_block_number(num), } } fn pending(&self) -> ProviderResult { Ok(Box::new(self.clone())) } fn pending_state_by_hash(&self, _block_hash: B256) -> ProviderResult> { Ok(Some(Box::new(self.clone()))) } } impl WithdrawalsProvider for MockEthProvider { fn withdrawals_by_block( &self, _id: BlockHashOrNumber, _timestamp: u64, ) -> ProviderResult> { Ok(None) } fn latest_withdrawal(&self) -> ProviderResult> { Ok(None) } } impl RequestsProvider for MockEthProvider { fn requests_by_block( &self, _id: BlockHashOrNumber, _timestamp: u64, ) -> ProviderResult> { Ok(None) } } impl ChangeSetReader for MockEthProvider { fn account_block_changeset( &self, _block_number: BlockNumber, ) -> ProviderResult> { Ok(Vec::default()) } }