Skip to main content

reth_trie_db/
proof.rs

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
14/// Extends [`Proof`] with operations specific for working with a database transaction.
15pub trait DatabaseProof<'a> {
16    /// Associated type for the database transaction.
17    type Tx;
18
19    /// Create a new [`Proof`] instance from database transaction.
20    fn from_tx(tx: &'a Self::Tx) -> Self;
21
22    /// Generates the state proof for target account based on [`TrieInput`].
23    fn overlay_account_proof(
24        &self,
25        input: TrieInput,
26        address: Address,
27        slots: &[B256],
28    ) -> Result<AccountProof, StateProofError>;
29
30    /// Generates the state [`MultiProof`] for target hashed account and storage keys.
31    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
64/// Extends [`StorageProof`] with operations specific for working with a database transaction.
65pub trait DatabaseStorageProof<'a, TX> {
66    /// Create a new [`StorageProof`] from database transaction and account address.
67    fn from_tx(tx: &'a TX, address: Address) -> Self;
68
69    /// Generates the storage proof for target slot based on [`TrieInput`].
70    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    /// Generates the storage multiproof for target slots based on [`TrieInput`].
78    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
141// Keeps sorting, cursor layering, and prefix invalidation identical across proof entrypoints.
142fn 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}