reth_revm/
test_utils.rs

1use alloc::vec::Vec;
2use alloy_primitives::{
3    keccak256, map::HashMap, Address, BlockNumber, Bytes, StorageKey, B256, U256,
4};
5use reth_primitives_traits::{Account, Bytecode};
6use reth_storage_api::{
7    AccountReader, BlockHashReader, HashedPostStateProvider, StateProofProvider, StateProvider,
8    StateRootProvider, StorageRootProvider,
9};
10use reth_storage_errors::provider::ProviderResult;
11use reth_trie::{
12    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, KeccakKeyHasher,
13    MultiProof, MultiProofTargets, StorageMultiProof, StorageProof, TrieInput,
14};
15
16/// Mock state for testing
17#[derive(Debug, Default, Clone, Eq, PartialEq)]
18pub struct StateProviderTest {
19    accounts: HashMap<Address, (HashMap<StorageKey, U256>, Account)>,
20    contracts: HashMap<B256, Bytecode>,
21    block_hash: HashMap<u64, B256>,
22}
23
24impl StateProviderTest {
25    /// Insert account.
26    pub fn insert_account(
27        &mut self,
28        address: Address,
29        mut account: Account,
30        bytecode: Option<Bytes>,
31        storage: HashMap<StorageKey, U256>,
32    ) {
33        if let Some(bytecode) = bytecode {
34            let hash = keccak256(&bytecode);
35            account.bytecode_hash = Some(hash);
36            self.contracts.insert(hash, Bytecode::new_raw(bytecode));
37        }
38        self.accounts.insert(address, (storage, account));
39    }
40
41    /// Insert a block hash.
42    pub fn insert_block_hash(&mut self, block_number: u64, block_hash: B256) {
43        self.block_hash.insert(block_number, block_hash);
44    }
45}
46
47impl AccountReader for StateProviderTest {
48    fn basic_account(&self, address: &Address) -> ProviderResult<Option<Account>> {
49        Ok(self.accounts.get(address).map(|(_, acc)| *acc))
50    }
51}
52
53impl BlockHashReader for StateProviderTest {
54    fn block_hash(&self, number: u64) -> ProviderResult<Option<B256>> {
55        Ok(self.block_hash.get(&number).copied())
56    }
57
58    fn canonical_hashes_range(
59        &self,
60        start: BlockNumber,
61        end: BlockNumber,
62    ) -> ProviderResult<Vec<B256>> {
63        let range = start..end;
64        Ok(self
65            .block_hash
66            .iter()
67            .filter_map(|(block, hash)| range.contains(block).then_some(*hash))
68            .collect())
69    }
70}
71
72impl StateRootProvider for StateProviderTest {
73    fn state_root(&self, _hashed_state: HashedPostState) -> ProviderResult<B256> {
74        unimplemented!("state root computation is not supported")
75    }
76
77    fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
78        unimplemented!("state root computation is not supported")
79    }
80
81    fn state_root_with_updates(
82        &self,
83        _hashed_state: HashedPostState,
84    ) -> ProviderResult<(B256, TrieUpdates)> {
85        unimplemented!("state root computation is not supported")
86    }
87
88    fn state_root_from_nodes_with_updates(
89        &self,
90        _input: TrieInput,
91    ) -> ProviderResult<(B256, TrieUpdates)> {
92        unimplemented!("state root computation is not supported")
93    }
94}
95
96impl StorageRootProvider for StateProviderTest {
97    fn storage_root(
98        &self,
99        _address: Address,
100        _hashed_storage: HashedStorage,
101    ) -> ProviderResult<B256> {
102        unimplemented!("storage root is not supported")
103    }
104
105    fn storage_proof(
106        &self,
107        _address: Address,
108        _slot: B256,
109        _hashed_storage: HashedStorage,
110    ) -> ProviderResult<StorageProof> {
111        unimplemented!("proof generation is not supported")
112    }
113
114    fn storage_multiproof(
115        &self,
116        _address: Address,
117        _slots: &[B256],
118        _hashed_storage: HashedStorage,
119    ) -> ProviderResult<StorageMultiProof> {
120        unimplemented!("proof generation is not supported")
121    }
122}
123
124impl StateProofProvider for StateProviderTest {
125    fn proof(
126        &self,
127        _input: TrieInput,
128        _address: Address,
129        _slots: &[B256],
130    ) -> ProviderResult<AccountProof> {
131        unimplemented!("proof generation is not supported")
132    }
133
134    fn multiproof(
135        &self,
136        _input: TrieInput,
137        _targets: MultiProofTargets,
138    ) -> ProviderResult<MultiProof> {
139        unimplemented!("proof generation is not supported")
140    }
141
142    fn witness(&self, _input: TrieInput, _target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
143        unimplemented!("witness generation is not supported")
144    }
145}
146
147impl HashedPostStateProvider for StateProviderTest {
148    fn hashed_post_state(&self, bundle_state: &revm::database::BundleState) -> HashedPostState {
149        HashedPostState::from_bundle_state::<KeccakKeyHasher>(bundle_state.state())
150    }
151}
152
153impl StateProvider for StateProviderTest {
154    fn storage(
155        &self,
156        account: Address,
157        storage_key: StorageKey,
158    ) -> ProviderResult<Option<alloy_primitives::StorageValue>> {
159        Ok(self.accounts.get(&account).and_then(|(storage, _)| storage.get(&storage_key).copied()))
160    }
161
162    fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
163        Ok(self.contracts.get(code_hash).cloned())
164    }
165}