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#[derive(Debug)]
30pub struct Proof<T, H> {
31 trie_cursor_factory: T,
33 hashed_cursor_factory: H,
35 prefix_sets: TriePrefixSetsMut,
37 collect_branch_node_masks: bool,
39}
40
41impl<T, H> Proof<T, H> {
42 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 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 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 pub fn with_prefix_sets_mut(mut self, prefix_sets: TriePrefixSetsMut) -> Self {
74 self.prefix_sets = prefix_sets;
75 self
76 }
77
78 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 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 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 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 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 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 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 if leaf_is_proof_target {
160 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#[derive(Debug)]
187pub struct StorageProof<T, H, K = AddedRemovedKeys> {
188 trie_cursor_factory: T,
190 hashed_cursor_factory: H,
192 hashed_address: B256,
194 prefix_set: PrefixSetMut,
196 collect_branch_node_masks: bool,
198 added_removed_keys: Option<K>,
200}
201
202impl<T, H> StorageProof<T, H> {
203 pub fn new(t: T, h: H, address: Address) -> Self {
205 Self::new_hashed(t, h, keccak256(address))
206 }
207
208 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 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 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 pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self {
251 self.prefix_set = prefix_set;
252 self
253 }
254
255 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 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 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 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 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}