use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
use alloy_primitives::{
keccak256,
map::{HashMap, HashSet},
Address, B256,
};
use reth_db_api::transaction::DbTx;
use reth_execution_errors::StateProofError;
use reth_trie::{
hashed_cursor::HashedPostStateCursorFactory,
proof::{Proof, StorageProof},
trie_cursor::InMemoryTrieCursorFactory,
HashedPostStateSorted, HashedStorage, MultiProof, TrieInput,
};
use reth_trie_common::AccountProof;
pub trait DatabaseProof<'a, TX> {
fn from_tx(tx: &'a TX) -> Self;
fn overlay_account_proof(
tx: &'a TX,
input: TrieInput,
address: Address,
slots: &[B256],
) -> Result<AccountProof, StateProofError>;
fn overlay_multiproof(
tx: &'a TX,
input: TrieInput,
targets: HashMap<B256, HashSet<B256>>,
) -> Result<MultiProof, StateProofError>;
}
impl<'a, TX: DbTx> DatabaseProof<'a, TX>
for Proof<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>
{
fn from_tx(tx: &'a TX) -> Self {
Self::new(DatabaseTrieCursorFactory::new(tx), DatabaseHashedCursorFactory::new(tx))
}
fn overlay_account_proof(
tx: &'a TX,
input: TrieInput,
address: Address,
slots: &[B256],
) -> Result<AccountProof, StateProofError> {
let nodes_sorted = input.nodes.into_sorted();
let state_sorted = input.state.into_sorted();
Self::from_tx(tx)
.with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
DatabaseTrieCursorFactory::new(tx),
&nodes_sorted,
))
.with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
DatabaseHashedCursorFactory::new(tx),
&state_sorted,
))
.with_prefix_sets_mut(input.prefix_sets)
.account_proof(address, slots)
}
fn overlay_multiproof(
tx: &'a TX,
input: TrieInput,
targets: HashMap<B256, HashSet<B256>>,
) -> Result<MultiProof, StateProofError> {
let nodes_sorted = input.nodes.into_sorted();
let state_sorted = input.state.into_sorted();
Self::from_tx(tx)
.with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
DatabaseTrieCursorFactory::new(tx),
&nodes_sorted,
))
.with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
DatabaseHashedCursorFactory::new(tx),
&state_sorted,
))
.with_prefix_sets_mut(input.prefix_sets)
.multiproof(targets)
}
}
pub trait DatabaseStorageProof<'a, TX> {
fn from_tx(tx: &'a TX, address: Address) -> Self;
fn overlay_storage_proof(
tx: &'a TX,
address: Address,
slot: B256,
storage: HashedStorage,
) -> Result<reth_trie_common::StorageProof, StateProofError>;
}
impl<'a, TX: DbTx> DatabaseStorageProof<'a, TX>
for StorageProof<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>
{
fn from_tx(tx: &'a TX, address: Address) -> Self {
Self::new(DatabaseTrieCursorFactory::new(tx), DatabaseHashedCursorFactory::new(tx), address)
}
fn overlay_storage_proof(
tx: &'a TX,
address: Address,
slot: B256,
storage: HashedStorage,
) -> Result<reth_trie_common::StorageProof, StateProofError> {
let hashed_address = keccak256(address);
let prefix_set = storage.construct_prefix_set();
let state_sorted = HashedPostStateSorted::new(
Default::default(),
HashMap::from([(hashed_address, storage.into_sorted())]),
);
Self::from_tx(tx, address)
.with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
DatabaseHashedCursorFactory::new(tx),
&state_sorted,
))
.with_prefix_set_mut(prefix_set)
.storage_proof(slot)
}
}