reth_trie/proof/
trie_node.rs

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/// Factory for instantiating providers capable of retrieving blinded trie nodes via proofs.
13#[derive(Debug, Clone)]
14pub struct ProofTrieNodeProviderFactory<T, H> {
15    /// The cursor factory for traversing trie nodes.
16    trie_cursor_factory: T,
17    /// The factory for hashed cursors.
18    hashed_cursor_factory: H,
19    /// A set of prefix sets that have changes.
20    prefix_sets: Arc<TriePrefixSetsMut>,
21}
22
23impl<T, H> ProofTrieNodeProviderFactory<T, H> {
24    /// Create new proof-based blinded provider factory.
25    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/// Blinded provider for retrieving account trie nodes by path.
61#[derive(Debug)]
62pub struct ProofBlindedAccountProvider<T, H> {
63    /// The cursor factory for traversing trie nodes.
64    trie_cursor_factory: T,
65    /// The factory for hashed cursors.
66    hashed_cursor_factory: H,
67    /// A set of prefix sets that have changes.
68    prefix_sets: Arc<TriePrefixSetsMut>,
69}
70
71impl<T, H> ProofBlindedAccountProvider<T, H> {
72    /// Create new proof-based blinded account node provider.
73    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/// Blinded provider for retrieving storage trie nodes by path.
114#[derive(Debug)]
115pub struct ProofBlindedStorageProvider<T, H> {
116    /// The cursor factory for traversing trie nodes.
117    trie_cursor_factory: T,
118    /// The factory for hashed cursors.
119    hashed_cursor_factory: H,
120    /// A set of prefix sets that have changes.
121    prefix_sets: Arc<TriePrefixSetsMut>,
122    /// Target account.
123    account: B256,
124}
125
126impl<T, H> ProofBlindedStorageProvider<T, H> {
127    /// Create new proof-based blinded storage node provider.
128    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}