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 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::state_trie(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::state_trie(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                    let is_private = false; // account leaves are always public. Their storage leaves can be private.
156                    hash_builder.add_leaf(
157                        Nibbles::unpack(hashed_address),
158                        &account_rlp,
159                        is_private,
160                    );
161
162                    // We might be adding leaves that are not necessarily our proof targets.
163                    if leaf_is_proof_target {
164                        // Overwrite storage multiproof.
165                        storages.insert(hashed_address, storage_multiproof);
166                    }
167                }
168            }
169        }
170        let _ = hash_builder.root();
171        let account_subtree = hash_builder.take_proof_nodes();
172        let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks {
173            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
174            (
175                updated_branch_nodes
176                    .iter()
177                    .map(|(path, node)| (path.clone(), node.hash_mask))
178                    .collect(),
179                updated_branch_nodes
180                    .into_iter()
181                    .map(|(path, node)| (path, node.tree_mask))
182                    .collect(),
183            )
184        } else {
185            (HashMap::default(), HashMap::default())
186        };
187
188        Ok(MultiProof { account_subtree, branch_node_hash_masks, branch_node_tree_masks, storages })
189    }
190}
191
192/// Generates storage merkle proofs.
193#[derive(Debug)]
194pub struct StorageProof<T, H> {
195    /// The factory for traversing trie nodes.
196    trie_cursor_factory: T,
197    /// The factory for hashed cursors.
198    hashed_cursor_factory: H,
199    /// The hashed address of an account.
200    hashed_address: B256,
201    /// The set of storage slot prefixes that have changed.
202    prefix_set: PrefixSetMut,
203    /// Flag indicating whether to include branch node masks in the proof.
204    collect_branch_node_masks: bool,
205}
206
207impl<T, H> StorageProof<T, H> {
208    /// Create a new [`StorageProof`] instance.
209    pub fn new(t: T, h: H, address: Address) -> Self {
210        Self::new_hashed(t, h, keccak256(address))
211    }
212
213    /// Create a new [`StorageProof`] instance with hashed address.
214    pub fn new_hashed(t: T, h: H, hashed_address: B256) -> Self {
215        Self {
216            trie_cursor_factory: t,
217            hashed_cursor_factory: h,
218            hashed_address,
219            prefix_set: PrefixSetMut::default(),
220            collect_branch_node_masks: false,
221        }
222    }
223
224    /// Set the trie cursor factory.
225    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> StorageProof<TF, H> {
226        StorageProof {
227            trie_cursor_factory,
228            hashed_cursor_factory: self.hashed_cursor_factory,
229            hashed_address: self.hashed_address,
230            prefix_set: self.prefix_set,
231            collect_branch_node_masks: self.collect_branch_node_masks,
232        }
233    }
234
235    /// Set the hashed cursor factory.
236    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> StorageProof<T, HF> {
237        StorageProof {
238            trie_cursor_factory: self.trie_cursor_factory,
239            hashed_cursor_factory,
240            hashed_address: self.hashed_address,
241            prefix_set: self.prefix_set,
242            collect_branch_node_masks: self.collect_branch_node_masks,
243        }
244    }
245
246    /// Set the changed prefixes.
247    pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self {
248        self.prefix_set = prefix_set;
249        self
250    }
251
252    /// Set the flag indicating whether to include branch node masks in the proof.
253    pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self {
254        self.collect_branch_node_masks = branch_node_masks;
255        self
256    }
257}
258
259impl<T, H> StorageProof<T, H>
260where
261    T: TrieCursorFactory,
262    H: HashedCursorFactory,
263{
264    /// Generate an account proof from intermediate nodes.
265    pub fn storage_proof(
266        self,
267        slot: B256,
268    ) -> Result<reth_trie_common::StorageProof, StateProofError> {
269        let targets = HashSet::from_iter([keccak256(slot)]);
270        Ok(self.storage_multiproof(targets)?.storage_proof(slot)?)
271    }
272
273    /// Generate storage proof.
274    pub fn storage_multiproof(
275        mut self,
276        targets: B256Set,
277    ) -> Result<StorageMultiProof, StateProofError> {
278        let mut hashed_storage_cursor =
279            self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
280
281        // short circuit on empty storage
282        if hashed_storage_cursor.is_storage_empty()? {
283            return Ok(StorageMultiProof::empty());
284        }
285
286        let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
287        self.prefix_set.extend_keys(target_nibbles.clone());
288
289        let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
290        let walker = TrieWalker::storage_trie(trie_cursor, self.prefix_set.freeze());
291
292        let retainer = ProofRetainer::from_iter(target_nibbles);
293        let mut hash_builder = HashBuilder::default()
294            .with_proof_retainer(retainer)
295            .with_updates(self.collect_branch_node_masks);
296        let mut storage_node_iter = TrieNodeIter::storage_trie(walker, hashed_storage_cursor);
297        while let Some(node) = storage_node_iter.try_next()? {
298            match node {
299                TrieElement::Branch(node) => {
300                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
301                }
302                TrieElement::Leaf(hashed_slot, value) => {
303                    hash_builder.add_leaf(
304                        Nibbles::unpack(hashed_slot),
305                        alloy_rlp::encode_fixed_size(&value.value).as_ref(),
306                        value.is_private,
307                    );
308                }
309            }
310        }
311
312        let root = hash_builder.root();
313        let subtree = hash_builder.take_proof_nodes();
314        let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks {
315            let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
316            (
317                updated_branch_nodes
318                    .iter()
319                    .map(|(path, node)| (path.clone(), node.hash_mask))
320                    .collect(),
321                updated_branch_nodes
322                    .into_iter()
323                    .map(|(path, node)| (path, node.tree_mask))
324                    .collect(),
325            )
326        } else {
327            (HashMap::default(), HashMap::default())
328        };
329
330        Ok(StorageMultiProof { root, subtree, branch_node_hash_masks, branch_node_tree_masks })
331    }
332}
reth_trie/proof/mod.rs

reth_trie/proof/
mod.rs
