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},
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};
17use alloy_rlp::{BufMut, Encodable};
18use alloy_trie::proof::AddedRemovedKeys;
19use reth_execution_errors::{StateRootError, StorageRootError};
20use reth_primitives_traits::Account;
21use tracing::{debug, trace, trace_span};
22
23const DEFAULT_INTERMEDIATE_THRESHOLD: u64 = 100_000;
26
27#[cfg(feature = "metrics")]
28use crate::metrics::{StateRootMetrics, TrieRootMetrics};
29
30#[derive(Debug)]
32pub struct StateRoot<T, H> {
33 pub trie_cursor_factory: T,
35 pub hashed_cursor_factory: H,
37 pub prefix_sets: TriePrefixSets,
39 previous_state: Option<IntermediateStateRootState>,
41 threshold: u64,
43 #[cfg(feature = "metrics")]
44 metrics: StateRootMetrics,
46}
47
48impl<T, H> StateRoot<T, H> {
49 pub fn new(trie_cursor_factory: T, hashed_cursor_factory: H) -> Self {
55 Self {
56 trie_cursor_factory,
57 hashed_cursor_factory,
58 prefix_sets: TriePrefixSets::default(),
59 previous_state: None,
60 threshold: DEFAULT_INTERMEDIATE_THRESHOLD,
61 #[cfg(feature = "metrics")]
62 metrics: StateRootMetrics::default(),
63 }
64 }
65
66 pub fn with_prefix_sets(mut self, prefix_sets: TriePrefixSets) -> Self {
68 self.prefix_sets = prefix_sets;
69 self
70 }
71
72 pub const fn with_threshold(mut self, threshold: u64) -> Self {
74 self.threshold = threshold;
75 self
76 }
77
78 pub const fn with_no_threshold(mut self) -> Self {
80 self.threshold = u64::MAX;
81 self
82 }
83
84 pub fn with_intermediate_state(mut self, state: Option<IntermediateStateRootState>) -> Self {
86 self.previous_state = state;
87 self
88 }
89
90 pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StateRoot<T, HF> {
92 StateRoot {
93 trie_cursor_factory: self.trie_cursor_factory,
94 hashed_cursor_factory,
95 prefix_sets: self.prefix_sets,
96 threshold: self.threshold,
97 previous_state: self.previous_state,
98 #[cfg(feature = "metrics")]
99 metrics: self.metrics,
100 }
101 }
102
103 pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StateRoot<TF, H> {
105 StateRoot {
106 trie_cursor_factory,
107 hashed_cursor_factory: self.hashed_cursor_factory,
108 prefix_sets: self.prefix_sets,
109 threshold: self.threshold,
110 previous_state: self.previous_state,
111 #[cfg(feature = "metrics")]
112 metrics: self.metrics,
113 }
114 }
115}
116
117impl<T, H> StateRoot<T, H>
118where
119 T: TrieCursorFactory + Clone,
120 H: HashedCursorFactory + Clone,
121{
122 pub fn root_with_updates(self) -> Result<(B256, TrieUpdates), StateRootError> {
131 match self.with_no_threshold().calculate(true)? {
132 StateRootProgress::Complete(root, _, updates) => Ok((root, updates)),
133 StateRootProgress::Progress(..) => unreachable!(), }
135 }
136
137 pub fn root(self) -> Result<B256, StateRootError> {
144 match self.calculate(false)? {
145 StateRootProgress::Complete(root, _, _) => Ok(root),
146 StateRootProgress::Progress(..) => unreachable!(), }
148 }
149
150 pub fn root_with_progress(self) -> Result<StateRootProgress, StateRootError> {
157 self.calculate(true)
158 }
159
160 fn calculate(self, retain_updates: bool) -> Result<StateRootProgress, StateRootError> {
161 trace!(target: "trie::state_root", "calculating state root");
162 let mut tracker = TrieTracker::default();
163
164 let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
165 let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
166
167 let mut storage_ctx = StateRootContext::new();
169
170 let (mut hash_builder, mut account_node_iter) = if let Some(state) = self.previous_state {
172 let IntermediateStateRootState { account_root_state, storage_root_state } = state;
173
174 let mut hash_builder = account_root_state.hash_builder.with_updates(retain_updates);
176 let walker = TrieWalker::<_>::state_trie_from_stack(
177 trie_cursor,
178 account_root_state.walker_stack,
179 self.prefix_sets.account_prefix_set,
180 )
181 .with_deletions_retained(retain_updates);
182 let account_node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor)
183 .with_last_hashed_key(account_root_state.last_hashed_key);
184
185 if let Some(storage_state) = storage_root_state {
187 let hashed_address = account_root_state.last_hashed_key;
188 let account = storage_state.account;
189
190 debug!(
191 target: "trie::state_root",
192 account_nonce = account.nonce,
193 account_balance = ?account.balance,
194 last_hashed_key = ?account_root_state.last_hashed_key,
195 "Resuming storage root calculation"
196 );
197
198 let remaining_threshold = self.threshold.saturating_sub(
200 storage_ctx.total_updates_len(&account_node_iter, &hash_builder),
201 );
202
203 let storage_root_calculator = StorageRoot::new_hashed(
204 self.trie_cursor_factory.clone(),
205 self.hashed_cursor_factory.clone(),
206 hashed_address,
207 self.prefix_sets
208 .storage_prefix_sets
209 .get(&hashed_address)
210 .cloned()
211 .unwrap_or_default(),
212 #[cfg(feature = "metrics")]
213 self.metrics.storage_trie.clone(),
214 )
215 .with_intermediate_state(Some(storage_state.state))
216 .with_threshold(remaining_threshold);
217
218 let storage_result = storage_root_calculator.calculate(retain_updates)?;
219 if let Some(storage_state) = storage_ctx.process_storage_root_result(
220 storage_result,
221 hashed_address,
222 account,
223 &mut hash_builder,
224 retain_updates,
225 )? {
226 return Ok(storage_ctx.create_progress_state(
228 account_node_iter,
229 hash_builder,
230 account_root_state.last_hashed_key,
231 Some(storage_state),
232 ))
233 }
234 }
235
236 (hash_builder, account_node_iter)
237 } else {
238 let hash_builder = HashBuilder::default().with_updates(retain_updates);
241 let walker = TrieWalker::state_trie(trie_cursor, self.prefix_sets.account_prefix_set)
242 .with_deletions_retained(retain_updates);
243 let node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor);
244 (hash_builder, node_iter)
245 };
246
247 while let Some(node) = account_node_iter.try_next()? {
248 match node {
249 TrieElement::Branch(node) => {
250 tracker.inc_branch();
251 hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
252 }
253 TrieElement::Leaf(hashed_address, account) => {
254 tracker.inc_leaf();
255 storage_ctx.hashed_entries_walked += 1;
256
257 let remaining_threshold = self.threshold.saturating_sub(
260 storage_ctx.total_updates_len(&account_node_iter, &hash_builder),
261 );
262
263 let storage_root_calculator = StorageRoot::new_hashed(
264 self.trie_cursor_factory.clone(),
265 self.hashed_cursor_factory.clone(),
266 hashed_address,
267 self.prefix_sets
268 .storage_prefix_sets
269 .get(&hashed_address)
270 .cloned()
271 .unwrap_or_default(),
272 #[cfg(feature = "metrics")]
273 self.metrics.storage_trie.clone(),
274 )
275 .with_threshold(remaining_threshold);
276
277 let storage_result = storage_root_calculator.calculate(retain_updates)?;
278 if let Some(storage_state) = storage_ctx.process_storage_root_result(
279 storage_result,
280 hashed_address,
281 account,
282 &mut hash_builder,
283 retain_updates,
284 )? {
285 return Ok(storage_ctx.create_progress_state(
287 account_node_iter,
288 hash_builder,
289 hashed_address,
290 Some(storage_state),
291 ))
292 }
293
294 let total_updates_len =
296 storage_ctx.total_updates_len(&account_node_iter, &hash_builder);
297 if retain_updates && total_updates_len >= self.threshold {
298 return Ok(storage_ctx.create_progress_state(
299 account_node_iter,
300 hash_builder,
301 hashed_address,
302 None,
303 ))
304 }
305 }
306 }
307 }
308
309 let root = hash_builder.root();
310
311 let removed_keys = account_node_iter.walker.take_removed_keys();
312 let StateRootContext { mut trie_updates, hashed_entries_walked, .. } = storage_ctx;
313 trie_updates.finalize(hash_builder, removed_keys, self.prefix_sets.destroyed_accounts);
314
315 let stats = tracker.finish();
316
317 #[cfg(feature = "metrics")]
318 self.metrics.state_trie.record(stats);
319
320 trace!(
321 target: "trie::state_root",
322 %root,
323 duration = ?stats.duration(),
324 branches_added = stats.branches_added(),
325 leaves_added = stats.leaves_added(),
326 "calculated state root"
327 );
328
329 Ok(StateRootProgress::Complete(root, hashed_entries_walked, trie_updates))
330 }
331}
332
333#[derive(Debug)]
335pub(crate) struct StateRootContext {
336 account_rlp: Vec<u8>,
338 trie_updates: TrieUpdates,
340 hashed_entries_walked: usize,
342 updated_storage_nodes: usize,
344}
345
346impl StateRootContext {
347 fn new() -> Self {
349 Self {
350 account_rlp: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
351 trie_updates: TrieUpdates::default(),
352 hashed_entries_walked: 0,
353 updated_storage_nodes: 0,
354 }
355 }
356
357 fn create_progress_state<C, H, K>(
360 mut self,
361 account_node_iter: TrieNodeIter<C, H, K>,
362 hash_builder: HashBuilder,
363 last_hashed_key: B256,
364 storage_state: Option<IntermediateStorageRootState>,
365 ) -> StateRootProgress
366 where
367 C: TrieCursor,
368 H: HashedCursor,
369 K: AsRef<AddedRemovedKeys>,
370 {
371 let (walker_stack, walker_deleted_keys) = account_node_iter.walker.split();
372 self.trie_updates.removed_nodes.extend(walker_deleted_keys);
373 let (hash_builder, hash_builder_updates) = hash_builder.split();
374 self.trie_updates.account_nodes.extend(hash_builder_updates);
375
376 let account_state = IntermediateRootState { hash_builder, walker_stack, last_hashed_key };
377
378 let state = IntermediateStateRootState {
379 account_root_state: account_state,
380 storage_root_state: storage_state,
381 };
382
383 StateRootProgress::Progress(Box::new(state), self.hashed_entries_walked, self.trie_updates)
384 }
385
386 fn total_updates_len<C, H, K>(
388 &self,
389 account_node_iter: &TrieNodeIter<C, H, K>,
390 hash_builder: &HashBuilder,
391 ) -> u64
392 where
393 C: TrieCursor,
394 H: HashedCursor,
395 K: AsRef<AddedRemovedKeys>,
396 {
397 (self.updated_storage_nodes +
398 account_node_iter.walker.removed_keys_len() +
399 hash_builder.updates_len()) as u64
400 }
401
402 fn process_storage_root_result(
412 &mut self,
413 storage_result: StorageRootProgress,
414 hashed_address: B256,
415 account: Account,
416 hash_builder: &mut HashBuilder,
417 retain_updates: bool,
418 ) -> Result<Option<IntermediateStorageRootState>, StateRootError> {
419 match storage_result {
420 StorageRootProgress::Complete(storage_root, storage_slots_walked, updates) => {
421 self.hashed_entries_walked += storage_slots_walked;
423 if retain_updates {
424 self.updated_storage_nodes += updates.len();
425 self.trie_updates.insert_storage_updates(hashed_address, updates);
426 }
427
428 self.account_rlp.clear();
430 let trie_account = account.into_trie_account(storage_root);
431 trie_account.encode(&mut self.account_rlp as &mut dyn BufMut);
432 hash_builder.add_leaf(Nibbles::unpack(hashed_address), &self.account_rlp);
433 Ok(None)
434 }
435 StorageRootProgress::Progress(state, storage_slots_walked, updates) => {
436 debug!(
438 target: "trie::state_root",
439 ?hashed_address,
440 storage_slots_walked,
441 last_storage_key = ?state.last_hashed_key,
442 ?account,
443 "Pausing storage root calculation"
444 );
445
446 self.hashed_entries_walked += storage_slots_walked;
447 if retain_updates {
448 self.trie_updates.insert_storage_updates(hashed_address, updates);
449 }
450
451 Ok(Some(IntermediateStorageRootState { state: *state, account }))
452 }
453 }
454 }
455}
456
457#[derive(Debug)]
459pub struct StorageRoot<T, H> {
460 pub trie_cursor_factory: T,
462 pub hashed_cursor_factory: H,
464 pub hashed_address: B256,
466 pub prefix_set: PrefixSet,
468 previous_state: Option<IntermediateRootState>,
470 threshold: u64,
472 #[cfg(feature = "metrics")]
474 metrics: TrieRootMetrics,
475}
476
477impl<T, H> StorageRoot<T, H> {
478 pub fn new(
480 trie_cursor_factory: T,
481 hashed_cursor_factory: H,
482 address: Address,
483 prefix_set: PrefixSet,
484 #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
485 ) -> Self {
486 Self::new_hashed(
487 trie_cursor_factory,
488 hashed_cursor_factory,
489 keccak256(address),
490 prefix_set,
491 #[cfg(feature = "metrics")]
492 metrics,
493 )
494 }
495
496 pub const fn new_hashed(
498 trie_cursor_factory: T,
499 hashed_cursor_factory: H,
500 hashed_address: B256,
501 prefix_set: PrefixSet,
502 #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
503 ) -> Self {
504 Self {
505 trie_cursor_factory,
506 hashed_cursor_factory,
507 hashed_address,
508 prefix_set,
509 previous_state: None,
510 threshold: DEFAULT_INTERMEDIATE_THRESHOLD,
511 #[cfg(feature = "metrics")]
512 metrics,
513 }
514 }
515
516 pub fn with_prefix_set(mut self, prefix_set: PrefixSet) -> Self {
518 self.prefix_set = prefix_set;
519 self
520 }
521
522 pub const fn with_threshold(mut self, threshold: u64) -> Self {
524 self.threshold = threshold;
525 self
526 }
527
528 pub const fn with_no_threshold(mut self) -> Self {
530 self.threshold = u64::MAX;
531 self
532 }
533
534 pub fn with_intermediate_state(mut self, state: Option<IntermediateRootState>) -> Self {
536 self.previous_state = state;
537 self
538 }
539
540 pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageRoot<T, HF> {
542 StorageRoot {
543 trie_cursor_factory: self.trie_cursor_factory,
544 hashed_cursor_factory,
545 hashed_address: self.hashed_address,
546 prefix_set: self.prefix_set,
547 previous_state: self.previous_state,
548 threshold: self.threshold,
549 #[cfg(feature = "metrics")]
550 metrics: self.metrics,
551 }
552 }
553
554 pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageRoot<TF, H> {
556 StorageRoot {
557 trie_cursor_factory,
558 hashed_cursor_factory: self.hashed_cursor_factory,
559 hashed_address: self.hashed_address,
560 prefix_set: self.prefix_set,
561 previous_state: self.previous_state,
562 threshold: self.threshold,
563 #[cfg(feature = "metrics")]
564 metrics: self.metrics,
565 }
566 }
567}
568
569impl<T, H> StorageRoot<T, H>
570where
571 T: TrieCursorFactory,
572 H: HashedCursorFactory,
573{
574 pub fn root_with_progress(self) -> Result<StorageRootProgress, StorageRootError> {
581 self.calculate(true)
582 }
583
584 pub fn root_with_updates(self) -> Result<(B256, usize, StorageTrieUpdates), StorageRootError> {
590 match self.with_no_threshold().calculate(true)? {
591 StorageRootProgress::Complete(root, walked, updates) => Ok((root, walked, updates)),
592 StorageRootProgress::Progress(..) => unreachable!(), }
594 }
595
596 pub fn root(self) -> Result<B256, StorageRootError> {
602 match self.calculate(false)? {
603 StorageRootProgress::Complete(root, _, _) => Ok(root),
604 StorageRootProgress::Progress(..) => unreachable!(), }
606 }
607
608 pub fn calculate(self, retain_updates: bool) -> Result<StorageRootProgress, StorageRootError> {
615 let span = trace_span!(target: "trie::storage_root", "Storage trie", hashed_address = ?self.hashed_address);
616 let _enter = span.enter();
617
618 trace!(target: "trie::storage_root", "calculating storage root");
619
620 let mut hashed_storage_cursor =
621 self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
622
623 if hashed_storage_cursor.is_storage_empty()? {
625 return Ok(StorageRootProgress::Complete(
626 EMPTY_ROOT_HASH,
627 0,
628 StorageTrieUpdates::deleted(),
629 ))
630 }
631
632 let mut tracker = TrieTracker::default();
633 let mut trie_updates = StorageTrieUpdates::default();
634
635 let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
636
637 let (mut hash_builder, mut storage_node_iter) = match self.previous_state {
638 Some(state) => {
639 let hash_builder = state.hash_builder.with_updates(retain_updates);
640 let walker = TrieWalker::<_>::storage_trie_from_stack(
641 trie_cursor,
642 state.walker_stack,
643 self.prefix_set,
644 )
645 .with_deletions_retained(retain_updates);
646 let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor)
647 .with_last_hashed_key(state.last_hashed_key);
648 (hash_builder, node_iter)
649 }
650 None => {
651 let hash_builder = HashBuilder::default().with_updates(retain_updates);
652 let walker = TrieWalker::storage_trie(trie_cursor, self.prefix_set)
653 .with_deletions_retained(retain_updates);
654 let node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor);
655 (hash_builder, node_iter)
656 }
657 };
658
659 let mut hashed_entries_walked = 0;
660 while let Some(node) = storage_node_iter.try_next()? {
661 match node {
662 TrieElement::Branch(node) => {
663 tracker.inc_branch();
664 hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
665 }
666 TrieElement::Leaf(hashed_slot, value) => {
667 tracker.inc_leaf();
668 hashed_entries_walked += 1;
669 hash_builder.add_leaf(
670 Nibbles::unpack(hashed_slot),
671 alloy_rlp::encode_fixed_size(&value).as_ref(),
672 );
673
674 let total_updates_len =
676 storage_node_iter.walker.removed_keys_len() + hash_builder.updates_len();
677 if retain_updates && total_updates_len as u64 >= self.threshold {
678 let (walker_stack, walker_deleted_keys) = storage_node_iter.walker.split();
679 trie_updates.removed_nodes.extend(walker_deleted_keys);
680 let (hash_builder, hash_builder_updates) = hash_builder.split();
681 trie_updates.storage_nodes.extend(hash_builder_updates);
682
683 let state = IntermediateRootState {
684 hash_builder,
685 walker_stack,
686 last_hashed_key: hashed_slot,
687 };
688
689 return Ok(StorageRootProgress::Progress(
690 Box::new(state),
691 hashed_entries_walked,
692 trie_updates,
693 ))
694 }
695 }
696 }
697 }
698
699 let root = hash_builder.root();
700
701 let removed_keys = storage_node_iter.walker.take_removed_keys();
702 trie_updates.finalize(hash_builder, removed_keys);
703
704 let stats = tracker.finish();
705
706 #[cfg(feature = "metrics")]
707 self.metrics.record(stats);
708
709 trace!(
710 target: "trie::storage_root",
711 %root,
712 hashed_address = %self.hashed_address,
713 duration = ?stats.duration(),
714 branches_added = stats.branches_added(),
715 leaves_added = stats.leaves_added(),
716 "calculated storage root"
717 );
718
719 let storage_slots_walked = stats.leaves_added() as usize;
720 Ok(StorageRootProgress::Complete(root, storage_slots_walked, trie_updates))
721 }
722}
723
724#[derive(Clone, Copy, Debug)]
726pub enum TrieType {
727 State,
729 Storage,
731}
732
733impl TrieType {
734 #[cfg(feature = "metrics")]
735 pub(crate) const fn as_str(&self) -> &'static str {
736 match self {
737 Self::State => "state",
738 Self::Storage => "storage",
739 }
740 }
741}