reth_trie_db/
state.rs

Help
1use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory, PrefixSetLoader};
2use alloy_primitives::{
3    map::{AddressMap, B256Map},
4    Address, BlockNumber, B256, U256,
5};
6use reth_db_api::{
7    cursor::DbCursorRO,
8    models::{AccountBeforeTx, BlockNumberAddress},
9    tables,
10    transaction::DbTx,
11    DatabaseError,
12};
13use reth_execution_errors::StateRootError;
14use reth_trie::{
15    hashed_cursor::HashedPostStateCursorFactory, trie_cursor::InMemoryTrieCursorFactory,
16    updates::TrieUpdates, HashedPostState, HashedStorage, KeccakKeyHasher, KeyHasher, StateRoot,
17    StateRootProgress, TrieInput,
18};
19use std::{collections::HashMap, ops::RangeInclusive};
20use tracing::debug;
21
22/// Extends [`StateRoot`] with operations specific for working with a database transaction.
23pub trait DatabaseStateRoot<'a, TX>: Sized {
24    /// Create a new [`StateRoot`] instance.
25    fn from_tx(tx: &'a TX) -> Self;
26
27    /// Given a block number range, identifies all the accounts and storage keys that
28    /// have changed.
29    ///
30    /// # Returns
31    ///
32    /// An instance of state root calculator with account and storage prefixes loaded.
33    fn incremental_root_calculator(
34        tx: &'a TX,
35        range: RangeInclusive<BlockNumber>,
36    ) -> Result<Self, StateRootError>;
37
38    /// Computes the state root of the trie with the changed account and storage prefixes and
39    /// existing trie nodes.
40    ///
41    /// # Returns
42    ///
43    /// The updated state root.
44    fn incremental_root(
45        tx: &'a TX,
46        range: RangeInclusive<BlockNumber>,
47    ) -> Result<B256, StateRootError>;
48
49    /// Computes the state root of the trie with the changed account and storage prefixes and
50    /// existing trie nodes collecting updates in the process.
51    ///
52    /// Ignores the threshold.
53    ///
54    /// # Returns
55    ///
56    /// The updated state root and the trie updates.
57    fn incremental_root_with_updates(
58        tx: &'a TX,
59        range: RangeInclusive<BlockNumber>,
60    ) -> Result<(B256, TrieUpdates), StateRootError>;
61
62    /// Computes the state root of the trie with the changed account and storage prefixes and
63    /// existing trie nodes collecting updates in the process.
64    ///
65    /// # Returns
66    ///
67    /// The intermediate progress of state root computation.
68    fn incremental_root_with_progress(
69        tx: &'a TX,
70        range: RangeInclusive<BlockNumber>,
71    ) -> Result<StateRootProgress, StateRootError>;
72
73    /// Calculate the state root for this [`HashedPostState`].
74    /// Internally, this method retrieves prefixsets and uses them
75    /// to calculate incremental state root.
76    ///
77    /// # Example
78    ///
79    /// ```
80    /// use alloy_primitives::U256;
81    /// use reth_db::test_utils::create_test_rw_db;
82    /// use reth_db_api::database::Database;
83    /// use reth_primitives_traits::Account;
84    /// use reth_trie::{updates::TrieUpdates, HashedPostState, StateRoot};
85    /// use reth_trie_db::DatabaseStateRoot;
86    ///
87    /// // Initialize the database
88    /// let db = create_test_rw_db();
89    ///
90    /// // Initialize hashed post state
91    /// let mut hashed_state = HashedPostState::default();
92    /// hashed_state.accounts.insert(
93    ///     [0x11; 32].into(),
94    ///     Some(Account { nonce: 1, balance: U256::from(10), bytecode_hash: None }),
95    /// );
96    ///
97    /// // Calculate the state root
98    /// let tx = db.tx().expect("failed to create transaction");
99    /// let state_root = StateRoot::overlay_root(&tx, hashed_state);
100    /// ```
101    ///
102    /// # Returns
103    ///
104    /// The state root for this [`HashedPostState`].
105    fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result<B256, StateRootError>;
106
107    /// Calculates the state root for this [`HashedPostState`] and returns it alongside trie
108    /// updates. See [`Self::overlay_root`] for more info.
109    fn overlay_root_with_updates(
110        tx: &'a TX,
111        post_state: HashedPostState,
112    ) -> Result<(B256, TrieUpdates), StateRootError>;
113
114    /// Calculates the state root for provided [`HashedPostState`] using cached intermediate nodes.
115    fn overlay_root_from_nodes(tx: &'a TX, input: TrieInput) -> Result<B256, StateRootError>;
116
117    /// Calculates the state root and trie updates for provided [`HashedPostState`] using
118    /// cached intermediate nodes.
119    fn overlay_root_from_nodes_with_updates(
120        tx: &'a TX,
121        input: TrieInput,
122    ) -> Result<(B256, TrieUpdates), StateRootError>;
123}
124
125/// Extends [`HashedPostState`] with operations specific for working with a database transaction.
126pub trait DatabaseHashedPostState<TX>: Sized {
127    /// Initializes [`HashedPostState`] from reverts. Iterates over state reverts from the specified
128    /// block up to the current tip and aggregates them into hashed state in reverse.
129    fn from_reverts<KH: KeyHasher>(tx: &TX, from: BlockNumber) -> Result<Self, DatabaseError>;
130}
131
132impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
133    for StateRoot<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>
134{
135    fn from_tx(tx: &'a TX) -> Self {
136        Self::new(DatabaseTrieCursorFactory::new(tx), DatabaseHashedCursorFactory::new(tx))
137    }
138
139    fn incremental_root_calculator(
140        tx: &'a TX,
141        range: RangeInclusive<BlockNumber>,
142    ) -> Result<Self, StateRootError> {
143        let loaded_prefix_sets = PrefixSetLoader::<_, KeccakKeyHasher>::new(tx).load(range)?;
144        Ok(Self::from_tx(tx).with_prefix_sets(loaded_prefix_sets))
145    }
146
147    fn incremental_root(
148        tx: &'a TX,
149        range: RangeInclusive<BlockNumber>,
150    ) -> Result<B256, StateRootError> {
151        debug!(target: "trie::loader", ?range, "incremental state root");
152        Self::incremental_root_calculator(tx, range)?.root()
153    }
154
155    fn incremental_root_with_updates(
156        tx: &'a TX,
157        range: RangeInclusive<BlockNumber>,
158    ) -> Result<(B256, TrieUpdates), StateRootError> {
159        debug!(target: "trie::loader", ?range, "incremental state root");
160        Self::incremental_root_calculator(tx, range)?.root_with_updates()
161    }
162
163    fn incremental_root_with_progress(
164        tx: &'a TX,
165        range: RangeInclusive<BlockNumber>,
166    ) -> Result<StateRootProgress, StateRootError> {
167        debug!(target: "trie::loader", ?range, "incremental state root with progress");
168        Self::incremental_root_calculator(tx, range)?.root_with_progress()
169    }
170
171    fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result<B256, StateRootError> {
172        let prefix_sets = post_state.construct_prefix_sets().freeze();
173        let state_sorted = post_state.into_sorted();
174        StateRoot::new(
175            DatabaseTrieCursorFactory::new(tx),
176            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
177        )
178        .with_prefix_sets(prefix_sets)
179        .root()
180    }
181
182    fn overlay_root_with_updates(
183        tx: &'a TX,
184        post_state: HashedPostState,
185    ) -> Result<(B256, TrieUpdates), StateRootError> {
186        let prefix_sets = post_state.construct_prefix_sets().freeze();
187        let state_sorted = post_state.into_sorted();
188        StateRoot::new(
189            DatabaseTrieCursorFactory::new(tx),
190            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
191        )
192        .with_prefix_sets(prefix_sets)
193        .root_with_updates()
194    }
195
196    fn overlay_root_from_nodes(tx: &'a TX, input: TrieInput) -> Result<B256, StateRootError> {
197        let state_sorted = input.state.into_sorted();
198        let nodes_sorted = input.nodes.into_sorted();
199        StateRoot::new(
200            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
201            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
202        )
203        .with_prefix_sets(input.prefix_sets.freeze())
204        .root()
205    }
206
207    fn overlay_root_from_nodes_with_updates(
208        tx: &'a TX,
209        input: TrieInput,
210    ) -> Result<(B256, TrieUpdates), StateRootError> {
211        let state_sorted = input.state.into_sorted();
212        let nodes_sorted = input.nodes.into_sorted();
213        StateRoot::new(
214            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
215            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
216        )
217        .with_prefix_sets(input.prefix_sets.freeze())
218        .root_with_updates()
219    }
220}
221
222impl<TX: DbTx> DatabaseHashedPostState<TX> for HashedPostState {
223    fn from_reverts<KH: KeyHasher>(tx: &TX, from: BlockNumber) -> Result<Self, DatabaseError> {
224        // Iterate over account changesets and record value before first occurring account change.
225        let mut accounts = HashMap::new();
226        let mut account_changesets_cursor = tx.cursor_read::<tables::AccountChangeSets>()?;
227        for entry in account_changesets_cursor.walk_range(from..)? {
228            let (_, AccountBeforeTx { address, info }) = entry?;
229            accounts.entry(address).or_insert(info);
230        }
231
232        // Iterate over storage changesets and record value before first occurring storage change.
233        let mut storages = AddressMap::<B256Map<U256>>::default();
234        let mut storage_changesets_cursor = tx.cursor_read::<tables::StorageChangeSets>()?;
235        for entry in
236            storage_changesets_cursor.walk_range(BlockNumberAddress((from, Address::ZERO))..)?
237        {
238            let (BlockNumberAddress((_, address)), storage) = entry?;
239            let account_storage = storages.entry(address).or_default();
240            account_storage.entry(storage.key).or_insert(storage.value);
241        }
242
243        let hashed_accounts =
244            accounts.into_iter().map(|(address, info)| (KH::hash_key(address), info)).collect();
245
246        let hashed_storages = storages
247            .into_iter()
248            .map(|(address, storage)| {
249                (
250                    KH::hash_key(address),
251                    HashedStorage::from_iter(
252                        // The `wiped` flag indicates only whether previous storage entries
253                        // should be looked up in db or not. For reverts it's a noop since all
254                        // wiped changes had been written as storage reverts.
255                        false,
256                        storage.into_iter().map(|(slot, value)| (KH::hash_key(slot), value)),
257                    ),
258                )
259            })
260            .collect();
261
262        Ok(Self { accounts: hashed_accounts, storages: hashed_storages })
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use alloy_primitives::{hex, map::HashMap, Address, U256};
270    use reth_db::test_utils::create_test_rw_db;
271    use reth_db_api::database::Database;
272    use reth_trie::KeccakKeyHasher;
273    use revm::state::AccountInfo;
274    use revm_database::BundleState;
275
276    #[test]
277    fn from_bundle_state_with_rayon() {
278        let address1 = Address::with_last_byte(1);
279        let address2 = Address::with_last_byte(2);
280        let slot1 = U256::from(1015);
281        let slot2 = U256::from(2015);
282
283        let account1 = AccountInfo { nonce: 1, ..Default::default() };
284        let account2 = AccountInfo { nonce: 2, ..Default::default() };
285
286        let bundle_state = BundleState::builder(2..=2)
287            .state_present_account_info(address1, account1)
288            .state_present_account_info(address2, account2)
289            .state_storage(address1, HashMap::from_iter([(slot1, (U256::ZERO, U256::from(10)))]))
290            .state_storage(address2, HashMap::from_iter([(slot2, (U256::ZERO, U256::from(20)))]))
291            .build();
292        assert_eq!(bundle_state.reverts.len(), 1);
293
294        let post_state = HashedPostState::from_bundle_state::<KeccakKeyHasher>(&bundle_state.state);
295        assert_eq!(post_state.accounts.len(), 2);
296        assert_eq!(post_state.storages.len(), 2);
297
298        let db = create_test_rw_db();
299        let tx = db.tx().expect("failed to create transaction");
300        assert_eq!(
301            StateRoot::overlay_root(&tx, post_state).unwrap(),
302            hex!("b464525710cafcf5d4044ac85b72c08b1e76231b8d91f288fe438cc41d8eaafd")
303        );
304    }
305}
reth_trie_db/state.rs

reth_trie_db/
state.rs
