reth_revm/
test_utils.rs

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