1use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory, TrieTableAdapter};
2use alloy_primitives::{keccak256, map::HashMap, Address, B256};
3use reth_db_api::transaction::DbTx;
4use reth_execution_errors::StateProofError;
5use reth_trie::{
6 hashed_cursor::HashedPostStateCursorFactory,
7 proof::{Proof, StorageProof},
8 trie_cursor::InMemoryTrieCursorFactory,
9 updates::TrieUpdatesSorted,
10 AccountProof, HashedPostStateSorted, HashedStorage, MultiProof, MultiProofTargets,
11 StorageMultiProof, TrieInput, TrieInputSorted,
12};
13
14pub trait DatabaseProof<'a> {
16 type Tx;
18
19 fn from_tx(tx: &'a Self::Tx) -> Self;
21
22 fn overlay_account_proof(
24 &self,
25 input: TrieInput,
26 address: Address,
27 slots: &[B256],
28 ) -> Result<AccountProof, StateProofError>;
29
30 fn overlay_multiproof(
32 &self,
33 input: TrieInput,
34 targets: MultiProofTargets,
35 ) -> Result<MultiProof, StateProofError>;
36}
37
38impl<'a, TX: DbTx, A: TrieTableAdapter> DatabaseProof<'a>
39 for Proof<DatabaseTrieCursorFactory<&'a TX, A>, DatabaseHashedCursorFactory<&'a TX>>
40{
41 type Tx = TX;
42
43 fn from_tx(tx: &'a Self::Tx) -> Self {
44 Self::new(DatabaseTrieCursorFactory::new(tx), DatabaseHashedCursorFactory::new(tx))
45 }
46 fn overlay_account_proof(
47 &self,
48 input: TrieInput,
49 address: Address,
50 slots: &[B256],
51 ) -> Result<AccountProof, StateProofError> {
52 with_overlay(self, input, |proof| proof.account_proof(address, slots))
53 }
54
55 fn overlay_multiproof(
56 &self,
57 input: TrieInput,
58 targets: MultiProofTargets,
59 ) -> Result<MultiProof, StateProofError> {
60 with_overlay(self, input, |proof| proof.multiproof(targets))
61 }
62}
63
64pub trait DatabaseStorageProof<'a, TX> {
66 fn from_tx(tx: &'a TX, address: Address) -> Self;
68
69 fn overlay_storage_proof(
71 tx: &'a TX,
72 address: Address,
73 slot: B256,
74 storage: HashedStorage,
75 ) -> Result<reth_trie::StorageProof, StateProofError>;
76
77 fn overlay_storage_multiproof(
79 tx: &'a TX,
80 address: Address,
81 slots: &[B256],
82 input: TrieInputSorted,
83 ) -> Result<StorageMultiProof, StateProofError>;
84}
85
86impl<'a, TX: DbTx, A: TrieTableAdapter> DatabaseStorageProof<'a, TX>
87 for StorageProof<DatabaseTrieCursorFactory<&'a TX, A>, DatabaseHashedCursorFactory<&'a TX>>
88{
89 fn from_tx(tx: &'a TX, address: Address) -> Self {
90 Self::new(
91 DatabaseTrieCursorFactory::<_, A>::new(tx),
92 DatabaseHashedCursorFactory::new(tx),
93 address,
94 )
95 }
96
97 fn overlay_storage_proof(
98 tx: &'a TX,
99 address: Address,
100 slot: B256,
101 storage: HashedStorage,
102 ) -> Result<reth_trie::StorageProof, StateProofError> {
103 let hashed_address = keccak256(address);
104 let prefix_set = storage.construct_prefix_set();
105 let state_sorted = HashedPostStateSorted::new(
106 Default::default(),
107 HashMap::from_iter([(hashed_address, storage.into_sorted())]),
108 );
109 StorageProof::new(
110 DatabaseTrieCursorFactory::<_, A>::new(tx),
111 HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
112 address,
113 )
114 .with_prefix_set_mut(prefix_set)
115 .storage_proof(slot)
116 }
117
118 fn overlay_storage_multiproof(
119 tx: &'a TX,
120 address: Address,
121 slots: &[B256],
122 mut input: TrieInputSorted,
123 ) -> Result<StorageMultiProof, StateProofError> {
124 let hashed_address = keccak256(address);
125 let targets = slots.iter().map(keccak256).collect();
126 StorageProof::new(
127 InMemoryTrieCursorFactory::new(
128 DatabaseTrieCursorFactory::<_, A>::new(tx),
129 &input.nodes,
130 ),
131 HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &input.state),
132 address,
133 )
134 .with_prefix_set_mut(
135 input.prefix_sets.storage_prefix_sets.remove(&hashed_address).unwrap_or_default(),
136 )
137 .storage_multiproof(targets)
138 }
139}
140
141fn with_overlay<T: Clone, H: Clone, R>(
143 proof: &Proof<T, H>,
144 input: TrieInput,
145 f: impl FnOnce(
146 Proof<
147 InMemoryTrieCursorFactory<T, &TrieUpdatesSorted>,
148 HashedPostStateCursorFactory<H, &HashedPostStateSorted>,
149 >,
150 ) -> R,
151) -> R {
152 let nodes = input.nodes.into_sorted();
153 let state = input.state.into_sorted();
154 f(Proof::new(
155 InMemoryTrieCursorFactory::new(proof.trie_cursor_factory().clone(), &nodes),
156 HashedPostStateCursorFactory::new(proof.hashed_cursor_factory().clone(), &state),
157 )
158 .with_prefix_sets_mut(input.prefix_sets))
159}