reth_trie_db/
prefix_set.rs

1use alloy_primitives::{
2    map::{HashMap, HashSet},
3    BlockNumber, B256,
4};
5use core::{
6    marker::PhantomData,
7    ops::{Deref, RangeInclusive},
8};
9use reth_db_api::{
10    cursor::DbCursorRO,
11    models::{AccountBeforeTx, BlockNumberAddress},
12    tables,
13    transaction::DbTx,
14    DatabaseError,
15};
16use reth_primitives_traits::StorageEntry;
17use reth_trie::{
18    prefix_set::{PrefixSetMut, TriePrefixSets},
19    KeyHasher, Nibbles,
20};
21
22/// A wrapper around a database transaction that loads prefix sets within a given block range.
23#[derive(Debug)]
24pub struct PrefixSetLoader<'a, TX, KH>(&'a TX, PhantomData<KH>);
25
26impl<'a, TX, KH> PrefixSetLoader<'a, TX, KH> {
27    /// Create a new loader.
28    pub const fn new(tx: &'a TX) -> Self {
29        Self(tx, PhantomData)
30    }
31}
32
33impl<TX, KH> Deref for PrefixSetLoader<'_, TX, KH> {
34    type Target = TX;
35
36    fn deref(&self) -> &Self::Target {
37        self.0
38    }
39}
40
41impl<TX: DbTx, KH: KeyHasher> PrefixSetLoader<'_, TX, KH> {
42    /// Load all account and storage changes for the given block range.
43    pub fn load(self, range: RangeInclusive<BlockNumber>) -> Result<TriePrefixSets, DatabaseError> {
44        // Initialize prefix sets.
45        let mut account_prefix_set = PrefixSetMut::default();
46        let mut storage_prefix_sets = HashMap::<B256, PrefixSetMut>::default();
47        let mut destroyed_accounts = HashSet::default();
48
49        // Walk account changeset and insert account prefixes.
50        let mut account_changeset_cursor = self.cursor_read::<tables::AccountChangeSets>()?;
51        let mut account_hashed_state_cursor = self.cursor_read::<tables::HashedAccounts>()?;
52        for account_entry in account_changeset_cursor.walk_range(range.clone())? {
53            let (_, AccountBeforeTx { address, .. }) = account_entry?;
54            let hashed_address = KH::hash_key(address);
55            account_prefix_set.insert(Nibbles::unpack(hashed_address));
56
57            if account_hashed_state_cursor.seek_exact(hashed_address)?.is_none() {
58                destroyed_accounts.insert(hashed_address);
59            }
60        }
61
62        // Walk storage changeset and insert storage prefixes as well as account prefixes if missing
63        // from the account prefix set.
64        let mut storage_cursor = self.cursor_dup_read::<tables::StorageChangeSets>()?;
65        let storage_range = BlockNumberAddress::range(range);
66        for storage_entry in storage_cursor.walk_range(storage_range)? {
67            let (BlockNumberAddress((_, address)), StorageEntry { key, .. }) = storage_entry?;
68            let hashed_address = KH::hash_key(address);
69            account_prefix_set.insert(Nibbles::unpack(hashed_address));
70            storage_prefix_sets
71                .entry(hashed_address)
72                .or_default()
73                .insert(Nibbles::unpack(KH::hash_key(key)));
74        }
75
76        Ok(TriePrefixSets {
77            account_prefix_set: account_prefix_set.freeze(),
78            storage_prefix_sets: storage_prefix_sets
79                .into_iter()
80                .map(|(k, v)| (k, v.freeze()))
81                .collect(),
82            destroyed_accounts,
83        })
84    }
85}