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#[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 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 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}