1use super::{Proof, StorageProof};
2use crate::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory};
3use alloy_primitives::{map::HashSet, B256};
4use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind};
5use reth_trie_common::{prefix_set::TriePrefixSetsMut, MultiProofTargets, Nibbles};
6use reth_trie_sparse::provider::{
7 pad_path_to_key, RevealedNode, TrieNodeProvider, TrieNodeProviderFactory,
8};
9use std::{sync::Arc, time::Instant};
10use tracing::{enabled, trace, Level};
11
12#[derive(Debug, Clone)]
14pub struct ProofTrieNodeProviderFactory<T, H> {
15 trie_cursor_factory: T,
17 hashed_cursor_factory: H,
19 prefix_sets: Arc<TriePrefixSetsMut>,
21}
22
23impl<T, H> ProofTrieNodeProviderFactory<T, H> {
24 pub const fn new(
26 trie_cursor_factory: T,
27 hashed_cursor_factory: H,
28 prefix_sets: Arc<TriePrefixSetsMut>,
29 ) -> Self {
30 Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets }
31 }
32}
33
34impl<T, H> TrieNodeProviderFactory for ProofTrieNodeProviderFactory<T, H>
35where
36 T: TrieCursorFactory + Clone + Send + Sync,
37 H: HashedCursorFactory + Clone + Send + Sync,
38{
39 type AccountNodeProvider = ProofBlindedAccountProvider<T, H>;
40 type StorageNodeProvider = ProofBlindedStorageProvider<T, H>;
41
42 fn account_node_provider(&self) -> Self::AccountNodeProvider {
43 ProofBlindedAccountProvider {
44 trie_cursor_factory: self.trie_cursor_factory.clone(),
45 hashed_cursor_factory: self.hashed_cursor_factory.clone(),
46 prefix_sets: self.prefix_sets.clone(),
47 }
48 }
49
50 fn storage_node_provider(&self, account: B256) -> Self::StorageNodeProvider {
51 ProofBlindedStorageProvider {
52 trie_cursor_factory: self.trie_cursor_factory.clone(),
53 hashed_cursor_factory: self.hashed_cursor_factory.clone(),
54 prefix_sets: self.prefix_sets.clone(),
55 account,
56 }
57 }
58}
59
60#[derive(Debug)]
62pub struct ProofBlindedAccountProvider<T, H> {
63 trie_cursor_factory: T,
65 hashed_cursor_factory: H,
67 prefix_sets: Arc<TriePrefixSetsMut>,
69}
70
71impl<T, H> ProofBlindedAccountProvider<T, H> {
72 pub const fn new(
74 trie_cursor_factory: T,
75 hashed_cursor_factory: H,
76 prefix_sets: Arc<TriePrefixSetsMut>,
77 ) -> Self {
78 Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets }
79 }
80}
81
82impl<T, H> TrieNodeProvider for ProofBlindedAccountProvider<T, H>
83where
84 T: TrieCursorFactory,
85 H: HashedCursorFactory,
86{
87 fn trie_node(&self, path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
88 let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now);
89
90 let targets = MultiProofTargets::from_iter([(pad_path_to_key(path), HashSet::default())]);
91 let mut proof = Proof::new(&self.trie_cursor_factory, &self.hashed_cursor_factory)
92 .with_prefix_sets_mut(self.prefix_sets.as_ref().clone())
93 .with_branch_node_masks(true)
94 .multiproof(targets)
95 .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
96 let node = proof.account_subtree.into_inner().remove(path);
97 let tree_mask = proof.branch_node_tree_masks.remove(path);
98 let hash_mask = proof.branch_node_hash_masks.remove(path);
99
100 trace!(
101 target: "trie::proof::blinded",
102 elapsed = ?start.unwrap().elapsed(),
103 ?path,
104 ?node,
105 ?tree_mask,
106 ?hash_mask,
107 "Blinded node for account trie"
108 );
109 Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask }))
110 }
111}
112
113#[derive(Debug)]
115pub struct ProofBlindedStorageProvider<T, H> {
116 trie_cursor_factory: T,
118 hashed_cursor_factory: H,
120 prefix_sets: Arc<TriePrefixSetsMut>,
122 account: B256,
124}
125
126impl<T, H> ProofBlindedStorageProvider<T, H> {
127 pub const fn new(
129 trie_cursor_factory: T,
130 hashed_cursor_factory: H,
131 prefix_sets: Arc<TriePrefixSetsMut>,
132 account: B256,
133 ) -> Self {
134 Self { trie_cursor_factory, hashed_cursor_factory, prefix_sets, account }
135 }
136}
137
138impl<T, H> TrieNodeProvider for ProofBlindedStorageProvider<T, H>
139where
140 T: TrieCursorFactory,
141 H: HashedCursorFactory,
142{
143 fn trie_node(&self, path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
144 let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now);
145
146 let targets = HashSet::from_iter([pad_path_to_key(path)]);
147 let storage_prefix_set =
148 self.prefix_sets.storage_prefix_sets.get(&self.account).cloned().unwrap_or_default();
149 let mut proof = StorageProof::new_hashed(
150 &self.trie_cursor_factory,
151 &self.hashed_cursor_factory,
152 self.account,
153 )
154 .with_prefix_set_mut(storage_prefix_set)
155 .with_branch_node_masks(true)
156 .storage_multiproof(targets)
157 .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
158 let node = proof.subtree.into_inner().remove(path);
159 let tree_mask = proof.branch_node_tree_masks.remove(path);
160 let hash_mask = proof.branch_node_hash_masks.remove(path);
161
162 trace!(
163 target: "trie::proof::blinded",
164 account = ?self.account,
165 elapsed = ?start.unwrap().elapsed(),
166 ?path,
167 ?node,
168 ?tree_mask,
169 ?hash_mask,
170 "Blinded node for storage trie"
171 );
172 Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask }))
173 }
174}