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