reth_trie/proof/
mod.rs

Help
1use crate::{
2    hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
3    node_iter::{TrieElement, TrieNodeIter},
4    prefix_set::{PrefixSetMut, TriePrefixSetsMut},
5    trie_cursor::TrieCursorFactory,
6    walker::TrieWalker,
7    HashBuilder, Nibbles, TRIE_ACCOUNT_RLP_MAX_SIZE,
8};
9use alloy_primitives::{
10    keccak256,
11    map::{B256Map, B256Set, HashMap, HashSet},
12    Address, B256,
13};
14use alloy_rlp::{BufMut, Encodable};
15use reth_execution_errors::trie::StateProofError;
16use reth_trie_common::{
17    proof::ProofRetainer, AccountProof, MultiProof, MultiProofTargets, StorageMultiProof,
18};
19
20mod blinded;
21pub use blinded::*;
22
23/// A struct for generating merkle proofs.
24///
25/// Proof generator adds the target address and slots to the prefix set, enables the proof retainer
26/// on the hash builder and follows the same algorithm as the state root calculator.
27/// See `StateRoot::root` for more info.
28#[derive(Debug)]
29pub struct Proof<T, H> {
30    /// The factory for traversing trie nodes.
31    trie_cursor_factory: T,
32    /// The factory for hashed cursors.
33    hashed_cursor_factory: H,
34    /// A set of prefix sets that have changes.
35    prefix_sets: TriePrefixSetsMut,
36    /// Flag indicating whether to include branch node masks in the proof.
37    collect_branch_node_masks: bool,
38}
39
40impl<T, H> Proof<T, H> {
41    /// Create a new [`Proof`] instance.
42    pub fn new(t: T, h: H) -> Self {
43        Self {
44            trie_cursor_factory: t,
45            hashed_cursor_factory: h,
46            prefix_sets: TriePrefixSetsMut::default(),
47            collect_branch_node_masks: false,
48        }
49    }
50
51    /// Set the trie cursor factory.
52    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> Proof<TF, H> {
53        Proof {
54            trie_cursor_factory,
55            hashed_cursor_factory: self.hashed_cursor_factory,
56            prefix_sets: self.prefix_sets,
57            collect_branch_node_masks: self.collect_branch_node_masks,
58        }
59    }
60
61    /// Set the hashed cursor factory.
62    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> Proof<T, HF> {
63        Proof {
64            trie_cursor_factory: self.trie_cursor_factory,
65            hashed_cursor_factory,
66            prefix_sets: self.prefix_sets,
67            collect_branch_node_masks: self.collect_branch_node_masks,
68        }
69    }
70
71    /// Set the prefix sets. They have to be mutable in order to allow extension with proof target.
72    pub fn with_prefix_sets_mut(mut self, prefix_sets: TriePrefixSetsMut) -> Self {
73        self.prefix_sets = prefix_sets;
74        self
75    }
76
77    /// Set the flag indicating whether to include branch node masks in the proof.
78    pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self {
79        self.collect_branch_node_masks = branch_node_masks;
80        self
81    }
82}
83
84impl<T, H> Proof<T, H>
85where
86    T: TrieCursorFactory + Clone,
87    H: HashedCursorFactory + Clone,
88{
89    /// Generate an account proof from intermediate nodes.
90    pub fn account_proof(
91        self,
92        address: Address,
93        slots: &[B256],
94    ) -> Result<AccountProof, StateProofError> {
95        Ok(self
96            .multiproof(MultiProofTargets::from_iter([(
97                keccak256(address),
98                slots.iter().map(keccak256).collect(),
99            )]))?
100            .account_proof(address, slots)?)
101    }
102
103    /// Generate a state multiproof according to specified targets.
104    pub fn multiproof(
105        mut self,
106        mut targets: MultiProofTargets,
107    ) -> Result<MultiProof, StateProofError> {
108        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
109        let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
110
111        // Create the walker.
112        let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
113        prefix_set.extend_keys(targets.keys().map(Nibbles::unpack));
114        let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
115
116        // Create a hash builder to rebuild the root node since it is not available in the database.
117        let retainer = targets.keys().map(Nibbles::unpack).collect();
118        let mut hash_builder = HashBuilder::default()
119            .with_proof_retainer(retainer)
120            .with_updates(self.collect_branch_node_masks);
121
122        // Initialize all storage multiproofs as empty.
123        // Storage multiproofs for non empty tries will be overwritten if necessary.
124        let mut storages: B256Map<_> =
125            targets.keys().map(|key| (*key, StorageMultiProof::empty())).collect();
126        let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
127        let mut account_node_iter = TrieNodeIter::new(walker, hashed_account_cursor);
128        while let Some(account_node) = account_node_iter.try_next()? {
129            match account_node {
130                TrieElement::Branch(node) => {
131                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
132                }
133                TrieElement::Leaf(hashed_address, account) => {
134                    let proof_targets = targets.remove(&hashed_address);
135                    let leaf_is_proof_target = proof_targets.is_some();
136                    let storage_prefix_set = self
137                        .prefix_sets
138                        .storage_prefix_sets
139                        .remove(&hashed_address)
140                        .unwrap_or_default();
141                    let storage_multiproof = StorageProof::new_hashed(
142                        self.trie_cursor_factory.clone(),
143                        self.hashed_cursor_factory.clone(),
144                        hashed_address,
145                    )
146                    .with_prefix_set_mut(storage_prefix_set)
147                    .with_branch_node_masks(self.collect_branch_node_masks)
148                    .storage_multiproof(proof_targets.unwrap_or_default())?;
149
150                    // Encode account
151                    account_rlp.clear();
152                    let account = account.into_trie_account(storage_multiproof.root);
153                    account.encode(&mut account_rlp as &mut dyn BufMut);
154
155                    hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
156
157                    // We might be adding leaves that are not necessarily our proof targets.
158                    if leaf_is_proof_target {
159                        // Overwrite storage multiproof.
160                        storages.insert(hashed_address, storage_multiproof);
161                    }
162                }
163            }
164        }
165        let _ = hash_builder.root();
166        let account_subtree = hash_builder.take_proof_nodes();
167        let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks {
168            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
169            (
170                updated_branch_nodes
171                    .iter()
172                    .map(|(path, node)| (path.clone(), node.hash_mask))
173                    .collect(),
174                updated_branch_nodes
175                    .into_iter()
176                    .map(|(path, node)| (path, node.tree_mask))
177                    .collect(),
178            )
179        } else {
180            (HashMap::default(), HashMap::default())
181        };
182
183        Ok(MultiProof { account_subtree, branch_node_hash_masks, branch_node_tree_masks, storages })
184    }
185}
186
187/// Generates storage merkle proofs.
188#[derive(Debug)]
189pub struct StorageProof<T, H> {
190    /// The factory for traversing trie nodes.
191    trie_cursor_factory: T,
192    /// The factory for hashed cursors.
193    hashed_cursor_factory: H,
194    /// The hashed address of an account.
195    hashed_address: B256,
196    /// The set of storage slot prefixes that have changed.
197    prefix_set: PrefixSetMut,
198    /// Flag indicating whether to include branch node masks in the proof.
199    collect_branch_node_masks: bool,
200}
201
202impl<T, H> StorageProof<T, H> {
203    /// Create a new [`StorageProof`] instance.
204    pub fn new(t: T, h: H, address: Address) -> Self {
205        Self::new_hashed(t, h, keccak256(address))
206    }
207
208    /// Create a new [`StorageProof`] instance with hashed address.
209    pub fn new_hashed(t: T, h: H, hashed_address: B256) -> Self {
210        Self {
211            trie_cursor_factory: t,
212            hashed_cursor_factory: h,
213            hashed_address,
214            prefix_set: PrefixSetMut::default(),
215            collect_branch_node_masks: false,
216        }
217    }
218
219    /// Set the trie cursor factory.
220    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageProof<TF, H> {
221        StorageProof {
222            trie_cursor_factory,
223            hashed_cursor_factory: self.hashed_cursor_factory,
224            hashed_address: self.hashed_address,
225            prefix_set: self.prefix_set,
226            collect_branch_node_masks: self.collect_branch_node_masks,
227        }
228    }
229
230    /// Set the hashed cursor factory.
231    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageProof<T, HF> {
232        StorageProof {
233            trie_cursor_factory: self.trie_cursor_factory,
234            hashed_cursor_factory,
235            hashed_address: self.hashed_address,
236            prefix_set: self.prefix_set,
237            collect_branch_node_masks: self.collect_branch_node_masks,
238        }
239    }
240
241    /// Set the changed prefixes.
242    pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self {
243        self.prefix_set = prefix_set;
244        self
245    }
246
247    /// Set the flag indicating whether to include branch node masks in the proof.
248    pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self {
249        self.collect_branch_node_masks = branch_node_masks;
250        self
251    }
252}
253
254impl<T, H> StorageProof<T, H>
255where
256    T: TrieCursorFactory,
257    H: HashedCursorFactory,
258{
259    /// Generate an account proof from intermediate nodes.
260    pub fn storage_proof(
261        self,
262        slot: B256,
263    ) -> Result<reth_trie_common::StorageProof, StateProofError> {
264        let targets = HashSet::from_iter([keccak256(slot)]);
265        Ok(self.storage_multiproof(targets)?.storage_proof(slot)?)
266    }
267
268    /// Generate storage proof.
269    pub fn storage_multiproof(
270        mut self,
271        targets: B256Set,
272    ) -> Result<StorageMultiProof, StateProofError> {
273        let mut hashed_storage_cursor =
274            self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
275
276        // short circuit on empty storage
277        if hashed_storage_cursor.is_storage_empty()? {
278            return Ok(StorageMultiProof::empty())
279        }
280
281        let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
282        self.prefix_set.extend_keys(target_nibbles.clone());
283
284        let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
285        let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze());
286
287        let retainer = ProofRetainer::from_iter(target_nibbles);
288        let mut hash_builder = HashBuilder::default()
289            .with_proof_retainer(retainer)
290            .with_updates(self.collect_branch_node_masks);
291        let mut storage_node_iter = TrieNodeIter::new(walker, hashed_storage_cursor);
292        while let Some(node) = storage_node_iter.try_next()? {
293            match node {
294                TrieElement::Branch(node) => {
295                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
296                }
297                TrieElement::Leaf(hashed_slot, value) => {
298                    hash_builder.add_leaf(
299                        Nibbles::unpack(hashed_slot),
300                        alloy_rlp::encode_fixed_size(&value).as_ref(),
301                    );
302                }
303            }
304        }
305
306        let root = hash_builder.root();
307        let subtree = hash_builder.take_proof_nodes();
308        let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks {
309            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
310            (
311                updated_branch_nodes
312                    .iter()
313                    .map(|(path, node)| (path.clone(), node.hash_mask))
314                    .collect(),
315                updated_branch_nodes
316                    .into_iter()
317                    .map(|(path, node)| (path, node.tree_mask))
318                    .collect(),
319            )
320        } else {
321            (HashMap::default(), HashMap::default())
322        };
323
324        Ok(StorageMultiProof { root, subtree, branch_node_hash_masks, branch_node_tree_masks })
325    }
326}
reth_trie/proof/mod.rs

reth_trie/proof/
mod.rs
