reth_stateless/
root.rs

1// Copied and modified from ress: https://github.com/paradigmxyz/ress/blob/06bf2c4788e45b8fcbd640e38b6243e6f87c4d0e/crates/engine/src/tree/root.rs
2
3use alloc::vec::Vec;
4use alloy_primitives::B256;
5use alloy_rlp::{Decodable, Encodable};
6use itertools::Itertools;
7use reth_trie_common::{
8    HashedPostState, Nibbles, TrieAccount, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
9};
10use reth_trie_sparse::{errors::SparseStateTrieResult, SparseStateTrie, SparseTrie};
11
12/// Calculates the post-execution state root by applying state changes to a sparse trie.
13///
14/// This function takes a [`SparseStateTrie`] with the pre-state and a [`HashedPostState`]
15/// containing account and storage changes resulting from block execution (state diff).
16///
17/// It modifies the input `trie` in place to reflect these changes and then calculates the
18/// final post-execution state root.
19pub(crate) fn calculate_state_root(
20    trie: &mut SparseStateTrie,
21    state: HashedPostState,
22) -> SparseStateTrieResult<B256> {
23    // 1. Apply storage‑slot updates and compute each contract’s storage root
24    //
25    //
26    // We walk over every (address, storage) pair in deterministic order
27    // and update the corresponding per‑account storage trie in‑place.
28    // When we’re done we collect (address, updated_storage_trie) in a `Vec`
29    // so that we can insert them back into the outer state trie afterwards ― this avoids
30    // borrowing issues.
31    let mut storage_results = Vec::with_capacity(state.storages.len());
32
33    for (address, storage) in state.storages.into_iter().sorted_unstable_by_key(|(addr, _)| *addr) {
34        // Take the existing storage trie (or create an empty, “revealed” one)
35        let mut storage_trie =
36            trie.take_storage_trie(&address).unwrap_or_else(SparseTrie::revealed_empty);
37
38        if storage.wiped {
39            storage_trie.wipe()?;
40        }
41
42        // Apply slot‑level changes
43        for (hashed_slot, value) in
44            storage.storage.into_iter().sorted_unstable_by_key(|(slot, _)| *slot)
45        {
46            let nibbles = Nibbles::unpack(hashed_slot);
47            if value.is_zero() {
48                storage_trie.remove_leaf(&nibbles)?;
49            } else {
50                storage_trie.update_leaf(
51                    nibbles,
52                    alloy_rlp::encode_fixed_size(&value).to_vec(),
53                    value.is_private,
54                )?;
55            }
56        }
57
58        // Finalise the storage‑trie root before pushing the result
59        storage_trie.root();
60        storage_results.push((address, storage_trie));
61    }
62
63    // Insert every updated storage trie back into the outer state trie
64    for (address, storage_trie) in storage_results {
65        trie.insert_storage_trie(address, storage_trie);
66    }
67
68    // 2. Apply account‑level updates and (re)encode the account nodes
69    // Update accounts with new values
70    // TODO: upstream changes into reth so that `SparseStateTrie::update_account` handles this
71    let mut account_rlp_buf = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
72
73    for (hashed_address, account) in
74        state.accounts.into_iter().sorted_unstable_by_key(|(addr, _)| *addr)
75    {
76        let nibbles = Nibbles::unpack(hashed_address);
77        let account = account.unwrap_or_default();
78
79        // Determine which storage root should be used for this account
80        let storage_root = if let Some(storage_trie) = trie.storage_trie_mut(&hashed_address) {
81            storage_trie.root()
82        } else if let Some(value) = trie.get_account_value(&hashed_address) {
83            TrieAccount::decode(&mut &value[..])?.storage_root
84        } else {
85            EMPTY_ROOT_HASH
86        };
87
88        // Decide whether to remove or update the account leaf
89        if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
90            trie.remove_account_leaf(&nibbles)?;
91        } else {
92            account_rlp_buf.clear();
93            account.into_trie_account(storage_root).encode(&mut account_rlp_buf);
94            trie.update_account_leaf(nibbles, account_rlp_buf.clone())?;
95        }
96    }
97
98    // Return new state root
99    trie.root()
100}