Skip to main content

reth_trie/
trie.rs

1use crate::{
2    hashed_cursor::{HashedCursor, HashedCursorFactory, HashedStorageCursor},
3    node_iter::{TrieElement, TrieNodeIter},
4    prefix_set::{PrefixSet, TriePrefixSets},
5    progress::{
6        IntermediateRootState, IntermediateStateRootState, IntermediateStorageRootState,
7        StateRootProgress, StorageRootProgress,
8    },
9    stats::TrieTracker,
10    trie_cursor::{TrieCursor, TrieCursorFactory, TrieStorageCursor},
11    updates::{StorageTrieUpdates, TrieUpdates},
12    walker::TrieWalker,
13    HashBuilder, Nibbles, TRIE_ACCOUNT_RLP_MAX_SIZE,
14};
15use alloy_consensus::EMPTY_ROOT_HASH;
16use alloy_primitives::{keccak256, Address, B256, U256};
17use alloy_rlp::{BufMut, Encodable};
18use alloy_trie::proof::AddedRemovedKeys;
19use reth_execution_errors::{StateRootError, StorageRootError};
20use reth_primitives_traits::Account;
21use tracing::{debug, instrument, trace, Span};
22
23/// The default updates after which root algorithms should return intermediate progress rather than
24/// finishing the computation.
25const DEFAULT_INTERMEDIATE_THRESHOLD: u64 = 100_000;
26
27#[cfg(feature = "metrics")]
28use crate::metrics::{StateRootMetrics, TrieRootMetrics};
29
30/// `StateRoot` is used to compute the root node of a state trie.
31#[derive(Debug)]
32pub struct StateRoot<T, H> {
33    /// The factory for trie cursors.
34    pub trie_cursor_factory: T,
35    /// The factory for hashed cursors.
36    pub hashed_cursor_factory: H,
37    /// A set of prefix sets that have changed.
38    pub prefix_sets: TriePrefixSets,
39    /// Whether every child under a branch whose path matches the prefix set should be walked.
40    walk_all_changed_branch_children: bool,
41    /// Previous intermediate state.
42    previous_state: Option<IntermediateStateRootState>,
43    /// The number of updates after which the intermediate progress should be returned.
44    threshold: u64,
45    #[cfg(feature = "metrics")]
46    /// State root metrics.
47    metrics: StateRootMetrics,
48}
49
50impl<T, H> StateRoot<T, H> {
51    /// Creates [`StateRoot`] with `trie_cursor_factory` and `hashed_cursor_factory`. All other
52    /// parameters are set to reasonable defaults.
53    ///
54    /// The cursors created by given factories are then used to walk through the accounts and
55    /// calculate the state root value with.
56    pub fn new(trie_cursor_factory: T, hashed_cursor_factory: H) -> Self {
57        Self {
58            trie_cursor_factory,
59            hashed_cursor_factory,
60            prefix_sets: TriePrefixSets::default(),
61            walk_all_changed_branch_children: false,
62            previous_state: None,
63            threshold: DEFAULT_INTERMEDIATE_THRESHOLD,
64            #[cfg(feature = "metrics")]
65            metrics: StateRootMetrics::default(),
66        }
67    }
68
69    /// Set the prefix sets.
70    pub fn with_prefix_sets(mut self, prefix_sets: TriePrefixSets) -> Self {
71        self.prefix_sets = prefix_sets;
72        self
73    }
74
75    /// Configures the state root walker to visit all children of changed branch paths.
76    pub const fn with_walk_all_changed_branch_children(mut self, enabled: bool) -> Self {
77        self.walk_all_changed_branch_children = enabled;
78        self
79    }
80
81    /// Set the threshold.
82    pub const fn with_threshold(mut self, threshold: u64) -> Self {
83        self.threshold = threshold;
84        self
85    }
86
87    /// Set the threshold to maximum value so that intermediate progress is not returned.
88    pub const fn with_no_threshold(mut self) -> Self {
89        self.threshold = u64::MAX;
90        self
91    }
92
93    /// Set the previously recorded intermediate state.
94    pub fn with_intermediate_state(mut self, state: Option<IntermediateStateRootState>) -> Self {
95        self.previous_state = state;
96        self
97    }
98
99    /// Set the hashed cursor factory.
100    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StateRoot<T, HF> {
101        StateRoot {
102            trie_cursor_factory: self.trie_cursor_factory,
103            hashed_cursor_factory,
104            prefix_sets: self.prefix_sets,
105            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
106            threshold: self.threshold,
107            previous_state: self.previous_state,
108            #[cfg(feature = "metrics")]
109            metrics: self.metrics,
110        }
111    }
112
113    /// Set the trie cursor factory.
114    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StateRoot<TF, H> {
115        StateRoot {
116            trie_cursor_factory,
117            hashed_cursor_factory: self.hashed_cursor_factory,
118            prefix_sets: self.prefix_sets,
119            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
120            threshold: self.threshold,
121            previous_state: self.previous_state,
122            #[cfg(feature = "metrics")]
123            metrics: self.metrics,
124        }
125    }
126}
127
128impl<T, H> StateRoot<T, H>
129where
130    T: TrieCursorFactory,
131    H: HashedCursorFactory,
132{
133    /// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
134    /// nodes into the hash builder. Collects the updates in the process.
135    ///
136    /// Ignores the threshold.
137    ///
138    /// # Returns
139    ///
140    /// The state root and the trie updates.
141    pub fn root_with_updates(self) -> Result<(B256, TrieUpdates), StateRootError> {
142        match self.with_no_threshold().calculate(true)? {
143            StateRootProgress::Complete(root, _, updates) => Ok((root, updates)),
144            StateRootProgress::Progress(..) => unreachable!(), // unreachable threshold
145        }
146    }
147
148    /// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
149    /// nodes into the hash builder.
150    ///
151    /// # Returns
152    ///
153    /// The state root hash.
154    pub fn root(self) -> Result<B256, StateRootError> {
155        match self.calculate(false)? {
156            StateRootProgress::Complete(root, _, _) => Ok(root),
157            StateRootProgress::Progress(..) => unreachable!(), // update retention is disabled
158        }
159    }
160
161    /// Walks the intermediate nodes of existing state trie (if any) and hashed entries. Feeds the
162    /// nodes into the hash builder. Collects the updates in the process.
163    ///
164    /// # Returns
165    ///
166    /// The intermediate progress of state root computation.
167    pub fn root_with_progress(self) -> Result<StateRootProgress, StateRootError> {
168        self.calculate(true)
169    }
170
171    fn calculate(self, retain_updates: bool) -> Result<StateRootProgress, StateRootError> {
172        trace!(target: "trie::state_root", "calculating state root");
173        let mut tracker = TrieTracker::default();
174
175        let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
176        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
177
178        // Shared storage cursors for reuse across all storage root calculations.
179        // Created lazily on first use to avoid issues with mock cursors.
180        let mut hashed_storage_cursor: Option<H::StorageCursor<'_>> = None;
181        let mut storage_trie_cursor: Option<T::StorageTrieCursor<'_>> = None;
182
183        // create state root context once for reuse
184        let mut storage_ctx = StateRootContext::new();
185
186        // first handle any in-progress storage root calculation
187        let (mut hash_builder, mut account_node_iter) = if let Some(state) = self.previous_state {
188            let IntermediateStateRootState { account_root_state, storage_root_state } = state;
189
190            // resume account trie iteration
191            let mut hash_builder = account_root_state.hash_builder.with_updates(retain_updates);
192            let walker = TrieWalker::<_>::state_trie_from_stack(
193                trie_cursor,
194                account_root_state.walker_stack,
195                self.prefix_sets.account_prefix_set,
196            )
197            .with_walk_all_changed_branch_children(self.walk_all_changed_branch_children)
198            .with_deletions_retained(retain_updates);
199            let account_node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor)
200                .with_last_hashed_key(account_root_state.last_hashed_key);
201
202            // if we have an in-progress storage root, complete it first
203            if let Some(storage_state) = storage_root_state {
204                let hashed_address = account_root_state.last_hashed_key;
205                let account = storage_state.account;
206
207                debug!(
208                    target: "trie::state_root",
209                    account_nonce = account.nonce,
210                    account_balance = ?account.balance,
211                    last_hashed_key = ?account_root_state.last_hashed_key,
212                    "Resuming storage root calculation"
213                );
214
215                let remaining_threshold = self.threshold.saturating_sub(
216                    storage_ctx.total_updates_len(&account_node_iter, &hash_builder),
217                );
218
219                if hashed_storage_cursor.is_none() {
220                    hashed_storage_cursor =
221                        Some(self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?);
222                }
223                if storage_trie_cursor.is_none() {
224                    storage_trie_cursor =
225                        Some(self.trie_cursor_factory.storage_trie_cursor(hashed_address)?);
226                }
227
228                let storage_result = StorageRoot::<T, H>::calculate_with_cursors(
229                    StorageRootCalculation {
230                        hashed_address,
231                        prefix_set: self
232                            .prefix_sets
233                            .storage_prefix_sets
234                            .get(&hashed_address)
235                            .cloned()
236                            .unwrap_or_default(),
237                        previous_state: Some(storage_state.state),
238                        walk_all_changed_branch_children: self.walk_all_changed_branch_children,
239                        threshold: remaining_threshold,
240                        retain_updates,
241                    },
242                    storage_trie_cursor.as_mut().expect("storage trie cursor is initialized"),
243                    hashed_storage_cursor.as_mut().expect("hashed storage cursor is initialized"),
244                    #[cfg(feature = "metrics")]
245                    &self.metrics.storage_trie,
246                )?;
247                if let Some(storage_state) = storage_ctx.process_storage_root_result(
248                    storage_result,
249                    hashed_address,
250                    account,
251                    &mut hash_builder,
252                    retain_updates,
253                )? {
254                    // still in progress, need to pause again
255                    return Ok(storage_ctx.create_progress_state(
256                        account_node_iter,
257                        hash_builder,
258                        account_root_state.last_hashed_key,
259                        Some(storage_state),
260                    ))
261                }
262            }
263
264            (hash_builder, account_node_iter)
265        } else {
266            // no intermediate state, create new hash builder and node iter for state root
267            // calculation
268            let hash_builder = HashBuilder::default().with_updates(retain_updates);
269            let walker = TrieWalker::state_trie(trie_cursor, self.prefix_sets.account_prefix_set)
270                .with_walk_all_changed_branch_children(self.walk_all_changed_branch_children)
271                .with_deletions_retained(retain_updates);
272            let node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor);
273            (hash_builder, node_iter)
274        };
275
276        while let Some(node) = account_node_iter.try_next()? {
277            match node {
278                TrieElement::Branch(node) => {
279                    tracker.inc_branch();
280                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
281                }
282                TrieElement::Leaf(hashed_address, account) => {
283                    tracker.inc_leaf();
284                    storage_ctx.hashed_entries_walked += 1;
285
286                    // calculate storage root, calculating the remaining threshold so we have
287                    // bounded memory usage even while in the middle of storage root calculation
288                    let remaining_threshold = self.threshold.saturating_sub(
289                        storage_ctx.total_updates_len(&account_node_iter, &hash_builder),
290                    );
291
292                    if hashed_storage_cursor.is_none() {
293                        hashed_storage_cursor =
294                            Some(self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?);
295                    }
296                    if storage_trie_cursor.is_none() {
297                        storage_trie_cursor =
298                            Some(self.trie_cursor_factory.storage_trie_cursor(hashed_address)?);
299                    }
300
301                    let storage_result = StorageRoot::<T, H>::calculate_with_cursors(
302                        StorageRootCalculation {
303                            hashed_address,
304                            prefix_set: self
305                                .prefix_sets
306                                .storage_prefix_sets
307                                .get(&hashed_address)
308                                .cloned()
309                                .unwrap_or_default(),
310                            previous_state: None,
311                            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
312                            threshold: remaining_threshold,
313                            retain_updates,
314                        },
315                        storage_trie_cursor.as_mut().expect("storage trie cursor is initialized"),
316                        hashed_storage_cursor
317                            .as_mut()
318                            .expect("hashed storage cursor is initialized"),
319                        #[cfg(feature = "metrics")]
320                        &self.metrics.storage_trie,
321                    )?;
322                    if let Some(storage_state) = storage_ctx.process_storage_root_result(
323                        storage_result,
324                        hashed_address,
325                        account,
326                        &mut hash_builder,
327                        retain_updates,
328                    )? {
329                        // storage root hit threshold, need to pause
330                        return Ok(storage_ctx.create_progress_state(
331                            account_node_iter,
332                            hash_builder,
333                            hashed_address,
334                            Some(storage_state),
335                        ))
336                    }
337
338                    // decide if we need to return intermediate progress
339                    let total_updates_len =
340                        storage_ctx.total_updates_len(&account_node_iter, &hash_builder);
341                    if retain_updates && total_updates_len >= self.threshold {
342                        return Ok(storage_ctx.create_progress_state(
343                            account_node_iter,
344                            hash_builder,
345                            hashed_address,
346                            None,
347                        ))
348                    }
349                }
350            }
351        }
352
353        let root = hash_builder.root();
354
355        let removed_keys = account_node_iter.walker.take_removed_keys();
356        let StateRootContext { mut trie_updates, hashed_entries_walked, .. } = storage_ctx;
357        trie_updates.finalize(hash_builder, removed_keys, self.prefix_sets.destroyed_accounts);
358
359        let stats = tracker.finish();
360
361        #[cfg(feature = "metrics")]
362        self.metrics.state_trie.record(stats);
363
364        trace!(
365            target: "trie::state_root",
366            %root,
367            duration = ?stats.duration(),
368            branches_added = stats.branches_added(),
369            leaves_added = stats.leaves_added(),
370            "calculated state root"
371        );
372
373        Ok(StateRootProgress::Complete(root, hashed_entries_walked, trie_updates))
374    }
375}
376
377/// Contains state mutated during state root calculation and storage root result handling.
378#[derive(Debug)]
379pub(crate) struct StateRootContext {
380    /// Reusable buffer for encoding account data.
381    account_rlp: Vec<u8>,
382    /// Accumulates updates from account and storage root calculation.
383    trie_updates: TrieUpdates,
384    /// Tracks total hashed entries walked.
385    hashed_entries_walked: usize,
386    /// Counts storage trie nodes updated.
387    updated_storage_nodes: usize,
388}
389
390impl StateRootContext {
391    /// Creates a new state root context.
392    fn new() -> Self {
393        Self {
394            account_rlp: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
395            trie_updates: TrieUpdates::default(),
396            hashed_entries_walked: 0,
397            updated_storage_nodes: 0,
398        }
399    }
400
401    /// Creates a [`StateRootProgress`] when the threshold is hit, from the state of the current
402    /// [`TrieNodeIter`], [`HashBuilder`], last hashed key and any storage root intermediate state.
403    fn create_progress_state<C, H, K>(
404        mut self,
405        account_node_iter: TrieNodeIter<C, H, K>,
406        hash_builder: HashBuilder,
407        last_hashed_key: B256,
408        storage_state: Option<IntermediateStorageRootState>,
409    ) -> StateRootProgress
410    where
411        C: TrieCursor,
412        H: HashedCursor,
413        K: AsRef<AddedRemovedKeys>,
414    {
415        let (walker_stack, walker_deleted_keys) = account_node_iter.walker.split();
416        self.trie_updates.removed_nodes.extend(walker_deleted_keys);
417        let (hash_builder, hash_builder_updates) = hash_builder.split();
418        self.trie_updates.account_nodes.extend(hash_builder_updates);
419
420        let account_state = IntermediateRootState { hash_builder, walker_stack, last_hashed_key };
421
422        let state = IntermediateStateRootState {
423            account_root_state: account_state,
424            storage_root_state: storage_state,
425        };
426
427        StateRootProgress::Progress(Box::new(state), self.hashed_entries_walked, self.trie_updates)
428    }
429
430    /// Calculates the total number of updated nodes.
431    fn total_updates_len<C, H, K>(
432        &self,
433        account_node_iter: &TrieNodeIter<C, H, K>,
434        hash_builder: &HashBuilder,
435    ) -> u64
436    where
437        C: TrieCursor,
438        H: HashedCursor,
439        K: AsRef<AddedRemovedKeys>,
440    {
441        (self.updated_storage_nodes +
442            account_node_iter.walker.removed_keys_len() +
443            hash_builder.updates_len()) as u64
444    }
445
446    /// Processes the result of a storage root calculation.
447    ///
448    /// Handles both completed and in-progress storage root calculations:
449    /// - For completed roots: encodes the account with the storage root, updates the hash builder
450    ///   with the new account, and updates metrics.
451    /// - For in-progress roots: returns the intermediate state for later resumption
452    ///
453    /// Returns an [`IntermediateStorageRootState`] if the calculation needs to be resumed later, or
454    /// `None` if the storage root was successfully computed and added to the trie.
455    fn process_storage_root_result(
456        &mut self,
457        storage_result: StorageRootProgress,
458        hashed_address: B256,
459        account: Account,
460        hash_builder: &mut HashBuilder,
461        retain_updates: bool,
462    ) -> Result<Option<IntermediateStorageRootState>, StateRootError> {
463        match storage_result {
464            StorageRootProgress::Complete(storage_root, storage_slots_walked, updates) => {
465                // Storage root completed
466                self.hashed_entries_walked += storage_slots_walked;
467                if retain_updates {
468                    self.updated_storage_nodes += updates.len();
469                    self.trie_updates.insert_storage_updates(hashed_address, updates);
470                }
471
472                // Encode the account with the computed storage root
473                self.account_rlp.clear();
474                let trie_account = account.into_trie_account(storage_root);
475                trie_account.encode(&mut self.account_rlp as &mut dyn BufMut);
476                hash_builder.add_leaf(Nibbles::unpack(hashed_address), &self.account_rlp);
477                Ok(None)
478            }
479            StorageRootProgress::Progress(state, storage_slots_walked, updates) => {
480                // Storage root hit threshold or resumed calculation hit threshold
481                debug!(
482                    target: "trie::state_root",
483                    ?hashed_address,
484                    storage_slots_walked,
485                    last_storage_key = ?state.last_hashed_key,
486                    ?account,
487                    "Pausing storage root calculation"
488                );
489
490                self.hashed_entries_walked += storage_slots_walked;
491                if retain_updates {
492                    self.trie_updates.insert_storage_updates(hashed_address, updates);
493                }
494
495                Ok(Some(IntermediateStorageRootState { state: *state, account }))
496            }
497        }
498    }
499}
500
501/// `StorageRoot` is used to compute the root node of an account storage trie.
502#[derive(Debug)]
503pub struct StorageRoot<T, H> {
504    /// A reference to the database transaction.
505    pub trie_cursor_factory: T,
506    /// The factory for hashed cursors.
507    pub hashed_cursor_factory: H,
508    /// The hashed address of an account.
509    pub hashed_address: B256,
510    /// The set of storage slot prefixes that have changed.
511    pub prefix_set: PrefixSet,
512    /// Whether every child under a branch whose path matches the prefix set should be walked.
513    walk_all_changed_branch_children: bool,
514    /// Previous intermediate state.
515    previous_state: Option<IntermediateRootState>,
516    /// The number of updates after which the intermediate progress should be returned.
517    threshold: u64,
518    /// Storage root metrics.
519    #[cfg(feature = "metrics")]
520    metrics: TrieRootMetrics,
521}
522
523impl<T, H> StorageRoot<T, H> {
524    /// Creates a new storage root calculator given a raw address.
525    pub fn new(
526        trie_cursor_factory: T,
527        hashed_cursor_factory: H,
528        address: Address,
529        prefix_set: PrefixSet,
530        #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
531    ) -> Self {
532        Self::new_hashed(
533            trie_cursor_factory,
534            hashed_cursor_factory,
535            keccak256(address),
536            prefix_set,
537            #[cfg(feature = "metrics")]
538            metrics,
539        )
540    }
541
542    /// Creates a new storage root calculator given a hashed address.
543    pub const fn new_hashed(
544        trie_cursor_factory: T,
545        hashed_cursor_factory: H,
546        hashed_address: B256,
547        prefix_set: PrefixSet,
548        #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
549    ) -> Self {
550        Self {
551            trie_cursor_factory,
552            hashed_cursor_factory,
553            hashed_address,
554            prefix_set,
555            walk_all_changed_branch_children: false,
556            previous_state: None,
557            threshold: DEFAULT_INTERMEDIATE_THRESHOLD,
558            #[cfg(feature = "metrics")]
559            metrics,
560        }
561    }
562
563    /// Set the changed prefixes.
564    pub fn with_prefix_set(mut self, prefix_set: PrefixSet) -> Self {
565        self.prefix_set = prefix_set;
566        self
567    }
568
569    /// Configures the storage root walker to visit all children of changed branch paths.
570    pub const fn with_walk_all_changed_branch_children(mut self, enabled: bool) -> Self {
571        self.walk_all_changed_branch_children = enabled;
572        self
573    }
574
575    /// Set the threshold.
576    pub const fn with_threshold(mut self, threshold: u64) -> Self {
577        self.threshold = threshold;
578        self
579    }
580
581    /// Set the threshold to maximum value so that intermediate progress is not returned.
582    pub const fn with_no_threshold(mut self) -> Self {
583        self.threshold = u64::MAX;
584        self
585    }
586
587    /// Set the previously recorded intermediate state.
588    pub fn with_intermediate_state(mut self, state: Option<IntermediateRootState>) -> Self {
589        self.previous_state = state;
590        self
591    }
592
593    /// Set the hashed cursor factory.
594    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageRoot<T, HF> {
595        StorageRoot {
596            trie_cursor_factory: self.trie_cursor_factory,
597            hashed_cursor_factory,
598            hashed_address: self.hashed_address,
599            prefix_set: self.prefix_set,
600            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
601            previous_state: self.previous_state,
602            threshold: self.threshold,
603            #[cfg(feature = "metrics")]
604            metrics: self.metrics,
605        }
606    }
607
608    /// Set the trie cursor factory.
609    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageRoot<TF, H> {
610        StorageRoot {
611            trie_cursor_factory,
612            hashed_cursor_factory: self.hashed_cursor_factory,
613            hashed_address: self.hashed_address,
614            prefix_set: self.prefix_set,
615            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
616            previous_state: self.previous_state,
617            threshold: self.threshold,
618            #[cfg(feature = "metrics")]
619            metrics: self.metrics,
620        }
621    }
622}
623
624impl<T, H> StorageRoot<T, H>
625where
626    T: TrieCursorFactory,
627    H: HashedCursorFactory,
628{
629    /// Walks the intermediate nodes of existing storage trie (if any) and hashed entries. Feeds the
630    /// nodes into the hash builder. Collects the updates in the process.
631    ///
632    /// # Returns
633    ///
634    /// The intermediate progress of state root computation.
635    pub fn root_with_progress(self) -> Result<StorageRootProgress, StorageRootError> {
636        self.calculate(true)
637    }
638
639    /// Walks the hashed storage table entries for a given address and calculates the storage root.
640    ///
641    /// # Returns
642    ///
643    /// The storage root and storage trie updates for a given address.
644    pub fn root_with_updates(self) -> Result<(B256, usize, StorageTrieUpdates), StorageRootError> {
645        match self.with_no_threshold().calculate(true)? {
646            StorageRootProgress::Complete(root, walked, updates) => Ok((root, walked, updates)),
647            StorageRootProgress::Progress(..) => unreachable!(), // unreachable threshold
648        }
649    }
650
651    /// Walks the hashed storage table entries for a given address and calculates the storage root.
652    ///
653    /// # Returns
654    ///
655    /// The storage root.
656    pub fn root(self) -> Result<B256, StorageRootError> {
657        match self.calculate(false)? {
658            StorageRootProgress::Complete(root, _, _) => Ok(root),
659            StorageRootProgress::Progress(..) => unreachable!(), // update retention is disabled
660        }
661    }
662
663    /// Walks the hashed storage table entries for a given address and calculates the storage root.
664    ///
665    /// # Returns
666    ///
667    /// The storage root, number of walked entries and trie updates
668    /// for a given address if requested.
669    #[instrument(skip_all, level = "trace", target = "trie::storage_root", name = "storage_trie", fields(addr = %self.hashed_address, storage_root))]
670    pub fn calculate(self, retain_updates: bool) -> Result<StorageRootProgress, StorageRootError> {
671        trace!(target: "trie::storage_root", "calculating storage root");
672
673        let Self {
674            trie_cursor_factory,
675            hashed_cursor_factory,
676            hashed_address,
677            prefix_set,
678            walk_all_changed_branch_children,
679            previous_state,
680            threshold,
681            #[cfg(feature = "metrics")]
682            metrics,
683        } = self;
684
685        let mut hashed_storage_cursor =
686            hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
687        let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(hashed_address)?;
688
689        Self::calculate_with_cursors(
690            StorageRootCalculation {
691                hashed_address,
692                prefix_set,
693                previous_state,
694                walk_all_changed_branch_children,
695                threshold,
696                retain_updates,
697            },
698            &mut storage_trie_cursor,
699            &mut hashed_storage_cursor,
700            #[cfg(feature = "metrics")]
701            &metrics,
702        )
703    }
704
705    /// Walks the hashed storage table entries for a given address and calculates the storage root
706    /// using a pre-created cursor. The cursor will be repositioned to the given hashed address.
707    ///
708    /// This method allows reusing a single cursor across multiple storage root calculations,
709    /// reducing overhead when computing storage roots for many accounts.
710    #[instrument(skip_all, level = "trace", target = "trie::storage_root", name = "storage_trie_with_cursor", fields(addr = %self.hashed_address, storage_root))]
711    pub fn calculate_with_cursor(
712        self,
713        hashed_storage_cursor: &mut H::StorageCursor<'_>,
714        retain_updates: bool,
715    ) -> Result<StorageRootProgress, StorageRootError> {
716        trace!(target: "trie::storage_root", "calculating storage root with cursor");
717
718        let Self {
719            trie_cursor_factory,
720            hashed_address,
721            prefix_set,
722            walk_all_changed_branch_children,
723            previous_state,
724            threshold,
725            #[cfg(feature = "metrics")]
726            metrics,
727            ..
728        } = self;
729
730        let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(hashed_address)?;
731
732        Self::calculate_with_cursors(
733            StorageRootCalculation {
734                hashed_address,
735                prefix_set,
736                previous_state,
737                walk_all_changed_branch_children,
738                threshold,
739                retain_updates,
740            },
741            &mut storage_trie_cursor,
742            hashed_storage_cursor,
743            #[cfg(feature = "metrics")]
744            &metrics,
745        )
746    }
747
748    /// Walks the hashed storage table entries for a given address and calculates the storage root
749    /// using pre-created cursors. The cursors will be repositioned to the given hashed address.
750    fn calculate_with_cursors<TC, HC>(
751        calculation: StorageRootCalculation,
752        trie_cursor: &mut TC,
753        hashed_storage_cursor: &mut HC,
754        #[cfg(feature = "metrics")] metrics: &TrieRootMetrics,
755    ) -> Result<StorageRootProgress, StorageRootError>
756    where
757        TC: TrieStorageCursor,
758        HC: HashedStorageCursor<Value = U256>,
759    {
760        let StorageRootCalculation {
761            hashed_address,
762            prefix_set,
763            previous_state,
764            walk_all_changed_branch_children,
765            threshold,
766            retain_updates,
767        } = calculation;
768        hashed_storage_cursor.set_hashed_address(hashed_address);
769
770        // Empty storage only needs to be walked when changed prefixes must produce trie updates.
771        if previous_state.is_none() &&
772            (!retain_updates || prefix_set.is_empty()) &&
773            hashed_storage_cursor.is_storage_empty()?
774        {
775            Span::current().record("storage_root", format!("{EMPTY_ROOT_HASH:?}"));
776            return Ok(StorageRootProgress::Complete(
777                EMPTY_ROOT_HASH,
778                0,
779                StorageTrieUpdates::default(),
780            ))
781        }
782
783        trie_cursor.set_hashed_address(hashed_address);
784
785        let mut tracker = TrieTracker::default();
786        let mut trie_updates = StorageTrieUpdates::default();
787
788        let (mut hash_builder, mut storage_node_iter) = match previous_state {
789            Some(state) => {
790                let hash_builder = state.hash_builder.with_updates(retain_updates);
791                let walker = TrieWalker::<_>::storage_trie_from_stack(
792                    trie_cursor,
793                    state.walker_stack,
794                    prefix_set,
795                )
796                .with_walk_all_changed_branch_children(walk_all_changed_branch_children)
797                .with_deletions_retained(retain_updates);
798                let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor)
799                    .with_last_hashed_key(state.last_hashed_key);
800                (hash_builder, node_iter)
801            }
802            None => {
803                let hash_builder = HashBuilder::default().with_updates(retain_updates);
804                let walker = TrieWalker::storage_trie(trie_cursor, prefix_set)
805                    .with_walk_all_changed_branch_children(walk_all_changed_branch_children)
806                    .with_deletions_retained(retain_updates);
807                let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor);
808                (hash_builder, node_iter)
809            }
810        };
811
812        let mut hashed_entries_walked = 0;
813        while let Some(node) = storage_node_iter.try_next()? {
814            match node {
815                TrieElement::Branch(node) => {
816                    tracker.inc_branch();
817                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
818                }
819                TrieElement::Leaf(hashed_slot, value) => {
820                    tracker.inc_leaf();
821                    hashed_entries_walked += 1;
822                    hash_builder.add_leaf(
823                        Nibbles::unpack(hashed_slot),
824                        alloy_rlp::encode_fixed_size(&value).as_ref(),
825                    );
826
827                    // Check if we need to return intermediate progress
828                    let total_updates_len =
829                        storage_node_iter.walker.removed_keys_len() + hash_builder.updates_len();
830                    if retain_updates && total_updates_len as u64 >= threshold {
831                        let (walker_stack, walker_deleted_keys) = storage_node_iter.walker.split();
832                        trie_updates.removed_nodes.extend(walker_deleted_keys);
833                        let (hash_builder, hash_builder_updates) = hash_builder.split();
834                        trie_updates.storage_nodes.extend(hash_builder_updates);
835
836                        let state = IntermediateRootState {
837                            hash_builder,
838                            walker_stack,
839                            last_hashed_key: hashed_slot,
840                        };
841
842                        return Ok(StorageRootProgress::Progress(
843                            Box::new(state),
844                            hashed_entries_walked,
845                            trie_updates,
846                        ))
847                    }
848                }
849            }
850        }
851
852        let root = hash_builder.root();
853        Span::current().record("storage_root", format!("{root:?}"));
854
855        let removed_keys = storage_node_iter.walker.take_removed_keys();
856        trie_updates.finalize(hash_builder, removed_keys);
857
858        let stats = tracker.finish();
859
860        #[cfg(feature = "metrics")]
861        metrics.record(stats);
862
863        trace!(
864            target: "trie::storage_root",
865            %root,
866            %hashed_address,
867            duration = ?stats.duration(),
868            branches_added = stats.branches_added(),
869            leaves_added = stats.leaves_added(),
870            "calculated storage root"
871        );
872
873        let storage_slots_walked = stats.leaves_added() as usize;
874        Ok(StorageRootProgress::Complete(root, storage_slots_walked, trie_updates))
875    }
876}
877
878/// Parameters for a storage root calculation using pre-created cursors.
879struct StorageRootCalculation {
880    hashed_address: B256,
881    prefix_set: PrefixSet,
882    previous_state: Option<IntermediateRootState>,
883    walk_all_changed_branch_children: bool,
884    threshold: u64,
885    retain_updates: bool,
886}
887
888/// Trie type for differentiating between various trie calculations.
889#[derive(Clone, Copy, Debug)]
890pub enum TrieType {
891    /// State trie type.
892    State,
893    /// Storage trie type.
894    Storage,
895    /// Custom trie type. Can be used in ExEx.
896    Custom(&'static str),
897}
898
899impl TrieType {
900    #[cfg(feature = "metrics")]
901    pub(crate) const fn as_str(&self) -> &'static str {
902        match self {
903            Self::State => "state",
904            Self::Storage => "storage",
905            Self::Custom(s) => s,
906        }
907    }
908}