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