reth_trie_parallel/
storage_root_targets.rs

1use alloy_primitives::{map::B256Map, B256};
2use derive_more::{Deref, DerefMut};
3use reth_trie::prefix_set::PrefixSet;
4
5/// Target accounts with corresponding prefix sets for storage root calculation.
6#[derive(Deref, DerefMut, Debug)]
7pub struct StorageRootTargets(B256Map<PrefixSet>);
8
9impl StorageRootTargets {
10    /// Create new storage root targets from updated post state accounts
11    /// and storage prefix sets.
12    ///
13    /// NOTE: Since updated accounts and prefix sets always overlap,
14    /// it's important that iterator over storage prefix sets takes precedence.
15    pub fn new(
16        changed_accounts: impl IntoIterator<Item = B256>,
17        storage_prefix_sets: impl IntoIterator<Item = (B256, PrefixSet)>,
18    ) -> Self {
19        Self(
20            changed_accounts
21                .into_iter()
22                .map(|address| (address, PrefixSet::default()))
23                .chain(storage_prefix_sets)
24                .collect(),
25        )
26    }
27
28    /// Returns the total number of unique storage root targets without allocating new maps.
29    pub fn count(
30        account_prefix_set: &PrefixSet,
31        storage_prefix_sets: &B256Map<PrefixSet>,
32    ) -> usize {
33        let mut count = storage_prefix_sets.len();
34
35        for nibbles in account_prefix_set {
36            let hashed_address = B256::from_slice(&nibbles.pack());
37            if !storage_prefix_sets.contains_key(&hashed_address) {
38                count += 1;
39            }
40        }
41
42        count
43    }
44}
45
46impl IntoIterator for StorageRootTargets {
47    type Item = (B256, PrefixSet);
48    type IntoIter = std::collections::hash_map::IntoIter<B256, PrefixSet>;
49
50    fn into_iter(self) -> Self::IntoIter {
51        self.0.into_iter()
52    }
53}
54
55impl rayon::iter::IntoParallelIterator for StorageRootTargets {
56    type Iter = rayon::collections::hash_map::IntoIter<B256, PrefixSet>;
57    type Item = (B256, PrefixSet);
58
59    fn into_par_iter(self) -> Self::Iter {
60        self.0.into_par_iter()
61    }
62}