reth_trie_db/
prefix_set.rs

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