Skip to main content

reth_revm/
test_utils.rs

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