reth_trie/proof/
mod.rs

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