1use crate::{Nibbles, TrieAccount};
4use alloc::{borrow::Cow, vec::Vec};
5use alloy_consensus::constants::KECCAK_EMPTY;
6use alloy_primitives::{
7 keccak256,
8 map::{hash_map, B256Map, B256Set, HashMap},
9 Address, Bytes, B256, U256,
10};
11use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE};
12use alloy_trie::{
13 nodes::TrieNode,
14 proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError},
15 TrieMask, EMPTY_ROOT_HASH,
16};
17use derive_more::{Deref, DerefMut, IntoIterator};
18use itertools::Itertools;
19use reth_primitives_traits::Account;
20
21#[derive(Deref, DerefMut, IntoIterator, Clone, PartialEq, Eq, Default, Debug)]
23pub struct MultiProofTargets(B256Map<B256Set>);
24
25impl FromIterator<(B256, B256Set)> for MultiProofTargets {
26 fn from_iter<T: IntoIterator<Item = (B256, B256Set)>>(iter: T) -> Self {
27 Self(B256Map::from_iter(iter))
28 }
29}
30
31impl MultiProofTargets {
32 pub fn with_capacity(capacity: usize) -> Self {
34 Self(B256Map::with_capacity_and_hasher(capacity, Default::default()))
35 }
36
37 pub fn account(hashed_address: B256) -> Self {
39 Self::accounts([hashed_address])
40 }
41
42 pub fn account_with_slots<I: IntoIterator<Item = B256>>(
44 hashed_address: B256,
45 slots_iter: I,
46 ) -> Self {
47 Self(B256Map::from_iter([(hashed_address, slots_iter.into_iter().collect())]))
48 }
49
50 pub fn accounts<I: IntoIterator<Item = B256>>(iter: I) -> Self {
52 Self(iter.into_iter().map(|hashed_address| (hashed_address, Default::default())).collect())
53 }
54
55 pub fn retain_difference(&mut self, other: &Self) {
58 self.0.retain(|hashed_address, hashed_slots| {
59 if let Some(other_hashed_slots) = other.get(hashed_address) {
60 hashed_slots.retain(|hashed_slot| !other_hashed_slots.contains(hashed_slot));
61 !hashed_slots.is_empty()
62 } else {
63 true
64 }
65 });
66 }
67
68 pub fn extend(&mut self, other: Self) {
70 self.extend_inner(Cow::Owned(other));
71 }
72
73 pub fn extend_ref(&mut self, other: &Self) {
77 self.extend_inner(Cow::Borrowed(other));
78 }
79
80 fn extend_inner(&mut self, other: Cow<'_, Self>) {
81 for (hashed_address, hashed_slots) in other.iter() {
82 self.entry(*hashed_address).or_default().extend(hashed_slots);
83 }
84 }
85
86 pub fn chunks(self, size: usize) -> ChunkedMultiProofTargets {
90 ChunkedMultiProofTargets::new(self, size)
91 }
92}
93
94#[derive(Debug)]
115pub struct ChunkedMultiProofTargets {
116 flattened_targets: alloc::vec::IntoIter<(B256, Option<B256>)>,
117 size: usize,
118}
119
120impl ChunkedMultiProofTargets {
121 fn new(targets: MultiProofTargets, size: usize) -> Self {
122 let flattened_targets = targets
123 .into_iter()
124 .flat_map(|(address, slots)| {
125 if slots.is_empty() {
126 itertools::Either::Left(core::iter::once((address, None)))
130 } else {
131 itertools::Either::Right(
132 slots.into_iter().map(move |slot| (address, Some(slot))),
133 )
134 }
135 })
136 .sorted();
137 Self { flattened_targets, size }
138 }
139}
140
141impl Iterator for ChunkedMultiProofTargets {
142 type Item = MultiProofTargets;
143
144 fn next(&mut self) -> Option<Self::Item> {
145 let chunk = self.flattened_targets.by_ref().take(self.size).fold(
146 MultiProofTargets::default(),
147 |mut acc, (address, slot)| {
148 let entry = acc.entry(address).or_default();
149 if let Some(slot) = slot {
150 entry.insert(slot);
151 }
152 acc
153 },
154 );
155
156 if chunk.is_empty() {
157 None
158 } else {
159 Some(chunk)
160 }
161 }
162}
163
164#[derive(Clone, Default, Debug, PartialEq, Eq)]
168pub struct MultiProof {
169 pub account_subtree: ProofNodes,
171 pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
173 pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
175 pub storages: B256Map<StorageMultiProof>,
177}
178
179impl MultiProof {
180 pub fn is_empty(&self) -> bool {
182 self.account_subtree.is_empty() &&
183 self.branch_node_hash_masks.is_empty() &&
184 self.branch_node_tree_masks.is_empty() &&
185 self.storages.is_empty()
186 }
187
188 pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, Bytes)> {
190 self.account_subtree.matching_nodes_sorted(path)
191 }
192
193 pub fn storage_proof_nodes(
195 &self,
196 hashed_address: B256,
197 slots: impl IntoIterator<Item = B256>,
198 ) -> Vec<(B256, Vec<(Nibbles, Bytes)>)> {
199 self.storages
200 .get(&hashed_address)
201 .map(|storage_mp| {
202 slots
203 .into_iter()
204 .map(|slot| {
205 let nibbles = Nibbles::unpack(slot);
206 (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles))
207 })
208 .collect()
209 })
210 .unwrap_or_default()
211 }
212
213 pub fn account_proof(
215 &self,
216 address: Address,
217 slots: &[B256],
218 ) -> Result<AccountProof, alloy_rlp::Error> {
219 let hashed_address = keccak256(address);
220 let nibbles = Nibbles::unpack(hashed_address);
221
222 let proof = self
224 .account_proof_nodes(&nibbles)
225 .into_iter()
226 .map(|(_, node)| node)
227 .collect::<Vec<_>>();
228
229 let info = 'info: {
232 if let Some(last) = proof.last() {
233 if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &last[..])? {
234 if nibbles.ends_with(&leaf.key) {
235 let account = TrieAccount::decode(&mut &leaf.value[..])?;
236 break 'info Some(Account {
237 balance: account.balance,
238 nonce: account.nonce,
239 bytecode_hash: (account.code_hash != KECCAK_EMPTY)
240 .then_some(account.code_hash),
241 });
242 }
243 }
244 }
245 None
246 };
247
248 let storage_multiproof = self.storages.get(&hashed_address);
250 let storage_root = storage_multiproof.map(|m| m.root).unwrap_or(EMPTY_ROOT_HASH);
251 let mut storage_proofs = Vec::with_capacity(slots.len());
252 for slot in slots {
253 let proof = if let Some(multiproof) = &storage_multiproof {
254 multiproof.storage_proof(*slot)?
255 } else {
256 StorageProof::new(*slot)
257 };
258 storage_proofs.push(proof);
259 }
260 Ok(AccountProof { address, info, proof, storage_root, storage_proofs })
261 }
262
263 pub fn extend(&mut self, other: Self) {
266 self.account_subtree.extend_from(other.account_subtree);
267
268 self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
269 self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
270
271 for (hashed_address, storage) in other.storages {
272 match self.storages.entry(hashed_address) {
273 hash_map::Entry::Occupied(mut entry) => {
274 debug_assert_eq!(entry.get().root, storage.root);
275 let entry = entry.get_mut();
276 entry.subtree.extend_from(storage.subtree);
277 entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks);
278 entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
279 }
280 hash_map::Entry::Vacant(entry) => {
281 entry.insert(storage);
282 }
283 }
284 }
285 }
286
287 pub fn from_storage_proof(hashed_address: B256, storage_proof: StorageMultiProof) -> Self {
289 Self {
290 storages: B256Map::from_iter([(hashed_address, storage_proof)]),
291 ..Default::default()
292 }
293 }
294}
295
296#[derive(Clone, Default, Debug, PartialEq, Eq)]
299pub struct DecodedMultiProof {
300 pub account_subtree: DecodedProofNodes,
302 pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
304 pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
306 pub storages: B256Map<DecodedStorageMultiProof>,
308}
309
310impl DecodedMultiProof {
311 pub fn is_empty(&self) -> bool {
313 self.account_subtree.is_empty() &&
314 self.branch_node_hash_masks.is_empty() &&
315 self.branch_node_tree_masks.is_empty() &&
316 self.storages.is_empty()
317 }
318
319 pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, TrieNode)> {
321 self.account_subtree.matching_nodes_sorted(path)
322 }
323
324 pub fn storage_proof_nodes(
326 &self,
327 hashed_address: B256,
328 slots: impl IntoIterator<Item = B256>,
329 ) -> Vec<(B256, Vec<(Nibbles, TrieNode)>)> {
330 self.storages
331 .get(&hashed_address)
332 .map(|storage_mp| {
333 slots
334 .into_iter()
335 .map(|slot| {
336 let nibbles = Nibbles::unpack(slot);
337 (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles))
338 })
339 .collect()
340 })
341 .unwrap_or_default()
342 }
343
344 pub fn account_proof(
346 &self,
347 address: Address,
348 slots: &[B256],
349 ) -> Result<DecodedAccountProof, alloy_rlp::Error> {
350 let hashed_address = keccak256(address);
351 let nibbles = Nibbles::unpack(hashed_address);
352
353 let proof = self
355 .account_proof_nodes(&nibbles)
356 .into_iter()
357 .map(|(_, node)| node)
358 .collect::<Vec<_>>();
359
360 let info = 'info: {
363 if let Some(TrieNode::Leaf(leaf)) = proof.last() {
364 if nibbles.ends_with(&leaf.key) {
365 let account = TrieAccount::decode(&mut &leaf.value[..])?;
366 break 'info Some(Account {
367 balance: account.balance,
368 nonce: account.nonce,
369 bytecode_hash: (account.code_hash != KECCAK_EMPTY)
370 .then_some(account.code_hash),
371 });
372 }
373 }
374 None
375 };
376
377 let storage_multiproof = self.storages.get(&hashed_address);
379 let storage_root = storage_multiproof.map(|m| m.root).unwrap_or(EMPTY_ROOT_HASH);
380 let mut storage_proofs = Vec::with_capacity(slots.len());
381 for slot in slots {
382 let proof = if let Some(multiproof) = &storage_multiproof {
383 multiproof.storage_proof(*slot)?
384 } else {
385 DecodedStorageProof::new(*slot)
386 };
387 storage_proofs.push(proof);
388 }
389 Ok(DecodedAccountProof { address, info, proof, storage_root, storage_proofs })
390 }
391
392 pub fn extend(&mut self, other: Self) {
395 self.account_subtree.extend_from(other.account_subtree);
396
397 self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
398 self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
399
400 for (hashed_address, storage) in other.storages {
401 match self.storages.entry(hashed_address) {
402 hash_map::Entry::Occupied(mut entry) => {
403 debug_assert_eq!(entry.get().root, storage.root);
404 let entry = entry.get_mut();
405 entry.subtree.extend_from(storage.subtree);
406 entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks);
407 entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
408 }
409 hash_map::Entry::Vacant(entry) => {
410 entry.insert(storage);
411 }
412 }
413 }
414 }
415
416 pub fn from_storage_proof(
418 hashed_address: B256,
419 storage_proof: DecodedStorageMultiProof,
420 ) -> Self {
421 Self {
422 storages: B256Map::from_iter([(hashed_address, storage_proof)]),
423 ..Default::default()
424 }
425 }
426}
427
428impl TryFrom<MultiProof> for DecodedMultiProof {
429 type Error = alloy_rlp::Error;
430
431 fn try_from(multi_proof: MultiProof) -> Result<Self, Self::Error> {
432 let account_subtree = DecodedProofNodes::try_from(multi_proof.account_subtree)?;
433 let storages = multi_proof
434 .storages
435 .into_iter()
436 .map(|(address, storage)| Ok((address, storage.try_into()?)))
437 .collect::<Result<B256Map<_>, alloy_rlp::Error>>()?;
438 Ok(Self {
439 account_subtree,
440 branch_node_hash_masks: multi_proof.branch_node_hash_masks,
441 branch_node_tree_masks: multi_proof.branch_node_tree_masks,
442 storages,
443 })
444 }
445}
446
447#[derive(Clone, Debug, PartialEq, Eq)]
449pub struct StorageMultiProof {
450 pub root: B256,
452 pub subtree: ProofNodes,
454 pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
456 pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
458}
459
460impl StorageMultiProof {
461 pub fn empty() -> Self {
463 Self {
464 root: EMPTY_ROOT_HASH,
465 subtree: ProofNodes::from_iter([(
466 Nibbles::default(),
467 Bytes::from([EMPTY_STRING_CODE]),
468 )]),
469 branch_node_hash_masks: HashMap::default(),
470 branch_node_tree_masks: HashMap::default(),
471 }
472 }
473
474 pub fn storage_proof(&self, slot: B256) -> Result<StorageProof, alloy_rlp::Error> {
476 let nibbles = Nibbles::unpack(keccak256(slot));
477
478 let proof = self
480 .subtree
481 .matching_nodes_iter(&nibbles)
482 .sorted_by(|a, b| a.0.cmp(b.0))
483 .map(|(_, node)| node.clone())
484 .collect::<Vec<_>>();
485
486 let mut is_private = false;
489 let value = 'value: {
490 if let Some(last) = proof.last() {
491 if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &last[..])? {
492 if nibbles.ends_with(&leaf.key) {
493 is_private = leaf.is_private;
494 break 'value U256::decode(&mut &leaf.value[..])?;
495 }
496 }
497 }
498 U256::ZERO
499 };
500
501 Ok(StorageProof { key: slot, nibbles, value, is_private, proof })
502 }
503}
504
505#[derive(Clone, Debug, PartialEq, Eq)]
507pub struct DecodedStorageMultiProof {
508 pub root: B256,
510 pub subtree: DecodedProofNodes,
512 pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
514 pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
516}
517
518impl DecodedStorageMultiProof {
519 pub fn empty() -> Self {
521 Self {
522 root: EMPTY_ROOT_HASH,
523 subtree: DecodedProofNodes::from_iter([(Nibbles::default(), TrieNode::EmptyRoot)]),
524 branch_node_hash_masks: HashMap::default(),
525 branch_node_tree_masks: HashMap::default(),
526 }
527 }
528
529 pub fn storage_proof(&self, slot: B256) -> Result<DecodedStorageProof, alloy_rlp::Error> {
531 let nibbles = Nibbles::unpack(keccak256(slot));
532
533 let proof = self
535 .subtree
536 .matching_nodes_iter(&nibbles)
537 .sorted_by(|a, b| a.0.cmp(b.0))
538 .map(|(_, node)| node.clone())
539 .collect::<Vec<_>>();
540
541 let value = 'value: {
544 if let Some(TrieNode::Leaf(leaf)) = proof.last() {
545 if nibbles.ends_with(&leaf.key) {
546 break 'value U256::decode(&mut &leaf.value[..])?;
547 }
548 }
549 U256::ZERO
550 };
551
552 Ok(DecodedStorageProof { key: slot, nibbles, value, proof })
553 }
554}
555
556impl TryFrom<StorageMultiProof> for DecodedStorageMultiProof {
557 type Error = alloy_rlp::Error;
558
559 fn try_from(multi_proof: StorageMultiProof) -> Result<Self, Self::Error> {
560 let subtree = DecodedProofNodes::try_from(multi_proof.subtree)?;
561 Ok(Self {
562 root: multi_proof.root,
563 subtree,
564 branch_node_hash_masks: multi_proof.branch_node_hash_masks,
565 branch_node_tree_masks: multi_proof.branch_node_tree_masks,
566 })
567 }
568}
569
570#[derive(Clone, PartialEq, Eq, Debug)]
572#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
573#[cfg_attr(any(test, feature = "serde"), serde(rename_all = "camelCase"))]
574pub struct AccountProof {
575 pub address: Address,
577 pub info: Option<Account>,
579 pub proof: Vec<Bytes>,
582 pub storage_root: B256,
584 pub storage_proofs: Vec<StorageProof>,
586}
587
588#[cfg(feature = "eip1186")]
589impl AccountProof {
590 pub fn into_eip1186_response(
592 self,
593 slots: Vec<alloy_serde::JsonStorageKey>,
594 ) -> alloy_rpc_types_eth::EIP1186AccountProofResponse {
595 let info = self.info.unwrap_or_default();
596 alloy_rpc_types_eth::EIP1186AccountProofResponse {
597 address: self.address,
598 balance: info.balance,
599 code_hash: info.get_bytecode_hash(),
600 nonce: info.nonce,
601 storage_hash: self.storage_root,
602 account_proof: self.proof,
603 storage_proof: self
604 .storage_proofs
605 .into_iter()
606 .filter_map(|proof| {
607 let input_slot = slots.iter().find(|s| s.as_b256() == proof.key)?;
608 Some(proof.into_eip1186_proof(*input_slot))
609 })
610 .collect(),
611 }
612 }
613
614 pub fn from_eip1186_proof(proof: alloy_rpc_types_eth::EIP1186AccountProofResponse) -> Self {
620 let alloy_rpc_types_eth::EIP1186AccountProofResponse {
621 nonce,
622 address,
623 balance,
624 code_hash,
625 storage_hash,
626 account_proof,
627 storage_proof,
628 ..
629 } = proof;
630 let storage_proofs = storage_proof.into_iter().map(Into::into).collect();
631
632 let (storage_root, info) = if nonce == 0 &&
633 balance.is_zero() &&
634 storage_hash.is_zero() &&
635 code_hash == KECCAK_EMPTY
636 {
637 (EMPTY_ROOT_HASH, None)
640 } else {
641 (storage_hash, Some(Account { nonce, balance, bytecode_hash: code_hash.into() }))
642 };
643
644 Self { address, info, proof: account_proof, storage_root, storage_proofs }
645 }
646}
647
648#[cfg(feature = "eip1186")]
649impl From<alloy_rpc_types_eth::EIP1186AccountProofResponse> for AccountProof {
650 fn from(proof: alloy_rpc_types_eth::EIP1186AccountProofResponse) -> Self {
651 Self::from_eip1186_proof(proof)
652 }
653}
654
655impl Default for AccountProof {
656 fn default() -> Self {
657 Self::new(Address::default())
658 }
659}
660
661impl AccountProof {
662 pub const fn new(address: Address) -> Self {
664 Self {
665 address,
666 info: None,
667 proof: Vec::new(),
668 storage_root: EMPTY_ROOT_HASH,
669 storage_proofs: Vec::new(),
670 }
671 }
672
673 #[expect(clippy::result_large_err)]
675 pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
676 for storage_proof in &self.storage_proofs {
678 storage_proof.verify(self.storage_root)?;
679 }
680
681 let expected = if self.info.is_none() && self.storage_root == EMPTY_ROOT_HASH {
683 None
684 } else {
685 Some(alloy_rlp::encode(
686 self.info.unwrap_or_default().into_trie_account(self.storage_root),
687 ))
688 };
689 let nibbles = Nibbles::unpack(keccak256(self.address));
690 let account_node_is_private = false; verify_proof(root, nibbles, expected, account_node_is_private, &self.proof)
692 }
693}
694
695#[derive(Clone, PartialEq, Eq, Debug)]
697pub struct DecodedAccountProof {
698 pub address: Address,
700 pub info: Option<Account>,
702 pub proof: Vec<TrieNode>,
705 pub storage_root: B256,
707 pub storage_proofs: Vec<DecodedStorageProof>,
709}
710
711impl Default for DecodedAccountProof {
712 fn default() -> Self {
713 Self::new(Address::default())
714 }
715}
716
717impl DecodedAccountProof {
718 pub const fn new(address: Address) -> Self {
720 Self {
721 address,
722 info: None,
723 proof: Vec::new(),
724 storage_root: EMPTY_ROOT_HASH,
725 storage_proofs: Vec::new(),
726 }
727 }
728}
729
730#[derive(Clone, PartialEq, Eq, Default, Debug)]
732#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
733pub struct StorageProof {
734 pub key: B256,
736 pub nibbles: Nibbles,
738 pub value: U256,
740 pub is_private: bool,
742 pub proof: Vec<Bytes>,
745}
746
747impl StorageProof {
748 pub fn new(key: B256) -> Self {
750 let nibbles = Nibbles::unpack(keccak256(key));
751 Self { key, nibbles, ..Default::default() }
752 }
753
754 pub fn new_with_hashed(key: B256, hashed_key: B256) -> Self {
756 Self { key, nibbles: Nibbles::unpack(hashed_key), ..Default::default() }
757 }
758
759 pub fn new_with_nibbles(key: B256, nibbles: Nibbles) -> Self {
761 Self { key, nibbles, ..Default::default() }
762 }
763
764 pub fn with_proof(mut self, proof: Vec<Bytes>) -> Self {
766 self.proof = proof;
767 self
768 }
769
770 #[expect(clippy::result_large_err)]
772 pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
773 let expected =
774 if self.value.is_zero() { None } else { Some(encode_fixed_size(&self.value).to_vec()) };
775 verify_proof(root, self.nibbles.clone(), expected, self.is_private, &self.proof)
776 }
777}
778
779#[cfg(feature = "eip1186")]
780impl StorageProof {
781 pub fn into_eip1186_proof(
783 self,
784 slot: alloy_serde::JsonStorageKey,
785 ) -> alloy_rpc_types_eth::EIP1186StorageProof {
786 alloy_rpc_types_eth::EIP1186StorageProof { key: slot, value: self.value, proof: self.proof }
787 }
788
789 pub fn from_eip1186_proof(storage_proof: alloy_rpc_types_eth::EIP1186StorageProof) -> Self {
794 Self {
795 value: storage_proof.value,
796 proof: storage_proof.proof,
797 ..Self::new(storage_proof.key.as_b256())
798 }
799 }
800}
801
802#[cfg(feature = "eip1186")]
803impl From<alloy_rpc_types_eth::EIP1186StorageProof> for StorageProof {
804 fn from(proof: alloy_rpc_types_eth::EIP1186StorageProof) -> Self {
805 Self::from_eip1186_proof(proof)
806 }
807}
808
809#[derive(Clone, PartialEq, Eq, Default, Debug)]
811pub struct DecodedStorageProof {
812 pub key: B256,
814 pub nibbles: Nibbles,
816 pub value: U256,
818 pub proof: Vec<TrieNode>,
821}
822
823impl DecodedStorageProof {
824 pub fn new(key: B256) -> Self {
826 let nibbles = Nibbles::unpack(keccak256(key));
827 Self { key, nibbles, ..Default::default() }
828 }
829
830 pub fn new_with_hashed(key: B256, hashed_key: B256) -> Self {
832 Self { key, nibbles: Nibbles::unpack(hashed_key), ..Default::default() }
833 }
834
835 pub fn new_with_nibbles(key: B256, nibbles: Nibbles) -> Self {
837 Self { key, nibbles, ..Default::default() }
838 }
839
840 pub fn with_proof(mut self, proof: Vec<TrieNode>) -> Self {
842 self.proof = proof;
843 self
844 }
845}
846
847#[cfg(any(test, feature = "test-utils"))]
850pub mod triehash {
851 use alloy_primitives::{keccak256, B256};
852 use alloy_rlp::RlpEncodable;
853 use hash_db::Hasher;
854 use plain_hasher::PlainHasher;
855
856 #[derive(Default, Debug, Clone, PartialEq, Eq, RlpEncodable)]
858 #[non_exhaustive]
859 pub struct KeccakHasher;
860
861 #[cfg(any(test, feature = "test-utils"))]
862 impl Hasher for KeccakHasher {
863 type Out = B256;
864 type StdHasher = PlainHasher;
865
866 const LENGTH: usize = 32;
867
868 fn hash(x: &[u8]) -> Self::Out {
869 keccak256(x)
870 }
871 }
872}
873
874#[cfg(test)]
875mod tests {
876 use super::*;
877
878 #[test]
879 fn test_multiproof_extend_account_proofs() {
880 let mut proof1 = MultiProof::default();
881 let mut proof2 = MultiProof::default();
882
883 let addr1 = B256::random();
884 let addr2 = B256::random();
885
886 proof1.account_subtree.insert(
887 Nibbles::unpack(addr1),
888 alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec().into(),
889 );
890 proof2.account_subtree.insert(
891 Nibbles::unpack(addr2),
892 alloy_rlp::encode_fixed_size(&U256::from(43)).to_vec().into(),
893 );
894
895 proof1.extend(proof2);
896
897 assert!(proof1.account_subtree.contains_key(&Nibbles::unpack(addr1)));
898 assert!(proof1.account_subtree.contains_key(&Nibbles::unpack(addr2)));
899 }
900
901 #[test]
902 fn test_multiproof_extend_storage_proofs() {
903 let mut proof1 = MultiProof::default();
904 let mut proof2 = MultiProof::default();
905
906 let addr = B256::random();
907 let root = B256::random();
908
909 let mut subtree1 = ProofNodes::default();
910 subtree1.insert(
911 Nibbles::from_nibbles(vec![0]),
912 alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec().into(),
913 );
914 proof1.storages.insert(
915 addr,
916 StorageMultiProof {
917 root,
918 subtree: subtree1,
919 branch_node_hash_masks: HashMap::default(),
920 branch_node_tree_masks: HashMap::default(),
921 },
922 );
923
924 let mut subtree2 = ProofNodes::default();
925 subtree2.insert(
926 Nibbles::from_nibbles(vec![1]),
927 alloy_rlp::encode_fixed_size(&U256::from(43)).to_vec().into(),
928 );
929 proof2.storages.insert(
930 addr,
931 StorageMultiProof {
932 root,
933 subtree: subtree2,
934 branch_node_hash_masks: HashMap::default(),
935 branch_node_tree_masks: HashMap::default(),
936 },
937 );
938
939 proof1.extend(proof2);
940
941 let storage = proof1.storages.get(&addr).unwrap();
942 assert_eq!(storage.root, root);
943 assert!(storage.subtree.contains_key(&Nibbles::from_nibbles(vec![0])));
944 assert!(storage.subtree.contains_key(&Nibbles::from_nibbles(vec![1])));
945 }
946
947 #[test]
948 fn test_multi_proof_retain_difference() {
949 let mut empty = MultiProofTargets::default();
950 empty.retain_difference(&Default::default());
951 assert!(empty.is_empty());
952
953 let targets = MultiProofTargets::accounts((0..10).map(B256::with_last_byte));
954
955 let mut diffed = targets.clone();
956 diffed.retain_difference(&MultiProofTargets::account(B256::with_last_byte(11)));
957 assert_eq!(diffed, targets);
958
959 diffed.retain_difference(&MultiProofTargets::accounts((0..5).map(B256::with_last_byte)));
960 assert_eq!(diffed, MultiProofTargets::accounts((5..10).map(B256::with_last_byte)));
961
962 diffed.retain_difference(&targets);
963 assert!(diffed.is_empty());
964
965 let mut targets = MultiProofTargets::default();
966 let (account1, account2, account3) =
967 (1..=3).map(B256::with_last_byte).collect_tuple().unwrap();
968 let account2_slots = (1..5).map(B256::with_last_byte).collect::<B256Set>();
969 targets.insert(account1, B256Set::from_iter([B256::with_last_byte(1)]));
970 targets.insert(account2, account2_slots.clone());
971 targets.insert(account3, B256Set::from_iter([B256::with_last_byte(1)]));
972
973 let mut diffed = targets.clone();
974 diffed.retain_difference(&MultiProofTargets::accounts((1..=3).map(B256::with_last_byte)));
975 assert_eq!(diffed, targets);
976
977 let mut account2_slots_expected_len = account2_slots.len();
979 for slot in account2_slots.iter().skip(1) {
980 diffed.retain_difference(&MultiProofTargets::account_with_slots(account2, [*slot]));
981 account2_slots_expected_len -= 1;
982 assert_eq!(
983 diffed.get(&account2).map(|slots| slots.len()),
984 Some(account2_slots_expected_len)
985 );
986 }
987
988 diffed.retain_difference(&targets);
989 assert!(diffed.is_empty());
990 }
991
992 #[test]
993 fn test_multi_proof_retain_difference_no_overlap() {
994 let mut targets = MultiProofTargets::default();
995
996 let (addr1, addr2) = (B256::random(), B256::random());
998 let (slot1, slot2) = (B256::random(), B256::random());
999 targets.insert(addr1, vec![slot1].into_iter().collect());
1000 targets.insert(addr2, vec![slot2].into_iter().collect());
1001
1002 let mut retained = targets.clone();
1003 retained.retain_difference(&Default::default());
1004 assert_eq!(retained, targets);
1005
1006 let mut other_targets = MultiProofTargets::default();
1008 let addr3 = B256::random();
1009 let slot3 = B256::random();
1010 other_targets.insert(addr3, B256Set::from_iter([slot3]));
1011
1012 let mut retained = targets.clone();
1015 retained.retain_difference(&other_targets);
1016 assert_eq!(retained, targets);
1017 }
1018
1019 #[test]
1020 fn test_get_prefetch_proof_targets_remove_subset() {
1021 let mut targets = MultiProofTargets::default();
1023 let (addr1, addr2) = (B256::random(), B256::random());
1024 let (slot1, slot2) = (B256::random(), B256::random());
1025 targets.insert(addr1, B256Set::from_iter([slot1]));
1026 targets.insert(addr2, B256Set::from_iter([slot2]));
1027
1028 let other_targets = MultiProofTargets::account_with_slots(addr1, [slot1]);
1030
1031 let mut retained = targets.clone();
1032 retained.retain_difference(&other_targets);
1033
1034 assert_eq!(retained.len(), 1);
1036 assert!(!retained.contains_key(&addr1));
1037 assert!(retained.contains_key(&addr2));
1038
1039 let slot3 = B256::random();
1041 targets.get_mut(&addr1).unwrap().insert(slot3);
1042
1043 let mut retained = targets.clone();
1044 retained.retain_difference(&other_targets);
1045
1046 assert_eq!(retained.len(), 2);
1049 assert!(retained.contains_key(&addr1));
1050 assert_eq!(retained.get(&addr1), Some(&B256Set::from_iter([slot3])));
1051 assert!(retained.contains_key(&addr2));
1052 assert_eq!(retained.get(&addr2), Some(&B256Set::from_iter([slot2])));
1053 }
1054
1055 #[test]
1056 #[cfg(feature = "eip1186")]
1057 fn eip_1186_roundtrip() {
1058 let mut acc = AccountProof {
1059 address: Address::random(),
1060 info: Some(
1061 Account { nonce: 100, balance: U256::ZERO, bytecode_hash: Some(KECCAK_EMPTY) },
1063 ),
1064 proof: vec![],
1065 storage_root: B256::ZERO,
1066 storage_proofs: vec![],
1067 };
1068
1069 let rpc_proof = acc.clone().into_eip1186_response(Vec::new());
1070 let inverse: AccountProof = rpc_proof.into();
1071 assert_eq!(acc, inverse);
1072
1073 acc.info.as_mut().unwrap().nonce = 0;
1075 let rpc_proof = acc.clone().into_eip1186_response(Vec::new());
1076 let inverse: AccountProof = rpc_proof.into();
1077 acc.info.take();
1078 acc.storage_root = EMPTY_ROOT_HASH;
1079 assert_eq!(acc, inverse);
1080 }
1081}