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, TrieCursorIter, 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, map::B256Map, 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 mut destroyed_storage_trie_nodes = B256Map::default();
356        if retain_updates && !self.prefix_sets.destroyed_accounts.is_empty() {
357            let mut storage_trie_cursor = match storage_trie_cursor {
358                Some(cursor) => cursor,
359                None => self.trie_cursor_factory.storage_trie_cursor(B256::ZERO)?,
360            };
361
362            for hashed_address in &self.prefix_sets.destroyed_accounts {
363                storage_trie_cursor.set_hashed_address(*hashed_address);
364                let nodes = TrieCursorIter::new(&mut storage_trie_cursor)
365                    .map(|node| node.map(|(path, _)| path))
366                    .collect::<Result<Vec<_>, _>>()?;
367
368                if !nodes.is_empty() {
369                    destroyed_storage_trie_nodes.insert(*hashed_address, nodes);
370                }
371            }
372        }
373
374        let removed_keys = account_node_iter.walker.take_removed_keys();
375        let StateRootContext { mut trie_updates, hashed_entries_walked, .. } = storage_ctx;
376        trie_updates.finalize(hash_builder, removed_keys, destroyed_storage_trie_nodes);
377
378        let stats = tracker.finish();
379
380        #[cfg(feature = "metrics")]
381        self.metrics.state_trie.record(stats);
382
383        trace!(
384            target: "trie::state_root",
385            %root,
386            duration = ?stats.duration(),
387            branches_added = stats.branches_added(),
388            leaves_added = stats.leaves_added(),
389            "calculated state root"
390        );
391
392        Ok(StateRootProgress::Complete(root, hashed_entries_walked, trie_updates))
393    }
394}
395
396/// Contains state mutated during state root calculation and storage root result handling.
397#[derive(Debug)]
398pub(crate) struct StateRootContext {
399    /// Reusable buffer for encoding account data.
400    account_rlp: Vec<u8>,
401    /// Accumulates updates from account and storage root calculation.
402    trie_updates: TrieUpdates,
403    /// Tracks total hashed entries walked.
404    hashed_entries_walked: usize,
405    /// Counts storage trie nodes updated.
406    updated_storage_nodes: usize,
407}
408
409impl StateRootContext {
410    /// Creates a new state root context.
411    fn new() -> Self {
412        Self {
413            account_rlp: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
414            trie_updates: TrieUpdates::default(),
415            hashed_entries_walked: 0,
416            updated_storage_nodes: 0,
417        }
418    }
419
420    /// Creates a [`StateRootProgress`] when the threshold is hit, from the state of the current
421    /// [`TrieNodeIter`], [`HashBuilder`], last hashed key and any storage root intermediate state.
422    fn create_progress_state<C, H, K>(
423        mut self,
424        account_node_iter: TrieNodeIter<C, H, K>,
425        hash_builder: HashBuilder,
426        last_hashed_key: B256,
427        storage_state: Option<IntermediateStorageRootState>,
428    ) -> StateRootProgress
429    where
430        C: TrieCursor,
431        H: HashedCursor,
432        K: AsRef<AddedRemovedKeys>,
433    {
434        let (walker_stack, walker_deleted_keys) = account_node_iter.walker.split();
435        self.trie_updates.removed_nodes.extend(walker_deleted_keys);
436        let (hash_builder, hash_builder_updates) = hash_builder.split();
437        self.trie_updates.account_nodes.extend(hash_builder_updates);
438
439        let account_state = IntermediateRootState { hash_builder, walker_stack, last_hashed_key };
440
441        let state = IntermediateStateRootState {
442            account_root_state: account_state,
443            storage_root_state: storage_state,
444        };
445
446        StateRootProgress::Progress(Box::new(state), self.hashed_entries_walked, self.trie_updates)
447    }
448
449    /// Calculates the total number of updated nodes.
450    fn total_updates_len<C, H, K>(
451        &self,
452        account_node_iter: &TrieNodeIter<C, H, K>,
453        hash_builder: &HashBuilder,
454    ) -> u64
455    where
456        C: TrieCursor,
457        H: HashedCursor,
458        K: AsRef<AddedRemovedKeys>,
459    {
460        (self.updated_storage_nodes +
461            account_node_iter.walker.removed_keys_len() +
462            hash_builder.updates_len()) as u64
463    }
464
465    /// Processes the result of a storage root calculation.
466    ///
467    /// Handles both completed and in-progress storage root calculations:
468    /// - For completed roots: encodes the account with the storage root, updates the hash builder
469    ///   with the new account, and updates metrics.
470    /// - For in-progress roots: returns the intermediate state for later resumption
471    ///
472    /// Returns an [`IntermediateStorageRootState`] if the calculation needs to be resumed later, or
473    /// `None` if the storage root was successfully computed and added to the trie.
474    fn process_storage_root_result(
475        &mut self,
476        storage_result: StorageRootProgress,
477        hashed_address: B256,
478        account: Account,
479        hash_builder: &mut HashBuilder,
480        retain_updates: bool,
481    ) -> Result<Option<IntermediateStorageRootState>, StateRootError> {
482        match storage_result {
483            StorageRootProgress::Complete(storage_root, storage_slots_walked, updates) => {
484                // Storage root completed
485                self.hashed_entries_walked += storage_slots_walked;
486                if retain_updates {
487                    self.updated_storage_nodes += updates.len();
488                    self.trie_updates.insert_storage_updates(hashed_address, updates);
489                }
490
491                // Encode the account with the computed storage root
492                self.account_rlp.clear();
493                let trie_account = account.into_trie_account(storage_root);
494                trie_account.encode(&mut self.account_rlp as &mut dyn BufMut);
495                hash_builder.add_leaf(Nibbles::unpack(hashed_address), &self.account_rlp);
496                Ok(None)
497            }
498            StorageRootProgress::Progress(state, storage_slots_walked, updates) => {
499                // Storage root hit threshold or resumed calculation hit threshold
500                debug!(
501                    target: "trie::state_root",
502                    ?hashed_address,
503                    storage_slots_walked,
504                    last_storage_key = ?state.last_hashed_key,
505                    ?account,
506                    "Pausing storage root calculation"
507                );
508
509                self.hashed_entries_walked += storage_slots_walked;
510                if retain_updates {
511                    self.trie_updates.insert_storage_updates(hashed_address, updates);
512                }
513
514                Ok(Some(IntermediateStorageRootState { state: *state, account }))
515            }
516        }
517    }
518}
519
520/// `StorageRoot` is used to compute the root node of an account storage trie.
521#[derive(Debug)]
522pub struct StorageRoot<T, H> {
523    /// A reference to the database transaction.
524    pub trie_cursor_factory: T,
525    /// The factory for hashed cursors.
526    pub hashed_cursor_factory: H,
527    /// The hashed address of an account.
528    pub hashed_address: B256,
529    /// The set of storage slot prefixes that have changed.
530    pub prefix_set: PrefixSet,
531    /// Whether every child under a branch whose path matches the prefix set should be walked.
532    walk_all_changed_branch_children: bool,
533    /// Previous intermediate state.
534    previous_state: Option<IntermediateRootState>,
535    /// The number of updates after which the intermediate progress should be returned.
536    threshold: u64,
537    /// Storage root metrics.
538    #[cfg(feature = "metrics")]
539    metrics: TrieRootMetrics,
540}
541
542impl<T, H> StorageRoot<T, H> {
543    /// Creates a new storage root calculator given a raw address.
544    pub fn new(
545        trie_cursor_factory: T,
546        hashed_cursor_factory: H,
547        address: Address,
548        prefix_set: PrefixSet,
549        #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
550    ) -> Self {
551        Self::new_hashed(
552            trie_cursor_factory,
553            hashed_cursor_factory,
554            keccak256(address),
555            prefix_set,
556            #[cfg(feature = "metrics")]
557            metrics,
558        )
559    }
560
561    /// Creates a new storage root calculator given a hashed address.
562    pub const fn new_hashed(
563        trie_cursor_factory: T,
564        hashed_cursor_factory: H,
565        hashed_address: B256,
566        prefix_set: PrefixSet,
567        #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
568    ) -> Self {
569        Self {
570            trie_cursor_factory,
571            hashed_cursor_factory,
572            hashed_address,
573            prefix_set,
574            walk_all_changed_branch_children: false,
575            previous_state: None,
576            threshold: DEFAULT_INTERMEDIATE_THRESHOLD,
577            #[cfg(feature = "metrics")]
578            metrics,
579        }
580    }
581
582    /// Set the changed prefixes.
583    pub fn with_prefix_set(mut self, prefix_set: PrefixSet) -> Self {
584        self.prefix_set = prefix_set;
585        self
586    }
587
588    /// Configures the storage root walker to visit all children of changed branch paths.
589    pub const fn with_walk_all_changed_branch_children(mut self, enabled: bool) -> Self {
590        self.walk_all_changed_branch_children = enabled;
591        self
592    }
593
594    /// Set the threshold.
595    pub const fn with_threshold(mut self, threshold: u64) -> Self {
596        self.threshold = threshold;
597        self
598    }
599
600    /// Set the threshold to maximum value so that intermediate progress is not returned.
601    pub const fn with_no_threshold(mut self) -> Self {
602        self.threshold = u64::MAX;
603        self
604    }
605
606    /// Set the previously recorded intermediate state.
607    pub fn with_intermediate_state(mut self, state: Option<IntermediateRootState>) -> Self {
608        self.previous_state = state;
609        self
610    }
611
612    /// Set the hashed cursor factory.
613    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageRoot<T, HF> {
614        StorageRoot {
615            trie_cursor_factory: self.trie_cursor_factory,
616            hashed_cursor_factory,
617            hashed_address: self.hashed_address,
618            prefix_set: self.prefix_set,
619            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
620            previous_state: self.previous_state,
621            threshold: self.threshold,
622            #[cfg(feature = "metrics")]
623            metrics: self.metrics,
624        }
625    }
626
627    /// Set the trie cursor factory.
628    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageRoot<TF, H> {
629        StorageRoot {
630            trie_cursor_factory,
631            hashed_cursor_factory: self.hashed_cursor_factory,
632            hashed_address: self.hashed_address,
633            prefix_set: self.prefix_set,
634            walk_all_changed_branch_children: self.walk_all_changed_branch_children,
635            previous_state: self.previous_state,
636            threshold: self.threshold,
637            #[cfg(feature = "metrics")]
638            metrics: self.metrics,
639        }
640    }
641}
642
643impl<T, H> StorageRoot<T, H>
644where
645    T: TrieCursorFactory,
646    H: HashedCursorFactory,
647{
648    /// Walks the intermediate nodes of existing storage trie (if any) and hashed entries. Feeds the
649    /// nodes into the hash builder. Collects the updates in the process.
650    ///
651    /// # Returns
652    ///
653    /// The intermediate progress of state root computation.
654    pub fn root_with_progress(self) -> Result<StorageRootProgress, StorageRootError> {
655        self.calculate(true)
656    }
657
658    /// Walks the hashed storage table entries for a given address and calculates the storage root.
659    ///
660    /// # Returns
661    ///
662    /// The storage root and storage trie updates for a given address.
663    pub fn root_with_updates(self) -> Result<(B256, usize, StorageTrieUpdates), StorageRootError> {
664        match self.with_no_threshold().calculate(true)? {
665            StorageRootProgress::Complete(root, walked, updates) => Ok((root, walked, updates)),
666            StorageRootProgress::Progress(..) => unreachable!(), // unreachable threshold
667        }
668    }
669
670    /// Walks the hashed storage table entries for a given address and calculates the storage root.
671    ///
672    /// # Returns
673    ///
674    /// The storage root.
675    pub fn root(self) -> Result<B256, StorageRootError> {
676        match self.calculate(false)? {
677            StorageRootProgress::Complete(root, _, _) => Ok(root),
678            StorageRootProgress::Progress(..) => unreachable!(), // update retention is disabled
679        }
680    }
681
682    /// Walks the hashed storage table entries for a given address and calculates the storage root.
683    ///
684    /// # Returns
685    ///
686    /// The storage root, number of walked entries and trie updates
687    /// for a given address if requested.
688    #[instrument(skip_all, level = "trace", target = "trie::storage_root", name = "storage_trie", fields(addr = %self.hashed_address, storage_root))]
689    pub fn calculate(self, retain_updates: bool) -> Result<StorageRootProgress, StorageRootError> {
690        trace!(target: "trie::storage_root", "calculating storage root");
691
692        let Self {
693            trie_cursor_factory,
694            hashed_cursor_factory,
695            hashed_address,
696            prefix_set,
697            walk_all_changed_branch_children,
698            previous_state,
699            threshold,
700            #[cfg(feature = "metrics")]
701            metrics,
702        } = self;
703
704        let mut hashed_storage_cursor =
705            hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
706        let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(hashed_address)?;
707
708        Self::calculate_with_cursors(
709            StorageRootCalculation {
710                hashed_address,
711                prefix_set,
712                previous_state,
713                walk_all_changed_branch_children,
714                threshold,
715                retain_updates,
716            },
717            &mut storage_trie_cursor,
718            &mut hashed_storage_cursor,
719            #[cfg(feature = "metrics")]
720            &metrics,
721        )
722    }
723
724    /// Walks the hashed storage table entries for a given address and calculates the storage root
725    /// using a pre-created cursor. The cursor will be repositioned to the given hashed address.
726    ///
727    /// This method allows reusing a single cursor across multiple storage root calculations,
728    /// reducing overhead when computing storage roots for many accounts.
729    #[instrument(skip_all, level = "trace", target = "trie::storage_root", name = "storage_trie_with_cursor", fields(addr = %self.hashed_address, storage_root))]
730    pub fn calculate_with_cursor(
731        self,
732        hashed_storage_cursor: &mut H::StorageCursor<'_>,
733        retain_updates: bool,
734    ) -> Result<StorageRootProgress, StorageRootError> {
735        trace!(target: "trie::storage_root", "calculating storage root with cursor");
736
737        let Self {
738            trie_cursor_factory,
739            hashed_address,
740            prefix_set,
741            walk_all_changed_branch_children,
742            previous_state,
743            threshold,
744            #[cfg(feature = "metrics")]
745            metrics,
746            ..
747        } = self;
748
749        let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(hashed_address)?;
750
751        Self::calculate_with_cursors(
752            StorageRootCalculation {
753                hashed_address,
754                prefix_set,
755                previous_state,
756                walk_all_changed_branch_children,
757                threshold,
758                retain_updates,
759            },
760            &mut storage_trie_cursor,
761            hashed_storage_cursor,
762            #[cfg(feature = "metrics")]
763            &metrics,
764        )
765    }
766
767    /// Walks the hashed storage table entries for a given address and calculates the storage root
768    /// using pre-created cursors. The cursors will be repositioned to the given hashed address.
769    fn calculate_with_cursors<TC, HC>(
770        calculation: StorageRootCalculation,
771        trie_cursor: &mut TC,
772        hashed_storage_cursor: &mut HC,
773        #[cfg(feature = "metrics")] metrics: &TrieRootMetrics,
774    ) -> Result<StorageRootProgress, StorageRootError>
775    where
776        TC: TrieStorageCursor,
777        HC: HashedStorageCursor<Value = U256>,
778    {
779        let StorageRootCalculation {
780            hashed_address,
781            prefix_set,
782            previous_state,
783            walk_all_changed_branch_children,
784            threshold,
785            retain_updates,
786        } = calculation;
787        hashed_storage_cursor.set_hashed_address(hashed_address);
788
789        // Empty storage only needs to be walked when changed prefixes must produce trie updates.
790        if previous_state.is_none() &&
791            (!retain_updates || prefix_set.is_empty()) &&
792            hashed_storage_cursor.is_storage_empty()?
793        {
794            Span::current().record("storage_root", tracing::field::debug(EMPTY_ROOT_HASH));
795            return Ok(StorageRootProgress::Complete(
796                EMPTY_ROOT_HASH,
797                0,
798                StorageTrieUpdates::default(),
799            ))
800        }
801
802        trie_cursor.set_hashed_address(hashed_address);
803
804        let mut tracker = TrieTracker::default();
805        let mut trie_updates = StorageTrieUpdates::default();
806
807        let (mut hash_builder, mut storage_node_iter) = match previous_state {
808            Some(state) => {
809                let hash_builder = state.hash_builder.with_updates(retain_updates);
810                let walker = TrieWalker::<_>::storage_trie_from_stack(
811                    trie_cursor,
812                    state.walker_stack,
813                    prefix_set,
814                )
815                .with_walk_all_changed_branch_children(walk_all_changed_branch_children)
816                .with_deletions_retained(retain_updates);
817                let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor)
818                    .with_last_hashed_key(state.last_hashed_key);
819                (hash_builder, node_iter)
820            }
821            None => {
822                let hash_builder = HashBuilder::default().with_updates(retain_updates);
823                let walker = TrieWalker::storage_trie(trie_cursor, prefix_set)
824                    .with_walk_all_changed_branch_children(walk_all_changed_branch_children)
825                    .with_deletions_retained(retain_updates);
826                let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor);
827                (hash_builder, node_iter)
828            }
829        };
830
831        let mut hashed_entries_walked = 0;
832        while let Some(node) = storage_node_iter.try_next()? {
833            match node {
834                TrieElement::Branch(node) => {
835                    tracker.inc_branch();
836                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
837                }
838                TrieElement::Leaf(hashed_slot, value) => {
839                    tracker.inc_leaf();
840                    hashed_entries_walked += 1;
841                    hash_builder.add_leaf(
842                        Nibbles::unpack(hashed_slot),
843                        alloy_rlp::encode_fixed_size(&value).as_ref(),
844                    );
845
846                    // Check if we need to return intermediate progress
847                    let total_updates_len =
848                        storage_node_iter.walker.removed_keys_len() + hash_builder.updates_len();
849                    if retain_updates && total_updates_len as u64 >= threshold {
850                        let (walker_stack, walker_deleted_keys) = storage_node_iter.walker.split();
851                        trie_updates.removed_nodes.extend(walker_deleted_keys);
852                        let (hash_builder, hash_builder_updates) = hash_builder.split();
853                        trie_updates.storage_nodes.extend(hash_builder_updates);
854
855                        let state = IntermediateRootState {
856                            hash_builder,
857                            walker_stack,
858                            last_hashed_key: hashed_slot,
859                        };
860
861                        return Ok(StorageRootProgress::Progress(
862                            Box::new(state),
863                            hashed_entries_walked,
864                            trie_updates,
865                        ))
866                    }
867                }
868            }
869        }
870
871        let root = hash_builder.root();
872        Span::current().record("storage_root", tracing::field::debug(root));
873
874        let removed_keys = storage_node_iter.walker.take_removed_keys();
875        trie_updates.finalize(hash_builder, removed_keys);
876
877        let stats = tracker.finish();
878
879        #[cfg(feature = "metrics")]
880        metrics.record(stats);
881
882        trace!(
883            target: "trie::storage_root",
884            %root,
885            %hashed_address,
886            duration = ?stats.duration(),
887            branches_added = stats.branches_added(),
888            leaves_added = stats.leaves_added(),
889            "calculated storage root"
890        );
891
892        let storage_slots_walked = stats.leaves_added() as usize;
893        Ok(StorageRootProgress::Complete(root, storage_slots_walked, trie_updates))
894    }
895}
896
897/// Parameters for a storage root calculation using pre-created cursors.
898struct StorageRootCalculation {
899    hashed_address: B256,
900    prefix_set: PrefixSet,
901    previous_state: Option<IntermediateRootState>,
902    walk_all_changed_branch_children: bool,
903    threshold: u64,
904    retain_updates: bool,
905}
906
907/// Trie type for differentiating between various trie calculations.
908#[derive(Clone, Copy, Debug)]
909pub enum TrieType {
910    /// State trie type.
911    State,
912    /// Storage trie type.
913    Storage,
914    /// Custom trie type. Can be used in ExEx.
915    Custom(&'static str),
916}
917
918impl TrieType {
919    #[cfg(feature = "metrics")]
920    pub(crate) const fn as_str(&self) -> &'static str {
921        match self {
922            Self::State => "state",
923            Self::Storage => "storage",
924            Self::Custom(s) => s,
925        }
926    }
927}