reth_stateless/
trie.rs

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
17/// Trait for stateless trie implementations that can be used for stateless validation.
18pub trait StatelessTrie: core::fmt::Debug {
19    /// Initialize the stateless trie using the `ExecutionWitness`
20    fn new(
21        witness: &ExecutionWitness,
22        pre_state_root: B256,
23    ) -> Result<(Self, B256Map<Bytecode>), StatelessValidationError>
24    where
25        Self: Sized;
26
27    /// Returns the `TrieAccount` that corresponds to the `Address`
28    ///
29    /// This method will error if the `ExecutionWitness` is not able to guarantee
30    /// that the account is missing from the Trie _and_ the witness was complete.
31    fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError>;
32
33    /// Returns the storage slot value that corresponds to the given (address, slot) tuple.
34    ///
35    /// This method will error if the `ExecutionWitness` is not able to guarantee
36    /// that the storage was missing from the Trie _and_ the witness was complete.
37    fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError>;
38
39    /// Computes the new state root from the `HashedPostState`.
40    fn calculate_state_root(
41        &mut self,
42        state: HashedPostState,
43    ) -> Result<B256, StatelessValidationError>;
44}
45
46/// `StatelessSparseTrie` structure for usage during stateless validation
47#[derive(Debug)]
48pub struct StatelessSparseTrie {
49    inner: SparseStateTrie,
50}
51
52impl StatelessSparseTrie {
53    /// Initialize the stateless trie using the `ExecutionWitness`
54    ///
55    /// Note: Currently this method does not check that the `ExecutionWitness`
56    /// is complete for all of the preimage keys.
57    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    /// Returns the `TrieAccount` that corresponds to the `Address`
66    ///
67    /// This method will error if the `ExecutionWitness` is not able to guarantee
68    /// that the account is missing from the Trie _and_ the witness was complete.
69    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    /// Returns the storage slot value that corresponds to the given (address, slot) tuple.
87    ///
88    /// This method will error if the `ExecutionWitness` is not able to guarantee
89    /// that the storage was missing from the Trie _and_ the witness was complete.
90    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        // Storage slot value is not present in the trie, validate that the witness is complete.
99        // If the account exists in the trie...
100        if let Some(bytes) = self.inner.get_account_value(&hashed_address) {
101            // ...check that its storage is either empty or the storage trie was sufficiently
102            // revealed...
103            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            // ...else if account is missing, validate that the account trie was sufficiently
113            // revealed.
114            return Err(ProviderError::TrieWitnessError(format!(
115                "incomplete account witness for {hashed_address:?}"
116            )));
117        }
118
119        Ok(U256::ZERO)
120    }
121
122    /// Computes the new state root from the `HashedPostState`.
123    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
156/// Verifies execution witness [`ExecutionWitness`] against an expected pre-state root.
157///
158/// This function takes the RLP-encoded values provided in [`ExecutionWitness`]
159/// (which includes state trie nodes, storage trie nodes, and contract bytecode)
160/// and uses it to populate a new [`SparseStateTrie`].
161///
162/// If the computed root hash matches the `pre_state_root`, it signifies that the
163/// provided execution witness is consistent with that pre-state root. In this case, the function
164/// returns the populated [`SparseStateTrie`] and a [`B256Map`] containing the
165/// contract bytecode (mapping code hash to [`Bytecode`]).
166///
167/// The bytecode has a separate mapping because the [`SparseStateTrie`] does not store the
168/// contract bytecode, only the hash of it (code hash).
169///
170/// If the roots do not match, it returns `None`, indicating the witness is invalid
171/// for the given `pre_state_root`.
172// Note: This approach might be inefficient for ZKVMs requiring minimal memory operations, which
173// would explain why they have for the most part re-implemented this function.
174fn 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    // Reveal the witness with our state root
193    // This method builds a trie using the sparse trie using the state_witness with
194    // the root being the pre_state_root.
195    // Here are some things to note:
196    // - You can pass in more witnesses than is needed for the block execution.
197    // - If you try to get an account and it has not been seen. This means that the account
198    // was not inserted into the Trie. It does not mean that the account does not exist.
199    // In order to determine an account not existing, we must do an exclusion proof.
200    trie.reveal_witness(pre_state_root, &state_witness)
201        .map_err(|_e| StatelessValidationError::WitnessRevealFailed { pre_state_root })?;
202
203    // Calculate the root
204    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
218// Copied and modified from ress: https://github.com/paradigmxyz/ress/blob/06bf2c4788e45b8fcbd640e38b6243e6f87c4d0e/crates/engine/src/tree/root.rs
219/// Calculates the post-execution state root by applying state changes to a sparse trie.
220///
221/// This function takes a [`SparseStateTrie`] with the pre-state and a [`HashedPostState`]
222/// containing account and storage changes resulting from block execution (state diff).
223///
224/// It modifies the input `trie` in place to reflect these changes and then calculates the
225/// final post-execution state root.
226fn calculate_state_root(
227    trie: &mut SparseStateTrie,
228    state: HashedPostState,
229) -> SparseStateTrieResult<B256> {
230    // 1. Apply storage‑slot updates and compute each contract’s storage root
231    //
232    //
233    // We walk over every (address, storage) pair in deterministic order
234    // and update the corresponding per‑account storage trie in‑place.
235    // When we’re done we collect (address, updated_storage_trie) in a `Vec`
236    // so that we can insert them back into the outer state trie afterwards ― this avoids
237    // borrowing issues.
238    let mut storage_results = Vec::with_capacity(state.storages.len());
239
240    // In `verify_execution_witness` a `DefaultTrieNodeProviderFactory` is used, so we use the same
241    // again in here.
242    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        // Take the existing storage trie (or create an empty, “revealed” one)
247        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        // Apply slot‑level changes
255        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        // Finalise the storage‑trie root before pushing the result
271        storage_trie.root();
272        storage_results.push((address, storage_trie));
273    }
274
275    // Insert every updated storage trie back into the outer state trie
276    for (address, storage_trie) in storage_results {
277        trie.insert_storage_trie(address, storage_trie);
278    }
279
280    // 2. Apply account‑level updates and (re)encode the account nodes
281    // Update accounts with new values
282    // TODO: upstream changes into reth so that `SparseStateTrie::update_account` handles this
283    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        // Determine which storage root should be used for this account
292        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        // Decide whether to remove or update the account leaf
301        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    // Return new state root
311    trie.root(&provider_factory)
312}