1use crate::{
2 blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
3 metrics::SparseStateTrieMetrics,
4 RevealedSparseTrie, SparseTrie, TrieMasks,
5};
6use alloy_primitives::{
7 hex,
8 map::{B256Map, HashMap, HashSet},
9 Bytes, B256,
10};
11use alloy_rlp::{Decodable, Encodable};
12use core::fmt;
13use reth_execution_errors::{SparseStateTrieErrorKind, SparseStateTrieResult, SparseTrieErrorKind};
14use reth_primitives_traits::Account;
15use reth_tracing::tracing::trace;
16use reth_trie_common::{
17 proof::ProofNodes,
18 updates::{StorageTrieUpdates, TrieUpdates},
19 MultiProof, Nibbles, RlpNode, StorageMultiProof, TrieAccount, TrieMask, TrieNode,
20 EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
21};
22use std::{collections::VecDeque, iter::Peekable};
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 metrics: SparseStateTrieMetrics,
42}
43
44#[cfg(test)]
45impl Default for SparseStateTrie {
46 fn default() -> Self {
47 Self {
48 provider_factory: DefaultBlindedProviderFactory,
49 state: Default::default(),
50 storages: Default::default(),
51 revealed_account_paths: Default::default(),
52 revealed_storage_paths: Default::default(),
53 retain_updates: false,
54 account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
55 metrics: Default::default(),
56 }
57 }
58}
59
60impl<P: BlindedProviderFactory> fmt::Debug for SparseStateTrie<P> {
61 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62 f.debug_struct("SparseStateTrie")
63 .field("state", &self.state)
64 .field("storages", &self.storages)
65 .field("revealed_account_paths", &self.revealed_account_paths)
66 .field("revealed_storage_paths", &self.revealed_storage_paths)
67 .field("retain_updates", &self.retain_updates)
68 .field("account_rlp_buf", &hex::encode(&self.account_rlp_buf))
69 .finish_non_exhaustive()
70 }
71}
72
73#[cfg(test)]
74impl SparseStateTrie {
75 pub fn from_state(state: SparseTrie) -> Self {
77 Self { state, ..Default::default() }
78 }
79}
80
81impl<F: BlindedProviderFactory> SparseStateTrie<F> {
82 pub fn new(provider_factory: F) -> Self {
84 Self {
85 provider_factory,
86 state: Default::default(),
87 storages: Default::default(),
88 revealed_account_paths: Default::default(),
89 revealed_storage_paths: Default::default(),
90 retain_updates: false,
91 account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
92 metrics: Default::default(),
93 }
94 }
95
96 pub const fn with_updates(mut self, retain_updates: bool) -> Self {
98 self.retain_updates = retain_updates;
99 self
100 }
101
102 pub fn is_account_revealed(&self, account: B256) -> bool {
104 self.revealed_account_paths.contains(&Nibbles::unpack(account))
105 }
106
107 pub fn is_storage_slot_revealed(&self, account: B256, slot: B256) -> bool {
109 self.revealed_storage_paths
110 .get(&account)
111 .is_some_and(|slots| slots.contains(&Nibbles::unpack(slot)))
112 }
113
114 pub fn get_account_value(&self, account: &B256) -> Option<&Vec<u8>> {
116 self.state.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(account))
117 }
118
119 pub fn get_storage_slot_value(&self, account: &B256, slot: &B256) -> Option<&Vec<u8>> {
121 self.storages.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
122 }
123
124 pub const fn state_trie_ref(&self) -> Option<&RevealedSparseTrie<F::AccountNodeProvider>> {
126 self.state.as_revealed_ref()
127 }
128
129 pub fn storage_trie_ref(
131 &self,
132 address: &B256,
133 ) -> Option<&RevealedSparseTrie<F::StorageNodeProvider>> {
134 self.storages.get(address).and_then(|e| e.as_revealed_ref())
135 }
136
137 pub fn storage_trie_mut(
139 &mut self,
140 address: &B256,
141 ) -> Option<&mut RevealedSparseTrie<F::StorageNodeProvider>> {
142 self.storages.get_mut(address).and_then(|e| e.as_revealed_mut())
143 }
144
145 pub fn take_storage_trie(
147 &mut self,
148 address: &B256,
149 ) -> Option<SparseTrie<F::StorageNodeProvider>> {
150 self.storages.remove(address)
151 }
152
153 pub fn insert_storage_trie(
155 &mut self,
156 address: B256,
157 storage_trie: SparseTrie<F::StorageNodeProvider>,
158 ) {
159 self.storages.insert(address, storage_trie);
160 }
161
162 pub fn reveal_account(
168 &mut self,
169 account: B256,
170 proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
171 ) -> SparseStateTrieResult<()> {
172 assert!(!self.retain_updates);
173
174 if self.is_account_revealed(account) {
175 return Ok(());
176 }
177
178 let mut proof = proof.into_iter().peekable();
179
180 let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
181
182 let trie = self.state.reveal_root_with_provider(
184 self.provider_factory.account_node_provider(),
185 root_node,
186 TrieMasks::none(),
187 self.retain_updates,
188 )?;
189
190 for (path, bytes) in proof {
192 if self.revealed_account_paths.contains(&path) {
193 continue
194 }
195 let node = TrieNode::decode(&mut &bytes[..])?;
196 trie.reveal_node(path.clone(), node, TrieMasks::none())?;
197
198 self.revealed_account_paths.insert(path);
200 }
201
202 Ok(())
203 }
204
205 pub fn reveal_storage_slot(
211 &mut self,
212 account: B256,
213 slot: B256,
214 proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
215 ) -> SparseStateTrieResult<()> {
216 assert!(!self.retain_updates);
217
218 if self.is_storage_slot_revealed(account, slot) {
219 return Ok(());
220 }
221
222 let mut proof = proof.into_iter().peekable();
223
224 let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
225
226 let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
228 self.provider_factory.storage_node_provider(account),
229 root_node,
230 TrieMasks::none(),
231 self.retain_updates,
232 )?;
233
234 let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
235
236 for (path, bytes) in proof {
238 if revealed_nodes.contains(&path) {
240 continue
241 }
242 let node = TrieNode::decode(&mut &bytes[..])?;
243 trie.reveal_node(path.clone(), node, TrieMasks::none())?;
244
245 revealed_nodes.insert(path);
247 }
248
249 Ok(())
250 }
251
252 pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
255 let MultiProof {
256 account_subtree,
257 storages,
258 branch_node_hash_masks,
259 branch_node_tree_masks,
260 } = multiproof;
261
262 self.reveal_account_multiproof(
264 account_subtree,
265 branch_node_hash_masks,
266 branch_node_tree_masks,
267 )?;
268
269 for (account, storage_subtree) in storages {
271 self.reveal_storage_multiproof(account, storage_subtree)?;
272 }
273
274 Ok(())
275 }
276
277 pub fn reveal_account_multiproof(
279 &mut self,
280 account_subtree: ProofNodes,
281 branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
282 branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
283 ) -> SparseStateTrieResult<()> {
284 let account_subtree = account_subtree.into_nodes_sorted();
285 let mut account_nodes = account_subtree.into_iter().peekable();
286
287 if let Some(root_node) = self.validate_root_node(&mut account_nodes)? {
288 let trie = self.state.reveal_root_with_provider(
290 self.provider_factory.account_node_provider(),
291 root_node,
292 TrieMasks {
293 hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
294 tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
295 },
296 self.retain_updates,
297 )?;
298
299 for (path, bytes) in account_nodes {
301 self.metrics.increment_total_account_nodes();
302 if self.revealed_account_paths.contains(&path) {
304 self.metrics.increment_skipped_account_nodes();
305 continue
306 }
307 let node = TrieNode::decode(&mut &bytes[..])?;
308 let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
309 (
310 branch_node_hash_masks.get(&path).copied(),
311 branch_node_tree_masks.get(&path).copied(),
312 )
313 } else {
314 (None, None)
315 };
316
317 trace!(target: "trie::sparse", ?path, ?node, ?hash_mask, ?tree_mask, "Revealing account node");
318 trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
319
320 self.revealed_account_paths.insert(path);
322 }
323 }
324
325 Ok(())
326 }
327
328 pub fn reveal_storage_multiproof(
330 &mut self,
331 account: B256,
332 storage_subtree: StorageMultiProof,
333 ) -> SparseStateTrieResult<()> {
334 let subtree = storage_subtree.subtree.into_nodes_sorted();
335 let mut nodes = subtree.into_iter().peekable();
336
337 if let Some(root_node) = self.validate_root_node(&mut nodes)? {
338 let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
340 self.provider_factory.storage_node_provider(account),
341 root_node,
342 TrieMasks {
343 hash_mask: storage_subtree
344 .branch_node_hash_masks
345 .get(&Nibbles::default())
346 .copied(),
347 tree_mask: storage_subtree
348 .branch_node_tree_masks
349 .get(&Nibbles::default())
350 .copied(),
351 },
352 self.retain_updates,
353 )?;
354 let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
355
356 for (path, bytes) in nodes {
358 self.metrics.increment_total_storage_nodes();
359 if revealed_nodes.contains(&path) {
361 self.metrics.increment_skipped_storage_nodes();
362 continue
363 }
364 let node = TrieNode::decode(&mut &bytes[..])?;
365 let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
366 (
367 storage_subtree.branch_node_hash_masks.get(&path).copied(),
368 storage_subtree.branch_node_tree_masks.get(&path).copied(),
369 )
370 } else {
371 (None, None)
372 };
373
374 trace!(target: "trie::sparse", ?account, ?path, ?node, ?hash_mask, ?tree_mask, "Revealing storage node");
375 trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
376
377 revealed_nodes.insert(path);
379 }
380 }
381
382 Ok(())
383 }
384
385 pub fn reveal_witness(
389 &mut self,
390 state_root: B256,
391 witness: &B256Map<Bytes>,
392 ) -> SparseStateTrieResult<()> {
393 let mut queue = VecDeque::from([(state_root, Nibbles::default(), None)]);
396
397 while let Some((hash, path, maybe_account)) = queue.pop_front() {
398 let Some(trie_node_bytes) = witness.get(&hash) else { continue };
400 let trie_node = TrieNode::decode(&mut &trie_node_bytes[..])?;
401
402 match &trie_node {
404 TrieNode::Branch(branch) => {
405 for (idx, maybe_child) in branch.as_ref().children() {
406 if let Some(child_hash) = maybe_child.and_then(RlpNode::as_hash) {
407 let mut child_path = path.clone();
408 child_path.push_unchecked(idx);
409 queue.push_back((child_hash, child_path, maybe_account));
410 }
411 }
412 }
413 TrieNode::Extension(ext) => {
414 if let Some(child_hash) = ext.child.as_hash() {
415 let mut child_path = path.clone();
416 child_path.extend_from_slice_unchecked(&ext.key);
417 queue.push_back((child_hash, child_path, maybe_account));
418 }
419 }
420 TrieNode::Leaf(leaf) => {
421 let mut full_path = path.clone();
422 full_path.extend_from_slice_unchecked(&leaf.key);
423 if maybe_account.is_none() {
424 let hashed_address = B256::from_slice(&full_path.pack());
425 let account = TrieAccount::decode(&mut &leaf.value[..])?;
426 if account.storage_root != EMPTY_ROOT_HASH {
427 queue.push_back((
428 account.storage_root,
429 Nibbles::default(),
430 Some(hashed_address),
431 ));
432 }
433 }
434 }
435 TrieNode::EmptyRoot => {} };
437
438 if let Some(account) = maybe_account {
440 if self
442 .revealed_storage_paths
443 .get(&account)
444 .is_none_or(|paths| !paths.contains(&path))
445 {
446 let storage_trie_entry = self.storages.entry(account).or_default();
447 if path.is_empty() {
448 storage_trie_entry.reveal_root_with_provider(
450 self.provider_factory.storage_node_provider(account),
451 trie_node,
452 TrieMasks::none(),
453 self.retain_updates,
454 )?;
455 } else {
456 storage_trie_entry
458 .as_revealed_mut()
459 .ok_or(SparseTrieErrorKind::Blind)?
460 .reveal_node(path.clone(), trie_node, TrieMasks::none())?;
461 }
462
463 self.revealed_storage_paths.entry(account).or_default().insert(path);
465 }
466 }
467 else if !self.revealed_account_paths.contains(&path) {
469 if path.is_empty() {
470 self.state.reveal_root_with_provider(
472 self.provider_factory.account_node_provider(),
473 trie_node,
474 TrieMasks::none(),
475 self.retain_updates,
476 )?;
477 } else {
478 self.state.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?.reveal_node(
480 path.clone(),
481 trie_node,
482 TrieMasks::none(),
483 )?;
484 }
485
486 self.revealed_account_paths.insert(path);
488 }
489 }
490
491 Ok(())
492 }
493
494 fn validate_root_node<I: Iterator<Item = (Nibbles, Bytes)>>(
496 &self,
497 proof: &mut Peekable<I>,
498 ) -> SparseStateTrieResult<Option<TrieNode>> {
499 let mut proof = proof.into_iter().peekable();
500
501 let Some((path, node)) = proof.next() else { return Ok(None) };
503 if !path.is_empty() {
504 return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into())
505 }
506
507 let root_node = TrieNode::decode(&mut &node[..])?;
509 if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
510 return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into())
511 }
512
513 Ok(Some(root_node))
514 }
515
516 pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
518 if let Some(trie) = self.storages.get_mut(&address) {
519 trie.wipe()?;
520 }
521 Ok(())
522 }
523
524 pub fn calculate_below_level(&mut self, level: usize) {
528 if let SparseTrie::Revealed(trie) = &mut self.state {
529 trie.update_rlp_node_level(level);
530 }
531 }
532
533 pub fn storage_root(&mut self, account: B256) -> Option<B256> {
535 self.storages.get_mut(&account).and_then(|trie| trie.root())
536 }
537
538 fn revealed_trie_mut(
542 &mut self,
543 ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
544 match self.state {
545 SparseTrie::Blind => {
546 let (root_node, hash_mask, tree_mask) = self
547 .provider_factory
548 .account_node_provider()
549 .blinded_node(&Nibbles::default())?
550 .map(|node| {
551 TrieNode::decode(&mut &node.node[..])
552 .map(|decoded| (decoded, node.hash_mask, node.tree_mask))
553 })
554 .transpose()?
555 .unwrap_or((TrieNode::EmptyRoot, None, None));
556 self.state
557 .reveal_root_with_provider(
558 self.provider_factory.account_node_provider(),
559 root_node,
560 TrieMasks { hash_mask, tree_mask },
561 self.retain_updates,
562 )
563 .map_err(Into::into)
564 }
565 SparseTrie::Revealed(ref mut trie) => Ok(trie),
566 }
567 }
568
569 pub fn root(&mut self) -> SparseStateTrieResult<B256> {
573 self.metrics.record();
575
576 Ok(self.revealed_trie_mut()?.root())
577 }
578
579 pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
581 self.metrics.record();
583
584 let storage_tries = self.storage_trie_updates();
585 let revealed = self.revealed_trie_mut()?;
586
587 let (root, updates) = (revealed.root(), revealed.take_updates());
588 let updates = TrieUpdates {
589 account_nodes: updates.updated_nodes,
590 removed_nodes: updates.removed_nodes,
591 storage_tries,
592 };
593 Ok((root, updates))
594 }
595
596 pub fn storage_trie_updates(&mut self) -> B256Map<StorageTrieUpdates> {
600 self.storages
601 .iter_mut()
602 .map(|(address, trie)| {
603 let trie = trie.as_revealed_mut().unwrap();
604 let updates = trie.take_updates();
605 let updates = StorageTrieUpdates {
606 is_deleted: updates.wiped,
607 storage_nodes: updates.updated_nodes,
608 removed_nodes: updates.removed_nodes,
609 };
610 (*address, updates)
611 })
612 .filter(|(_, updates)| !updates.is_empty())
613 .collect()
614 }
615
616 pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
620 let storage_tries = self.storage_trie_updates();
621 self.state.as_revealed_mut().map(|state| {
622 let updates = state.take_updates();
623 TrieUpdates {
624 account_nodes: updates.updated_nodes,
625 removed_nodes: updates.removed_nodes,
626 storage_tries,
627 }
628 })
629 }
630
631 pub fn update_account_leaf(
633 &mut self,
634 path: Nibbles,
635 value: Vec<u8>,
636 ) -> SparseStateTrieResult<()> {
637 if !self.revealed_account_paths.contains(&path) {
638 self.revealed_account_paths.insert(path.clone());
639 }
640
641 self.state.update_leaf(path, value)?;
642 Ok(())
643 }
644
645 pub fn update_storage_leaf(
647 &mut self,
648 address: B256,
649 slot: Nibbles,
650 value: Vec<u8>,
651 ) -> SparseStateTrieResult<()> {
652 if !self.revealed_storage_paths.get(&address).is_some_and(|slots| slots.contains(&slot)) {
653 self.revealed_storage_paths.entry(address).or_default().insert(slot.clone());
654 }
655
656 let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
657 storage_trie.update_leaf(slot, value)?;
658 Ok(())
659 }
660
661 pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> {
666 let nibbles = Nibbles::unpack(address);
667
668 let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
669 trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
670 storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
671 } else if self.is_account_revealed(address) {
672 trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account");
673 if let Some(value) = self.get_account_value(&address) {
675 TrieAccount::decode(&mut &value[..])?.storage_root
677 } else {
678 EMPTY_ROOT_HASH
680 }
681 } else {
682 return Err(SparseTrieErrorKind::Blind.into())
683 };
684
685 if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
686 trace!(target: "trie::sparse", ?address, "Removing account");
687 self.remove_account_leaf(&nibbles)
688 } else {
689 trace!(target: "trie::sparse", ?address, "Updating account");
690 self.account_rlp_buf.clear();
691 account.into_trie_account(storage_root).encode(&mut self.account_rlp_buf);
692 self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
693 }
694 }
695
696 pub fn update_account_storage_root(&mut self, address: B256) -> SparseStateTrieResult<()> {
703 if !self.is_account_revealed(address) {
704 return Err(SparseTrieErrorKind::Blind.into())
705 }
706
707 let Some(mut trie_account) = self
709 .get_account_value(&address)
710 .map(|v| TrieAccount::decode(&mut &v[..]))
711 .transpose()?
712 else {
713 trace!(target: "trie::sparse", ?address, "Account not found in trie, skipping storage root update");
714 return Ok(())
715 };
716
717 let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
720 trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
721 storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
722 } else {
723 EMPTY_ROOT_HASH
724 };
725
726 trie_account.storage_root = storage_root;
728
729 let nibbles = Nibbles::unpack(address);
730 if trie_account == TrieAccount::default() {
731 trace!(target: "trie::sparse", ?address, "Removing account because the storage root is empty");
733 self.remove_account_leaf(&nibbles)?;
734 } else {
735 trace!(target: "trie::sparse", ?address, "Updating account with the new storage root");
737 self.account_rlp_buf.clear();
738 trie_account.encode(&mut self.account_rlp_buf);
739 self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?;
740 }
741
742 Ok(())
743 }
744
745 pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
747 self.state.remove_leaf(path)?;
748 Ok(())
749 }
750
751 pub fn remove_storage_leaf(
753 &mut self,
754 address: B256,
755 slot: &Nibbles,
756 ) -> SparseStateTrieResult<()> {
757 let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
758 storage_trie.remove_leaf(slot)?;
759 Ok(())
760 }
761}
762
763#[cfg(test)]
764mod tests {
765 use super::*;
766 use alloy_primitives::{
767 b256,
768 map::{HashMap, HashSet},
769 Bytes, U256,
770 };
771 use alloy_rlp::EMPTY_STRING_CODE;
772 use arbitrary::Arbitrary;
773 use assert_matches::assert_matches;
774 use rand::{rngs::StdRng, Rng, SeedableRng};
775 use reth_primitives_traits::Account;
776 use reth_trie::{updates::StorageTrieUpdates, HashBuilder, EMPTY_ROOT_HASH};
777 use reth_trie_common::{
778 proof::{ProofNodes, ProofRetainer},
779 BranchNode, LeafNode, StorageMultiProof, TrieMask,
780 };
781
782 #[test]
783 fn validate_root_node_first_node_not_root() {
784 let sparse = SparseStateTrie::default();
785 let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
786 assert_matches!(
787 sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
788 Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
789 );
790 }
791
792 #[test]
793 fn validate_root_node_invalid_proof_with_empty_root() {
794 let sparse = SparseStateTrie::default();
795 let proof = [
796 (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
797 (Nibbles::from_nibbles([0x1]), Bytes::new()),
798 ];
799 assert_matches!(
800 sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
801 Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
802 );
803 }
804
805 #[test]
806 fn reveal_account_empty() {
807 let retainer = ProofRetainer::from_iter([Nibbles::default()]);
808 let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
809 hash_builder.root();
810 let proofs = hash_builder.take_proof_nodes();
811 assert_eq!(proofs.len(), 1);
812
813 let mut sparse = SparseStateTrie::default();
814 assert_eq!(sparse.state, SparseTrie::Blind);
815
816 sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
817 assert_eq!(sparse.state, SparseTrie::revealed_empty());
818 }
819
820 #[test]
821 fn reveal_storage_slot_empty() {
822 let retainer = ProofRetainer::from_iter([Nibbles::default()]);
823 let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
824 hash_builder.root();
825 let proofs = hash_builder.take_proof_nodes();
826 assert_eq!(proofs.len(), 1);
827
828 let mut sparse = SparseStateTrie::default();
829 assert!(sparse.storages.is_empty());
830
831 sparse
832 .reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
833 .unwrap();
834 assert_eq!(
835 sparse.storages,
836 HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
837 );
838 }
839
840 #[test]
841 fn reveal_account_path_twice() {
842 let mut sparse = SparseStateTrie::default();
843
844 let leaf_value = alloy_rlp::encode(TrieAccount::default());
845 let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
846 Nibbles::default(),
847 leaf_value.clone(),
848 )));
849 let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
850 Nibbles::default(),
851 leaf_value.clone(),
852 )));
853
854 let multiproof = MultiProof {
855 account_subtree: ProofNodes::from_iter([
856 (
857 Nibbles::default(),
858 alloy_rlp::encode(TrieNode::Branch(BranchNode {
859 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
860 state_mask: TrieMask::new(0b11),
861 }))
862 .into(),
863 ),
864 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
865 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
866 ]),
867 ..Default::default()
868 };
869
870 sparse.reveal_multiproof(multiproof.clone()).unwrap();
872 assert!(sparse
873 .state_trie_ref()
874 .unwrap()
875 .nodes_ref()
876 .contains_key(&Nibbles::from_nibbles([0x0])),);
877 assert_eq!(
878 sparse.state_trie_ref().unwrap().get_leaf_value(&Nibbles::from_nibbles([0x0])),
879 Some(&leaf_value)
880 );
881
882 sparse.remove_account_leaf(&Nibbles::from_nibbles([0x0])).unwrap();
885 assert!(!sparse
886 .state_trie_ref()
887 .unwrap()
888 .nodes_ref()
889 .contains_key(&Nibbles::from_nibbles([0x0])),);
890 assert!(sparse
891 .state_trie_ref()
892 .unwrap()
893 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
894 .is_none());
895
896 sparse.reveal_multiproof(multiproof).unwrap();
899 assert!(!sparse
900 .state_trie_ref()
901 .unwrap()
902 .nodes_ref()
903 .contains_key(&Nibbles::from_nibbles([0x0])));
904 assert!(sparse
905 .state_trie_ref()
906 .unwrap()
907 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
908 .is_none());
909 }
910
911 #[test]
912 fn reveal_storage_path_twice() {
913 let mut sparse = SparseStateTrie::default();
914
915 let leaf_value = alloy_rlp::encode(TrieAccount::default());
916 let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
917 Nibbles::default(),
918 leaf_value.clone(),
919 )));
920 let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
921 Nibbles::default(),
922 leaf_value.clone(),
923 )));
924
925 let multiproof = MultiProof {
926 storages: HashMap::from_iter([(
927 B256::ZERO,
928 StorageMultiProof {
929 root: B256::ZERO,
930 subtree: ProofNodes::from_iter([
931 (
932 Nibbles::default(),
933 alloy_rlp::encode(TrieNode::Branch(BranchNode {
934 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
935 state_mask: TrieMask::new(0b11),
936 }))
937 .into(),
938 ),
939 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
940 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
941 ]),
942 branch_node_hash_masks: Default::default(),
943 branch_node_tree_masks: Default::default(),
944 },
945 )]),
946 ..Default::default()
947 };
948
949 sparse.reveal_multiproof(multiproof.clone()).unwrap();
951 assert!(sparse
952 .storage_trie_ref(&B256::ZERO)
953 .unwrap()
954 .nodes_ref()
955 .contains_key(&Nibbles::from_nibbles([0x0])),);
956 assert_eq!(
957 sparse
958 .storage_trie_ref(&B256::ZERO)
959 .unwrap()
960 .get_leaf_value(&Nibbles::from_nibbles([0x0])),
961 Some(&leaf_value)
962 );
963
964 sparse.remove_storage_leaf(B256::ZERO, &Nibbles::from_nibbles([0x0])).unwrap();
967 assert!(!sparse
968 .storage_trie_ref(&B256::ZERO)
969 .unwrap()
970 .nodes_ref()
971 .contains_key(&Nibbles::from_nibbles([0x0])),);
972 assert!(sparse
973 .storage_trie_ref(&B256::ZERO)
974 .unwrap()
975 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
976 .is_none());
977
978 sparse.reveal_multiproof(multiproof).unwrap();
981 assert!(!sparse
982 .storage_trie_ref(&B256::ZERO)
983 .unwrap()
984 .nodes_ref()
985 .contains_key(&Nibbles::from_nibbles([0x0])));
986 assert!(sparse
987 .storage_trie_ref(&B256::ZERO)
988 .unwrap()
989 .get_leaf_value(&Nibbles::from_nibbles([0x0]))
990 .is_none());
991 }
992
993 #[test]
994 fn take_trie_updates() {
995 reth_tracing::init_test_tracing();
996
997 let mut rng = StdRng::seed_from_u64(1);
999
1000 let mut bytes = [0u8; 1024];
1001 rng.fill(bytes.as_mut_slice());
1002
1003 let slot_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1004 let slot_path_1 = Nibbles::unpack(slot_1);
1005 let value_1 = U256::from(rng.gen::<u64>());
1006 let slot_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1007 let slot_path_2 = Nibbles::unpack(slot_2);
1008 let value_2 = U256::from(rng.gen::<u64>());
1009 let slot_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1010 let slot_path_3 = Nibbles::unpack(slot_3);
1011 let value_3 = U256::from(rng.gen::<u64>());
1012
1013 let mut storage_hash_builder =
1014 HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1015 slot_path_1.clone(),
1016 slot_path_2.clone(),
1017 ]));
1018 storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1));
1019 storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2));
1020
1021 let storage_root = storage_hash_builder.root();
1022 let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
1023 let storage_branch_node_hash_masks = HashMap::from_iter([
1024 (Nibbles::default(), TrieMask::new(0b010)),
1025 (Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)),
1026 ]);
1027
1028 let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1029 let address_path_1 = Nibbles::unpack(address_1);
1030 let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1031 let mut trie_account_1 = account_1.into_trie_account(storage_root);
1032 let address_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1033 let address_path_2 = Nibbles::unpack(address_2);
1034 let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1035 let mut trie_account_2 = account_2.into_trie_account(EMPTY_ROOT_HASH);
1036
1037 let mut hash_builder =
1038 HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1039 address_path_1.clone(),
1040 address_path_2.clone(),
1041 ]));
1042 hash_builder.add_leaf(address_path_1.clone(), &alloy_rlp::encode(trie_account_1));
1043 hash_builder.add_leaf(address_path_2.clone(), &alloy_rlp::encode(trie_account_2));
1044
1045 let root = hash_builder.root();
1046 let proof_nodes = hash_builder.take_proof_nodes();
1047
1048 let mut sparse = SparseStateTrie::default().with_updates(true);
1049 sparse
1050 .reveal_multiproof(MultiProof {
1051 account_subtree: proof_nodes,
1052 branch_node_hash_masks: HashMap::from_iter([(
1053 Nibbles::from_nibbles([0x1]),
1054 TrieMask::new(0b00),
1055 )]),
1056 branch_node_tree_masks: HashMap::default(),
1057 storages: HashMap::from_iter([
1058 (
1059 address_1,
1060 StorageMultiProof {
1061 root,
1062 subtree: storage_proof_nodes.clone(),
1063 branch_node_hash_masks: storage_branch_node_hash_masks.clone(),
1064 branch_node_tree_masks: HashMap::default(),
1065 },
1066 ),
1067 (
1068 address_2,
1069 StorageMultiProof {
1070 root,
1071 subtree: storage_proof_nodes,
1072 branch_node_hash_masks: storage_branch_node_hash_masks,
1073 branch_node_tree_masks: HashMap::default(),
1074 },
1075 ),
1076 ]),
1077 })
1078 .unwrap();
1079
1080 assert_eq!(sparse.root().unwrap(), root);
1081
1082 let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1083 let address_path_3 = Nibbles::unpack(address_3);
1084 let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
1085 let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH);
1086
1087 sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap();
1088
1089 sparse.update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3)).unwrap();
1090 trie_account_1.storage_root = sparse.storage_root(address_1).unwrap();
1091 sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap();
1092
1093 sparse.wipe_storage(address_2).unwrap();
1094 trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
1095 sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap();
1096
1097 sparse.root().unwrap();
1098
1099 let sparse_updates = sparse.take_trie_updates().unwrap();
1100 pretty_assertions::assert_eq!(
1102 sparse_updates,
1103 TrieUpdates {
1104 account_nodes: HashMap::default(),
1105 storage_tries: HashMap::from_iter([(
1106 b256!("0x1100000000000000000000000000000000000000000000000000000000000000"),
1107 StorageTrieUpdates {
1108 is_deleted: true,
1109 storage_nodes: HashMap::default(),
1110 removed_nodes: HashSet::default()
1111 }
1112 )]),
1113 removed_nodes: HashSet::default()
1114 }
1115 );
1116 }
1117}