1use crate::validation::StatelessValidationError;
2use alloc::{format, vec::Vec};
3use alloy_primitives::{keccak256, map::B256Map, Address, B256, U256};
4use alloy_rlp::{Decodable, Encodable};
5use alloy_rpc_types_debug::ExecutionWitness;
6use alloy_trie::{TrieAccount, EMPTY_ROOT_HASH};
7use itertools::Itertools;
8use reth_errors::ProviderError;
9use reth_revm::state::Bytecode;
10use reth_trie_common::{HashedPostState, Nibbles, TRIE_ACCOUNT_RLP_MAX_SIZE};
11use reth_trie_sparse::{
12 errors::SparseStateTrieResult,
13 provider::{DefaultTrieNodeProvider, DefaultTrieNodeProviderFactory},
14 SparseStateTrie, SparseTrie, SparseTrieInterface,
15};
16
17pub trait StatelessTrie: core::fmt::Debug {
19 fn new(
21 witness: &ExecutionWitness,
22 pre_state_root: B256,
23 ) -> Result<(Self, B256Map<Bytecode>), StatelessValidationError>
24 where
25 Self: Sized;
26
27 fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError>;
32
33 fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError>;
38
39 fn calculate_state_root(
41 &mut self,
42 state: HashedPostState,
43 ) -> Result<B256, StatelessValidationError>;
44}
45
46#[derive(Debug)]
48pub struct StatelessSparseTrie {
49 inner: SparseStateTrie,
50}
51
52impl StatelessSparseTrie {
53 pub fn new(
58 witness: &ExecutionWitness,
59 pre_state_root: B256,
60 ) -> Result<(Self, B256Map<Bytecode>), StatelessValidationError> {
61 verify_execution_witness(witness, pre_state_root)
62 .map(|(inner, bytecode)| (Self { inner }, bytecode))
63 }
64
65 pub fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError> {
70 let hashed_address = keccak256(address);
71
72 if let Some(bytes) = self.inner.get_account_value(&hashed_address) {
73 let account = TrieAccount::decode(&mut bytes.as_slice())?;
74 return Ok(Some(account))
75 }
76
77 if !self.inner.check_valid_account_witness(hashed_address) {
78 return Err(ProviderError::TrieWitnessError(format!(
79 "incomplete account witness for {hashed_address:?}"
80 )));
81 }
82
83 Ok(None)
84 }
85
86 pub fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError> {
91 let hashed_address = keccak256(address);
92 let hashed_slot = keccak256(B256::from(slot));
93
94 if let Some(raw) = self.inner.get_storage_slot_value(&hashed_address, &hashed_slot) {
95 return Ok(U256::decode(&mut raw.as_slice())?)
96 }
97
98 if let Some(bytes) = self.inner.get_account_value(&hashed_address) {
101 let account = TrieAccount::decode(&mut bytes.as_slice())?;
104 if account.storage_root != EMPTY_ROOT_HASH &&
105 !self.inner.check_valid_storage_witness(hashed_address, hashed_slot)
106 {
107 return Err(ProviderError::TrieWitnessError(format!(
108 "incomplete storage witness: prover must supply exclusion proof for slot {hashed_slot:?} in account {hashed_address:?}"
109 )));
110 }
111 } else if !self.inner.check_valid_account_witness(hashed_address) {
112 return Err(ProviderError::TrieWitnessError(format!(
115 "incomplete account witness for {hashed_address:?}"
116 )));
117 }
118
119 Ok(U256::ZERO)
120 }
121
122 pub fn calculate_state_root(
124 &mut self,
125 state: HashedPostState,
126 ) -> Result<B256, StatelessValidationError> {
127 calculate_state_root(&mut self.inner, state)
128 .map_err(|_e| StatelessValidationError::StatelessStateRootCalculationFailed)
129 }
130}
131
132impl StatelessTrie for StatelessSparseTrie {
133 fn new(
134 witness: &ExecutionWitness,
135 pre_state_root: B256,
136 ) -> Result<(Self, B256Map<Bytecode>), StatelessValidationError> {
137 Self::new(witness, pre_state_root)
138 }
139
140 fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError> {
141 self.account(address)
142 }
143
144 fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError> {
145 self.storage(address, slot)
146 }
147
148 fn calculate_state_root(
149 &mut self,
150 state: HashedPostState,
151 ) -> Result<B256, StatelessValidationError> {
152 self.calculate_state_root(state)
153 }
154}
155
156fn verify_execution_witness(
175 witness: &ExecutionWitness,
176 pre_state_root: B256,
177) -> Result<(SparseStateTrie, B256Map<Bytecode>), StatelessValidationError> {
178 let provider_factory = DefaultTrieNodeProviderFactory;
179 let mut trie = SparseStateTrie::new();
180 let mut state_witness = B256Map::default();
181 let mut bytecode = B256Map::default();
182
183 for rlp_encoded in &witness.state {
184 let hash = keccak256(rlp_encoded);
185 state_witness.insert(hash, rlp_encoded.clone());
186 }
187 for rlp_encoded in &witness.codes {
188 let hash = keccak256(rlp_encoded);
189 bytecode.insert(hash, Bytecode::new_raw(rlp_encoded.clone()));
190 }
191
192 trie.reveal_witness(pre_state_root, &state_witness)
201 .map_err(|_e| StatelessValidationError::WitnessRevealFailed { pre_state_root })?;
202
203 let computed_root = trie
205 .root(&provider_factory)
206 .map_err(|_e| StatelessValidationError::StatelessPreStateRootCalculationFailed)?;
207
208 if computed_root == pre_state_root {
209 Ok((trie, bytecode))
210 } else {
211 Err(StatelessValidationError::PreStateRootMismatch {
212 got: computed_root,
213 expected: pre_state_root,
214 })
215 }
216}
217
218fn calculate_state_root(
227 trie: &mut SparseStateTrie,
228 state: HashedPostState,
229) -> SparseStateTrieResult<B256> {
230 let mut storage_results = Vec::with_capacity(state.storages.len());
239
240 let provider_factory = DefaultTrieNodeProviderFactory;
243 let storage_provider = DefaultTrieNodeProvider;
244
245 for (address, storage) in state.storages.into_iter().sorted_unstable_by_key(|(addr, _)| *addr) {
246 let mut storage_trie =
248 trie.take_storage_trie(&address).unwrap_or_else(SparseTrie::revealed_empty);
249
250 if storage.wiped {
251 storage_trie.wipe()?;
252 }
253
254 for (hashed_slot, value) in
256 storage.storage.into_iter().sorted_unstable_by_key(|(slot, _)| *slot)
257 {
258 let nibbles = Nibbles::unpack(hashed_slot);
259 if value.is_zero() {
260 storage_trie.remove_leaf(&nibbles, &storage_provider)?;
261 } else {
262 storage_trie.update_leaf(
263 nibbles,
264 alloy_rlp::encode_fixed_size(&value).to_vec(),
265 &storage_provider,
266 )?;
267 }
268 }
269
270 storage_trie.root();
272 storage_results.push((address, storage_trie));
273 }
274
275 for (address, storage_trie) in storage_results {
277 trie.insert_storage_trie(address, storage_trie);
278 }
279
280 let mut account_rlp_buf = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
284
285 for (hashed_address, account) in
286 state.accounts.into_iter().sorted_unstable_by_key(|(addr, _)| *addr)
287 {
288 let nibbles = Nibbles::unpack(hashed_address);
289 let account = account.unwrap_or_default();
290
291 let storage_root = if let Some(storage_trie) = trie.storage_trie_mut(&hashed_address) {
293 storage_trie.root()
294 } else if let Some(value) = trie.get_account_value(&hashed_address) {
295 TrieAccount::decode(&mut &value[..])?.storage_root
296 } else {
297 EMPTY_ROOT_HASH
298 };
299
300 if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
302 trie.remove_account_leaf(&nibbles, &provider_factory)?;
303 } else {
304 account_rlp_buf.clear();
305 account.into_trie_account(storage_root).encode(&mut account_rlp_buf);
306 trie.update_account_leaf(nibbles, account_rlp_buf.clone(), &provider_factory)?;
307 }
308 }
309
310 trie.root(&provider_factory)
312}