1use crate::{
2 blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
3 LeafLookup, RevealedSparseTrie, SparseTrie, TrieMasks,
4};
5use alloc::{collections::VecDeque, vec::Vec};
6use alloy_primitives::{
7 hex,
8 map::{B256Map, HashMap, HashSet},
9 Bytes, B256,
10};
11use alloy_rlp::{Decodable, Encodable};
12use alloy_trie::proof::DecodedProofNodes;
13use core::{fmt, iter::Peekable};
14use reth_execution_errors::{SparseStateTrieErrorKind, SparseStateTrieResult, SparseTrieErrorKind};
15use reth_primitives_traits::Account;
16use reth_trie_common::{
17 proof::ProofNodes,
18 updates::{StorageTrieUpdates, TrieUpdates},
19 DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, RlpNode, StorageMultiProof,
20 TrieAccount, TrieMask, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
21};
22use tracing::trace;
23
24pub struct SparseStateTrie<F: BlindedProviderFactory = DefaultBlindedProviderFactory> {
26 provider_factory: F,
28 state: SparseTrie<F::AccountNodeProvider>,
30 storages: B256Map<SparseTrie<F::StorageNodeProvider>>,
32 revealed_account_paths: HashSet<Nibbles>,
34 revealed_storage_paths: B256Map<HashSet<Nibbles>>,
36 retain_updates: bool,
38 account_rlp_buf: Vec<u8>,
40 #[cfg(feature = "metrics")]
42 metrics: crate::metrics::SparseStateTrieMetrics,
43}
44
45#[cfg(test)]
46impl Default for SparseStateTrie {
47 fn default() -> Self {
48 Self {
49 provider_factory: DefaultBlindedProviderFactory,
50 state: Default::default(),
51 storages: Default::default(),
52 revealed_account_paths: Default::default(),
53 revealed_storage_paths: Default::default(),
54 retain_updates: false,
55 account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
56 #[cfg(feature = "metrics")]
57 metrics: Default::default(),
58 }
59 }
60}
61
62impl<P: BlindedProviderFactory> fmt::Debug for SparseStateTrie<P> {
63 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64 f.debug_struct("SparseStateTrie")
65 .field("state", &self.state)
66 .field("storages", &self.storages)
67 .field("revealed_account_paths", &self.revealed_account_paths)
68 .field("revealed_storage_paths", &self.revealed_storage_paths)
69 .field("retain_updates", &self.retain_updates)
70 .field("account_rlp_buf", &hex::encode(&self.account_rlp_buf))
71 .finish_non_exhaustive()
72 }
73}
74
75#[cfg(test)]
76impl SparseStateTrie {
77 pub fn from_state(state: SparseTrie) -> Self {
79 Self { state, ..Default::default() }
80 }
81}
82
83impl<F: BlindedProviderFactory> SparseStateTrie<F> {
84 pub fn new(provider_factory: F) -> Self {
86 Self {
87 provider_factory,
88 state: Default::default(),
89 storages: Default::default(),
90 revealed_account_paths: Default::default(),
91 revealed_storage_paths: Default::default(),
92 retain_updates: false,
93 account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
94 #[cfg(feature = "metrics")]
95 metrics: Default::default(),
96 }
97 }
98
99 pub const fn with_updates(mut self, retain_updates: bool) -> Self {
101 self.retain_updates = retain_updates;
102 self
103 }
104
105 pub fn is_account_revealed(&self, account: B256) -> bool {
107 self.revealed_account_paths.contains(&Nibbles::unpack(account))
108 }
109
110 pub fn check_valid_account_witness(&self, address: B256) -> bool {
112 let path = Nibbles::unpack(address);
113 let trie = match self.state_trie_ref() {
114 Some(t) => t,
115 None => return false,
116 };
117
118 matches!(
119 trie.find_leaf(&path, None),
120 Ok(LeafLookup::Exists | LeafLookup::NonExistent { .. })
121 )
122 }
123
124 pub fn check_valid_storage_witness(&self, address: B256, slot: B256) -> bool {
126 let path = Nibbles::unpack(slot);
127 let trie = match self.storage_trie_ref(&address) {
128 Some(t) => t,
129 None => return false,
130 };
131
132 matches!(
133 trie.find_leaf(&path, None),
134 Ok(LeafLookup::Exists | LeafLookup::NonExistent { .. })
135 )
136 }
137
138 pub fn is_storage_slot_revealed(&self, account: B256, slot: B256) -> bool {
140 self.revealed_storage_paths
141 .get(&account)
142 .is_some_and(|slots| slots.contains(&Nibbles::unpack(slot)))
143 }
144
145 pub fn get_account_value(&self, account: &B256) -> Option<&Vec<u8>> {
147 self.state.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(account))
148 }
149
150 pub fn get_storage_slot_value(&self, account: &B256, slot: &B256) -> Option<&Vec<u8>> {
152 self.storages.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
153 }
154
155 pub const fn state_trie_ref(&self) -> Option<&RevealedSparseTrie<F::AccountNodeProvider>> {
157 self.state.as_revealed_ref()
158 }
159
160 pub fn storage_trie_ref(
162 &self,
163 address: &B256,
164 ) -> Option<&RevealedSparseTrie<F::StorageNodeProvider>> {
165 self.storages.get(address).and_then(|e| e.as_revealed_ref())
166 }
167
168 pub fn storage_trie_mut(
170 &mut self,
171 address: &B256,
172 ) -> Option<&mut RevealedSparseTrie<F::StorageNodeProvider>> {
173 self.storages.get_mut(address).and_then(|e| e.as_revealed_mut())
174 }
175
176 pub fn take_storage_trie(
178 &mut self,
179 address: &B256,
180 ) -> Option<SparseTrie<F::StorageNodeProvider>> {
181 self.storages.remove(address)
182 }
183
184 pub fn insert_storage_trie(
186 &mut self,
187 address: B256,
188 storage_trie: SparseTrie<F::StorageNodeProvider>,
189 ) {
190 self.storages.insert(address, storage_trie);
191 }
192
193 pub fn reveal_account(
199 &mut self,
200 account: B256,
201 proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
202 ) -> SparseStateTrieResult<()> {
203 assert!(!self.retain_updates);
204
205 if self.is_account_revealed(account) {
206 return Ok(());
207 }
208
209 let mut proof = proof.into_iter().peekable();
210
211 let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
212
213 let trie = self.state.reveal_root_with_provider(
215 self.provider_factory.account_node_provider(),
216 root_node,
217 TrieMasks::none(),
218 self.retain_updates,
219 )?;
220
221 for (path, bytes) in proof {
223 if self.revealed_account_paths.contains(&path) {
224 continue;
225 }
226 let node = TrieNode::decode(&mut &bytes[..])?;
227 trie.reveal_node(path.clone(), node, TrieMasks::none())?;
228
229 self.revealed_account_paths.insert(path);
231 }
232
233 Ok(())
234 }
235
236 pub fn reveal_storage_slot(
242 &mut self,
243 account: B256,
244 slot: B256,
245 proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
246 ) -> SparseStateTrieResult<()> {
247 assert!(!self.retain_updates);
248
249 if self.is_storage_slot_revealed(account, slot) {
250 return Ok(());
251 }
252
253 let mut proof = proof.into_iter().peekable();
254
255 let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
256
257 let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
259 self.provider_factory.storage_node_provider(account),
260 root_node,
261 TrieMasks::none(),
262 self.retain_updates,
263 )?;
264
265 let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
266
267 for (path, bytes) in proof {
269 if revealed_nodes.contains(&path) {
271 continue;
272 }
273 let node = TrieNode::decode(&mut &bytes[..])?;
274 trie.reveal_node(path.clone(), node, TrieMasks::none())?;
275
276 revealed_nodes.insert(path);
278 }
279
280 Ok(())
281 }
282
283 pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
286 let decoded_multiproof = multiproof.try_into()?;
288
289 self.reveal_decoded_multiproof(decoded_multiproof)
291 }
292
293 pub fn reveal_decoded_multiproof(
296 &mut self,
297 multiproof: DecodedMultiProof,
298 ) -> SparseStateTrieResult<()> {
299 let DecodedMultiProof {
300 account_subtree,
301 storages,
302 branch_node_hash_masks,
303 branch_node_tree_masks,
304 } = multiproof;
305
306 self.reveal_decoded_account_multiproof(
308 account_subtree,
309 branch_node_hash_masks,
310 branch_node_tree_masks,
311 )?;
312
313 for (account, storage_subtree) in storages {
315 self.reveal_decoded_storage_multiproof(account, storage_subtree)?;
316 }
317
318 Ok(())
319 }
320
321 pub fn reveal_account_multiproof(
323 &mut self,
324 account_subtree: ProofNodes,
325 branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
326 branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
327 ) -> SparseStateTrieResult<()> {
328 let decoded_multiproof = account_subtree.try_into()?;
330 self.reveal_decoded_account_multiproof(
331 decoded_multiproof,
332 branch_node_hash_masks,
333 branch_node_tree_masks,
334 )
335 }
336
337 pub fn reveal_decoded_account_multiproof(
339 &mut self,
340 account_subtree: DecodedProofNodes,
341 branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
342 branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
343 ) -> SparseStateTrieResult<()> {
344 let FilteredProofNodes {
345 nodes,
346 new_nodes,
347 total_nodes: _total_nodes,
348 skipped_nodes: _skipped_nodes,
349 } = filter_revealed_nodes(account_subtree, &self.revealed_account_paths)?;
350 #[cfg(feature = "metrics")]
351 {
352 self.metrics.increment_total_account_nodes(_total_nodes as u64);
353 self.metrics.increment_skipped_account_nodes(_skipped_nodes as u64);
354 }
355 let mut account_nodes = nodes.into_iter().peekable();
356
357 if let Some(root_node) = Self::validate_root_node_decoded(&mut account_nodes)? {
358 let trie = self.state.reveal_root_with_provider(
360 self.provider_factory.account_node_provider(),
361 root_node,
362 TrieMasks {
363 hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
364 tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
365 },
366 self.retain_updates,
367 )?;
368
369 trie.reserve_nodes(new_nodes);
371
372 for (path, node) in account_nodes {
374 let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
383 (
384 branch_node_hash_masks.get(&path).copied(),
385 branch_node_tree_masks.get(&path).copied(),
386 )
387 } else {
388 (None, None)
389 };
390
391 trace!(target: "trie::sparse", ?path, ?node, ?hash_mask, ?tree_mask, "Revealing account node");
392 trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
393
394 self.revealed_account_paths.insert(path);
396 }
397 }
398
399 Ok(())
400 }
401
402 pub fn reveal_storage_multiproof(
404 &mut self,
405 account: B256,
406 storage_subtree: StorageMultiProof,
407 ) -> SparseStateTrieResult<()> {
408 let decoded_multiproof = storage_subtree.try_into()?;
410 self.reveal_decoded_storage_multiproof(account, decoded_multiproof)
411 }
412
413 pub fn reveal_decoded_storage_multiproof(
415 &mut self,
416 account: B256,
417 storage_subtree: DecodedStorageMultiProof,
418 ) -> SparseStateTrieResult<()> {
419 let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
420
421 let FilteredProofNodes {
422 nodes,
423 new_nodes,
424 total_nodes: _total_nodes,
425 skipped_nodes: _skipped_nodes,
426 } = filter_revealed_nodes(storage_subtree.subtree, revealed_nodes)?;
427 #[cfg(feature = "metrics")]
428 {
429 self.metrics.increment_total_storage_nodes(_total_nodes as u64);
430 self.metrics.increment_skipped_storage_nodes(_skipped_nodes as u64);
431 }
432 let mut nodes = nodes.into_iter().peekable();
433
434 if let Some(root_node) = Self::validate_root_node_decoded(&mut nodes)? {
435 let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
437 self.provider_factory.storage_node_provider(account),
438 root_node,
439 TrieMasks {
440 hash_mask: storage_subtree
441 .branch_node_hash_masks
442 .get(&Nibbles::default())
443 .copied(),
444 tree_mask: storage_subtree
445 .branch_node_tree_masks
446 .get(&Nibbles::default())
447 .copied(),
448 },
449 self.retain_updates,
450 )?;
451
452 trie.reserve_nodes(new_nodes);
454
455 for (path, node) in nodes {
457 let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
466 (
467 storage_subtree.branch_node_hash_masks.get(&path).copied(),
468 storage_subtree.branch_node_tree_masks.get(&path).copied(),
469 )
470 } else {
471 (None, None)
472 };
473
474 trace!(target: "trie::sparse", ?account, ?path, ?node, ?hash_mask, ?tree_mask, "Revealing storage node");
475 trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
476
477 revealed_nodes.insert(path);
479 }
480 }
481
482 Ok(())
483 }
484
485 pub fn reveal_witness(
489 &mut self,
490 state_root: B256,
491 witness: &B256Map<Bytes>,
492 ) -> SparseStateTrieResult<()> {
493 let mut queue = VecDeque::from([(state_root, Nibbles::default(), None)]);
496
497 while let Some((hash, path, maybe_account)) = queue.pop_front() {
498 let Some(trie_node_bytes) = witness.get(&hash) else { continue };
500 let trie_node = TrieNode::decode(&mut &trie_node_bytes[..])?;
501
502 match &trie_node {
504 TrieNode::Branch(branch) => {
505 for (idx, maybe_child) in branch.as_ref().children() {
506 if let Some(child_hash) = maybe_child.and_then(RlpNode::as_hash) {
507 let mut child_path = path.clone();
508 child_path.push_unchecked(idx);
509 queue.push_back((child_hash, child_path, maybe_account));
510 }
511 }
512 }
513 TrieNode::Extension(ext) => {
514 if let Some(child_hash) = ext.child.as_hash() {
515 let mut child_path = path.clone();
516 child_path.extend_from_slice_unchecked(&ext.key);
517 queue.push_back((child_hash, child_path, maybe_account));
518 }
519 }
520 TrieNode::Leaf(leaf) => {
521 let mut full_path = path.clone();
522 full_path.extend_from_slice_unchecked(&leaf.key);
523 if maybe_account.is_none() {
524 let hashed_address = B256::from_slice(&full_path.pack());
525 let account = TrieAccount::decode(&mut &leaf.value[..])?;
526 if account.storage_root != EMPTY_ROOT_HASH {
527 queue.push_back((
528 account.storage_root,
529 Nibbles::default(),
530 Some(hashed_address),
531 ));
532 }
533 }
534 }
535 TrieNode::EmptyRoot => {} };
537
538 if let Some(account) = maybe_account {
540 if self
542 .revealed_storage_paths
543 .get(&account)
544 .is_none_or(|paths| !paths.contains(&path))
545 {
546 let storage_trie_entry = self.storages.entry(account).or_default();
547 if path.is_empty() {
548 storage_trie_entry.reveal_root_with_provider(
550 self.provider_factory.storage_node_provider(account),
551 trie_node,
552 TrieMasks::none(),
553 self.retain_updates,
554 )?;
555 } else {
556 storage_trie_entry
558 .as_revealed_mut()
559 .ok_or(SparseTrieErrorKind::Blind)?
560 .reveal_node(path.clone(), trie_node, TrieMasks::none())?;
561 }
562
563 self.revealed_storage_paths.entry(account).or_default().insert(path);
565 }
566 }
567 else if !self.revealed_account_paths.contains(&path) {
569 if path.is_empty() {
570 self.state.reveal_root_with_provider(
572 self.provider_factory.account_node_provider(),
573 trie_node,
574 TrieMasks::none(),
575 self.retain_updates,
576 )?;
577 } else {
578 self.state.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?.reveal_node(
580 path.clone(),
581 trie_node,
582 TrieMasks::none(),
583 )?;
584 }
585
586 self.revealed_account_paths.insert(path);
588 }
589 }
590
591 Ok(())
592 }
593
594 fn validate_root_node<I: Iterator<Item = (Nibbles, Bytes)>>(
596 &self,
597 proof: &mut Peekable<I>,
598 ) -> SparseStateTrieResult<Option<TrieNode>> {
599 let Some((path, node)) = proof.next() else { return Ok(None) };
601 if !path.is_empty() {
602 return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into());
603 }
604
605 let root_node = TrieNode::decode(&mut &node[..])?;
607 if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
608 return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into());
609 }
610
611 Ok(Some(root_node))
612 }
613
614 fn validate_root_node_decoded<I: Iterator<Item = (Nibbles, TrieNode)>>(
616 proof: &mut Peekable<I>,
617 ) -> SparseStateTrieResult<Option<TrieNode>> {
618 let Some((path, root_node)) = proof.next() else { return Ok(None) };
620 if !path.is_empty() {
621 return Err(SparseStateTrieErrorKind::InvalidRootNode {
622 path,
623 node: alloy_rlp::encode(&root_node).into(),
624 }
625 .into())
626 }
627
628 if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
630 return Err(SparseStateTrieErrorKind::InvalidRootNode {
631 path,
632 node: alloy_rlp::encode(&root_node).into(),
633 }
634 .into())
635 }
636
637 Ok(Some(root_node))
638 }
639
640 pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
642 if let Some(trie) = self.storages.get_mut(&address) {
643 trie.wipe()?;
644 }
645 Ok(())
646 }
647
648 pub fn calculate_below_level(&mut self, level: usize) {
652 if let SparseTrie::Revealed(trie) = &mut self.state {
653 trie.update_rlp_node_level(level);
654 }
655 }
656
657 pub fn storage_root(&mut self, account: B256) -> Option<B256> {
659 self.storages.get_mut(&account).and_then(|trie| trie.root())
660 }
661
662 fn revealed_trie_mut(
666 &mut self,
667 ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
668 match self.state {
669 SparseTrie::Blind => {
670 let (root_node, hash_mask, tree_mask) = self
671 .provider_factory
672 .account_node_provider()
673 .blinded_node(&Nibbles::default())?
674 .map(|node| {
675 TrieNode::decode(&mut &node.node[..])
676 .map(|decoded| (decoded, node.hash_mask, node.tree_mask))
677 })
678 .transpose()?
679 .unwrap_or((TrieNode::EmptyRoot, None, None));
680 self.state
681 .reveal_root_with_provider(
682 self.provider_factory.account_node_provider(),
683 root_node,
684 TrieMasks { hash_mask, tree_mask },
685 self.retain_updates,
686 )
687 .map_err(Into::into)
688 }
689 SparseTrie::Revealed(ref mut trie) => Ok(trie),
690 }
691 }
692
693 pub fn root(&mut self) -> SparseStateTrieResult<B256> {
697 #[cfg(feature = "metrics")]
699 self.metrics.record();
700
701 Ok(self.revealed_trie_mut()?.root())
702 }
703
704 pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
706 #[cfg(feature = "metrics")]
708 self.metrics.record();
709
710 let storage_tries = self.storage_trie_updates();
711 let revealed = self.revealed_trie_mut()?;
712
713 let (root, updates) = (revealed.root(), revealed.take_updates());
714 let updates = TrieUpdates {
715 account_nodes: updates.updated_nodes,
716 removed_nodes: updates.removed_nodes,
717 storage_tries,
718 };
719 Ok((root, updates))
720 }
721
722 pub fn storage_trie_updates(&mut self) -> B256Map<StorageTrieUpdates> {
726 self.storages
727 .iter_mut()
728 .map(|(address, trie)| {
729 let trie = trie.as_revealed_mut().unwrap();
730 let updates = trie.take_updates();
731 let updates = StorageTrieUpdates {
732 is_deleted: updates.wiped,
733 storage_nodes: updates.updated_nodes,
734 removed_nodes: updates.removed_nodes,
735 };
736 (*address, updates)
737 })
738 .filter(|(_, updates)| !updates.is_empty())
739 .collect()
740 }
741
742 pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
746 let storage_tries = self.storage_trie_updates();
747 self.state.as_revealed_mut().map(|state| {
748 let updates = state.take_updates();
749 TrieUpdates {
750 account_nodes: updates.updated_nodes,
751 removed_nodes: updates.removed_nodes,
752 storage_tries,
753 }
754 })
755 }
756
757 pub fn update_account_leaf(
759 &mut self,
760 path: Nibbles,
761 value: Vec<u8>,
762 ) -> SparseStateTrieResult<()> {
763 if !self.revealed_account_paths.contains(&path) {
764 self.revealed_account_paths.insert(path.clone());
765 }
766 let is_private = false; self.state.update_leaf(path, value, is_private)?;
769 Ok(())
770 }
771
772 pub fn update_storage_leaf(
774 &mut self,
775 address: B256,
776 slot: Nibbles,
777 value: Vec<u8>,
778 is_private: bool,
779 ) -> SparseStateTrieResult<()> {
780 if !self.revealed_storage_paths.get(&address).is_some_and(|slots| slots.contains(&slot)) {
781 self.revealed_storage_paths.entry(address).or_default().insert(slot.clone());
782 }
783
784 let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
785 storage_trie.update_leaf(slot, value, is_private)?;
786 Ok(())
787 }
788
789 pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> {
794 let nibbles = Nibbles::unpack(address);
795
796 let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
797 trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
798 storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
799 } else if self.is_account_revealed(address) {
800 trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account");
801 if let Some(value) = self.get_account_value(&address) {
803 TrieAccount::decode(&mut &value[..])?.storage_root
805 } else {
806 EMPTY_ROOT_HASH
808 }
809 } else {
810 return Err(SparseTrieErrorKind::Blind.into());
811 };
812
813 if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
814 trace!(target: "trie::sparse", ?address, "Removing account");
815 self.remove_account_leaf(&nibbles)
816 } else {
817 trace!(target: "trie::sparse", ?address, "Updating account");
818 self.account_rlp_buf.clear();
819 account.into_trie_account(storage_root).encode(&mut self.account_rlp_buf);
820 self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
821 }
822 }
823
824 pub fn update_account_storage_root(&mut self, address: B256) -> SparseStateTrieResult<()> {
831 if !self.is_account_revealed(address) {
832 return Err(SparseTrieErrorKind::Blind.into());
833 }
834
835 let Some(mut trie_account) = self
837 .get_account_value(&address)
838 .map(|v| TrieAccount::decode(&mut &v[..]))
839 .transpose()?
840 else {
841 trace!(target: "trie::sparse", ?address, "Account not found in trie, skipping storage root update");
842 return Ok(());
843 };
844
845 let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
848 trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
849 storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
850 } else {
851 EMPTY_ROOT_HASH
852 };
853
854 trie_account.storage_root = storage_root;
856
857 let nibbles = Nibbles::unpack(address);
858 if trie_account == TrieAccount::default() {
859 trace!(target: "trie::sparse", ?address, "Removing account because the storage root is empty");
861 self.remove_account_leaf(&nibbles)?;
862 } else {
863 trace!(target: "trie::sparse", ?address, "Updating account with the new storage root");
865 self.account_rlp_buf.clear();
866 trie_account.encode(&mut self.account_rlp_buf);
867 self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?;
868 }
869
870 Ok(())
871 }
872
873 pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
875 self.state.remove_leaf(path)?;
876 Ok(())
877 }
878
879 pub fn remove_storage_leaf(
881 &mut self,
882 address: B256,
883 slot: &Nibbles,
884 ) -> SparseStateTrieResult<()> {
885 let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
886 storage_trie.remove_leaf(slot)?;
887 Ok(())
888 }
889}
890
891#[derive(Debug, PartialEq, Eq)]
893struct FilteredProofNodes {
894 nodes: Vec<(Nibbles, TrieNode)>,
896 total_nodes: usize,
898 skipped_nodes: usize,
900 new_nodes: usize,
903}
904
905fn filter_revealed_nodes(
908 proof_nodes: DecodedProofNodes,
909 revealed_nodes: &HashSet<Nibbles>,
910) -> alloy_rlp::Result<FilteredProofNodes> {
911 let mut result = FilteredProofNodes {
912 nodes: Vec::with_capacity(proof_nodes.len()),
913 total_nodes: 0,
914 skipped_nodes: 0,
915 new_nodes: 0,
916 };
917
918 for (path, node) in proof_nodes.into_inner() {
919 result.total_nodes += 1;
920 if revealed_nodes.contains(&path) {
922 result.skipped_nodes += 1;
923 continue
924 }
925
926 result.new_nodes += 1;
927 if let TrieNode::Branch(branch) = &node {
930 result.new_nodes += branch.state_mask.count_ones() as usize;
931 }
932
933 result.nodes.push((path, node));
934 }
935
936 result.nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
937 Ok(result)
938}
939
940#[cfg(test)]
941mod tests {
942 use super::*;
943 use alloy_primitives::{
944 b256,
945 map::{HashMap, HashSet},
946 Bytes, U256,
947 };
948 use alloy_rlp::EMPTY_STRING_CODE;
949 use arbitrary::Arbitrary;
950 use assert_matches::assert_matches;
951 use rand::{rngs::StdRng, Rng, SeedableRng};
952 use reth_primitives_traits::Account;
953 use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH};
954 use reth_trie_common::{
955 proof::{ProofNodes, ProofRetainer},
956 BranchNode, LeafNode, StorageMultiProof, TrieMask,
957 };
958
959 #[test]
960 fn validate_root_node_first_node_not_root() {
961 let sparse = SparseStateTrie::default();
962 let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
963 assert_matches!(
964 sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
965 Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
966 );
967 }
968
969 #[test]
970 fn validate_root_node_invalid_proof_with_empty_root() {
971 let sparse = SparseStateTrie::default();
972 let proof = [
973 (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
974 (Nibbles::from_nibbles([0x1]), Bytes::new()),
975 ];
976 assert_matches!(
977 sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
978 Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
979 );
980 }
981
982 #[test]
983 fn reveal_account_empty() {
984 let retainer = ProofRetainer::from_iter([Nibbles::default()]);
985 let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
986 hash_builder.root();
987 let proofs = hash_builder.take_proof_nodes();
988 assert_eq!(proofs.len(), 1);
989
990 let mut sparse = SparseStateTrie::default();
991 assert_eq!(sparse.state, SparseTrie::Blind);
992
993 sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
994 assert_eq!(sparse.state, SparseTrie::revealed_empty());
995 }
996
997 #[test]
998 fn reveal_storage_slot_empty() {
999 let retainer = ProofRetainer::from_iter([Nibbles::default()]);
1000 let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
1001 hash_builder.root();
1002 let proofs = hash_builder.take_proof_nodes();
1003 assert_eq!(proofs.len(), 1);
1004
1005 let mut sparse = SparseStateTrie::default();
1006 assert!(sparse.storages.is_empty());
1007
1008 sparse
1009 .reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
1010 .unwrap();
1011 assert_eq!(
1012 sparse.storages,
1013 HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
1014 );
1015 }
1016
1017 #[test]
1018 fn reveal_account_path_twice() {
1019 let is_private = false; let mut sparse = SparseStateTrie::default();
1021
1022 let leaf_value = alloy_rlp::encode(TrieAccount::default());
1023 let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1024 Nibbles::default(),
1025 leaf_value.clone(),
1026 is_private,
1027 )));
1028 let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1029 Nibbles::default(),
1030 leaf_value.clone(),
1031 is_private,
1032 )));
1033
1034 let multiproof = MultiProof {
1035 account_subtree: ProofNodes::from_iter([
1036 (
1037 Nibbles::default(),
1038 alloy_rlp::encode(TrieNode::Branch(BranchNode {
1039 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
1040 state_mask: TrieMask::new(0b11),
1041 }))
1042 .into(),
1043 ),
1044 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
1045 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
1046 ]),
1047 ..Default::default()
1048 };
1049
1050 sparse.reveal_decoded_multiproof(multiproof.clone().try_into().unwrap()).unwrap();
1052 assert!(sparse
1053 .state_trie_ref()
1054 .unwrap()
1055 .nodes_ref()
1056 .contains_key(&Nibbles::from_nibbles([0x0])),);
1057 assert_eq!(
1058 sparse.state_trie_ref().unwrap().get_leaf_value(&Nibbles::from_nibbles([0x0])),
1059 Some(&leaf_value)
1060 );
1061
1062 sparse.remove_account_leaf(&Nibbles::from_nibbles([0x0])).unwrap();
1065 assert!(!sparse
1066 .state_trie_ref()
1067 .unwrap()
1068 .nodes_ref()
1069 .contains_key(&Nibbles::from_nibbles([0x0])),);
1070 assert!(sparse
1071 .state_trie_ref()
1072 .unwrap()
1073 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1074 .is_none());
1075
1076 sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
1079 assert!(!sparse
1080 .state_trie_ref()
1081 .unwrap()
1082 .nodes_ref()
1083 .contains_key(&Nibbles::from_nibbles([0x0])));
1084 assert!(sparse
1085 .state_trie_ref()
1086 .unwrap()
1087 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1088 .is_none());
1089 }
1090
1091 #[test]
1092 fn reveal_storage_path_twice() {
1093 let is_private = false; let mut sparse = SparseStateTrie::default();
1095
1096 let leaf_value = alloy_rlp::encode(TrieAccount::default());
1097 let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1098 Nibbles::default(),
1099 leaf_value.clone(),
1100 is_private,
1101 )));
1102 let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1103 Nibbles::default(),
1104 leaf_value.clone(),
1105 is_private,
1106 )));
1107
1108 let multiproof = MultiProof {
1109 storages: HashMap::from_iter([(
1110 B256::ZERO,
1111 StorageMultiProof {
1112 root: B256::ZERO,
1113 subtree: ProofNodes::from_iter([
1114 (
1115 Nibbles::default(),
1116 alloy_rlp::encode(TrieNode::Branch(BranchNode {
1117 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
1118 state_mask: TrieMask::new(0b11),
1119 }))
1120 .into(),
1121 ),
1122 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
1123 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
1124 ]),
1125 branch_node_hash_masks: Default::default(),
1126 branch_node_tree_masks: Default::default(),
1127 },
1128 )]),
1129 ..Default::default()
1130 };
1131
1132 sparse.reveal_decoded_multiproof(multiproof.clone().try_into().unwrap()).unwrap();
1134 assert!(sparse
1135 .storage_trie_ref(&B256::ZERO)
1136 .unwrap()
1137 .nodes_ref()
1138 .contains_key(&Nibbles::from_nibbles([0x0])),);
1139 assert_eq!(
1140 sparse
1141 .storage_trie_ref(&B256::ZERO)
1142 .unwrap()
1143 .get_leaf_value(&Nibbles::from_nibbles([0x0])),
1144 Some(&leaf_value)
1145 );
1146
1147 sparse.remove_storage_leaf(B256::ZERO, &Nibbles::from_nibbles([0x0])).unwrap();
1150 assert!(!sparse
1151 .storage_trie_ref(&B256::ZERO)
1152 .unwrap()
1153 .nodes_ref()
1154 .contains_key(&Nibbles::from_nibbles([0x0])),);
1155 assert!(sparse
1156 .storage_trie_ref(&B256::ZERO)
1157 .unwrap()
1158 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1159 .is_none());
1160
1161 sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
1164 assert!(!sparse
1165 .storage_trie_ref(&B256::ZERO)
1166 .unwrap()
1167 .nodes_ref()
1168 .contains_key(&Nibbles::from_nibbles([0x0])));
1169 assert!(sparse
1170 .storage_trie_ref(&B256::ZERO)
1171 .unwrap()
1172 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1173 .is_none());
1174 }
1175
1176 #[test]
1177 fn take_trie_updates() {
1178 reth_tracing::init_test_tracing();
1179
1180 let mut rng = StdRng::seed_from_u64(1);
1182
1183 let mut bytes = [0u8; 1024];
1184 rng.fill(bytes.as_mut_slice());
1185
1186 let slot_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1187 let slot_path_1 = Nibbles::unpack(slot_1);
1188 let value_1 = U256::from(rng.random::<u64>());
1189 let slot_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1190 let slot_path_2 = Nibbles::unpack(slot_2);
1191 let value_2 = U256::from(rng.random::<u64>());
1192 let slot_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1193 let slot_path_3 = Nibbles::unpack(slot_3);
1194 let value_3 = U256::from(rng.random::<u64>());
1195
1196 let mut storage_hash_builder =
1197 HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1198 slot_path_1.clone(),
1199 slot_path_2.clone(),
1200 ]));
1201 let is_private = false; storage_hash_builder.add_leaf(
1203 slot_path_1,
1204 &alloy_rlp::encode_fixed_size(&value_1),
1205 is_private,
1206 );
1207 storage_hash_builder.add_leaf(
1208 slot_path_2,
1209 &alloy_rlp::encode_fixed_size(&value_2),
1210 is_private,
1211 );
1212
1213 let storage_root = storage_hash_builder.root();
1214 let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
1215 let storage_branch_node_hash_masks = HashMap::from_iter([
1216 (Nibbles::default(), TrieMask::new(0b010)),
1217 (Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)),
1218 ]);
1219
1220 let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1221 let address_path_1 = Nibbles::unpack(address_1);
1222 let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1223 let mut trie_account_1 = account_1.into_trie_account(storage_root);
1224 let address_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1225 let address_path_2 = Nibbles::unpack(address_2);
1226 let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1227 let mut trie_account_2 = account_2.into_trie_account(EMPTY_ROOT_HASH);
1228
1229 let mut hash_builder =
1230 HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1231 address_path_1.clone(),
1232 address_path_2.clone(),
1233 ]));
1234 let is_private = false; hash_builder.add_leaf(
1236 address_path_1.clone(),
1237 &alloy_rlp::encode(trie_account_1),
1238 is_private,
1239 );
1240 hash_builder.add_leaf(
1241 address_path_2.clone(),
1242 &alloy_rlp::encode(trie_account_2),
1243 is_private,
1244 );
1245
1246 let root = hash_builder.root();
1247 let proof_nodes = hash_builder.take_proof_nodes();
1248
1249 let mut sparse = SparseStateTrie::default().with_updates(true);
1250 sparse
1251 .reveal_decoded_multiproof(
1252 MultiProof {
1253 account_subtree: proof_nodes,
1254 branch_node_hash_masks: HashMap::from_iter([(
1255 Nibbles::from_nibbles([0x1]),
1256 TrieMask::new(0b00),
1257 )]),
1258 branch_node_tree_masks: HashMap::default(),
1259 storages: HashMap::from_iter([
1260 (
1261 address_1,
1262 StorageMultiProof {
1263 root,
1264 subtree: storage_proof_nodes.clone(),
1265 branch_node_hash_masks: storage_branch_node_hash_masks.clone(),
1266 branch_node_tree_masks: HashMap::default(),
1267 },
1268 ),
1269 (
1270 address_2,
1271 StorageMultiProof {
1272 root,
1273 subtree: storage_proof_nodes,
1274 branch_node_hash_masks: storage_branch_node_hash_masks,
1275 branch_node_tree_masks: HashMap::default(),
1276 },
1277 ),
1278 ]),
1279 }
1280 .try_into()
1281 .unwrap(),
1282 )
1283 .unwrap();
1284
1285 assert_eq!(sparse.root().unwrap(), root);
1286
1287 let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1288 let address_path_3 = Nibbles::unpack(address_3);
1289 let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
1290 let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH);
1291
1292 sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap();
1293
1294 let is_private = false; sparse
1296 .update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3), is_private)
1297 .unwrap();
1298 trie_account_1.storage_root = sparse.storage_root(address_1).unwrap();
1299 sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap();
1300
1301 sparse.wipe_storage(address_2).unwrap();
1302 trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
1303 sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap();
1304
1305 sparse.root().unwrap();
1306
1307 let sparse_updates = sparse.take_trie_updates().unwrap();
1308 pretty_assertions::assert_eq!(
1310 sparse_updates,
1311 TrieUpdates {
1312 account_nodes: HashMap::default(),
1313 storage_tries: HashMap::from_iter([(
1314 b256!("0x1100000000000000000000000000000000000000000000000000000000000000"),
1315 StorageTrieUpdates {
1316 is_deleted: true,
1317 storage_nodes: HashMap::default(),
1318 removed_nodes: HashSet::default()
1319 }
1320 )]),
1321 removed_nodes: HashSet::default()
1322 }
1323 );
1324 }
1325
1326 #[test]
1327 fn test_filter_revealed_nodes() {
1328 let is_private = false; let revealed_nodes = HashSet::from_iter([Nibbles::from_nibbles([0x0])]);
1330 let leaf =
1331 TrieNode::Leaf(LeafNode::new(Nibbles::default(), alloy_rlp::encode([]), is_private));
1332 let leaf_encoded = alloy_rlp::encode(&leaf);
1333 let branch = TrieNode::Branch(BranchNode::new(
1334 vec![RlpNode::from_rlp(&leaf_encoded), RlpNode::from_rlp(&leaf_encoded)],
1335 TrieMask::new(0b11),
1336 ));
1337 let proof_nodes = alloy_trie::proof::DecodedProofNodes::from_iter([
1338 (Nibbles::default(), branch.clone()),
1339 (Nibbles::from_nibbles([0x0]), leaf.clone()),
1340 (Nibbles::from_nibbles([0x1]), leaf.clone()),
1341 ]);
1342
1343 let decoded = filter_revealed_nodes(proof_nodes, &revealed_nodes).unwrap();
1344
1345 assert_eq!(
1346 decoded,
1347 FilteredProofNodes {
1348 nodes: vec![(Nibbles::default(), branch), (Nibbles::from_nibbles([0x1]), leaf)],
1349 total_nodes: 3,
1351 skipped_nodes: 1,
1353 new_nodes: 4
1355 }
1356 );
1357 }
1358}