reth_trie_common/
hashed_state.rs

1use core::ops::Not;
2
3use crate::{
4    added_removed_keys::MultiAddedRemovedKeys,
5    prefix_set::{PrefixSetMut, TriePrefixSetsMut},
6    KeyHasher, MultiProofTargets, Nibbles,
7};
8use alloc::{borrow::Cow, vec::Vec};
9use alloy_primitives::{
10    keccak256,
11    map::{hash_map, B256Map, B256Set, HashMap, HashSet},
12    Address, B256, U256,
13};
14use itertools::Itertools;
15#[cfg(feature = "rayon")]
16pub use rayon::*;
17use reth_primitives_traits::Account;
18
19#[cfg(feature = "rayon")]
20use rayon::prelude::{IntoParallelIterator, ParallelIterator};
21
22use revm_database::{AccountStatus, BundleAccount};
23
24/// Representation of in-memory hashed state.
25#[derive(PartialEq, Eq, Clone, Default, Debug)]
26#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
27pub struct HashedPostState {
28    /// Mapping of hashed address to account info, `None` if destroyed.
29    pub accounts: B256Map<Option<Account>>,
30    /// Mapping of hashed address to hashed storage.
31    pub storages: B256Map<HashedStorage>,
32}
33
34impl HashedPostState {
35    /// Create new instance of [`HashedPostState`].
36    pub fn with_capacity(capacity: usize) -> Self {
37        Self {
38            accounts: B256Map::with_capacity_and_hasher(capacity, Default::default()),
39            storages: B256Map::with_capacity_and_hasher(capacity, Default::default()),
40        }
41    }
42
43    /// Initialize [`HashedPostState`] from bundle state.
44    /// Hashes all changed accounts and storage entries that are currently stored in the bundle
45    /// state.
46    #[inline]
47    #[cfg(feature = "rayon")]
48    pub fn from_bundle_state<'a, KH: KeyHasher>(
49        state: impl IntoParallelIterator<Item = (&'a Address, &'a BundleAccount)>,
50    ) -> Self {
51        let hashed = state
52            .into_par_iter()
53            .map(|(address, account)| {
54                let hashed_address = KH::hash_key(address);
55                let hashed_account = account.info.as_ref().map(Into::into);
56                let hashed_storage = HashedStorage::from_plain_storage(
57                    account.status,
58                    account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
59                );
60                (hashed_address, (hashed_account, hashed_storage))
61            })
62            .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
63
64        let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
65        let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
66        for (address, (account, storage)) in hashed {
67            accounts.insert(address, account);
68            if !storage.is_empty() {
69                storages.insert(address, storage);
70            }
71        }
72        Self { accounts, storages }
73    }
74
75    /// Initialize [`HashedPostState`] from bundle state.
76    /// Hashes all changed accounts and storage entries that are currently stored in the bundle
77    /// state.
78    #[cfg(not(feature = "rayon"))]
79    pub fn from_bundle_state<'a, KH: KeyHasher>(
80        state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
81    ) -> Self {
82        let hashed = state
83            .into_iter()
84            .map(|(address, account)| {
85                let hashed_address = KH::hash_key(address);
86                let hashed_account = account.info.as_ref().map(Into::into);
87                let hashed_storage = HashedStorage::from_plain_storage(
88                    account.status,
89                    account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
90                );
91                (hashed_address, (hashed_account, hashed_storage))
92            })
93            .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
94
95        let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
96        let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
97        for (address, (account, storage)) in hashed {
98            accounts.insert(address, account);
99            if !storage.is_empty() {
100                storages.insert(address, storage);
101            }
102        }
103        Self { accounts, storages }
104    }
105
106    /// Construct [`HashedPostState`] from a single [`HashedStorage`].
107    pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
108        Self {
109            accounts: HashMap::default(),
110            storages: HashMap::from_iter([(hashed_address, storage)]),
111        }
112    }
113
114    /// Set account entries on hashed state.
115    pub fn with_accounts(
116        mut self,
117        accounts: impl IntoIterator<Item = (B256, Option<Account>)>,
118    ) -> Self {
119        self.accounts = HashMap::from_iter(accounts);
120        self
121    }
122
123    /// Set storage entries on hashed state.
124    pub fn with_storages(
125        mut self,
126        storages: impl IntoIterator<Item = (B256, HashedStorage)>,
127    ) -> Self {
128        self.storages = HashMap::from_iter(storages);
129        self
130    }
131
132    /// Returns `true` if the hashed state is empty.
133    pub fn is_empty(&self) -> bool {
134        self.accounts.is_empty() && self.storages.is_empty()
135    }
136
137    /// Construct [`TriePrefixSetsMut`] from hashed post state.
138    /// The prefix sets contain the hashed account and storage keys that have been changed in the
139    /// post state.
140    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
141        // Populate account prefix set.
142        let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
143        let mut destroyed_accounts = HashSet::default();
144        for (hashed_address, account) in &self.accounts {
145            account_prefix_set.insert(Nibbles::unpack(hashed_address));
146
147            if account.is_none() {
148                destroyed_accounts.insert(*hashed_address);
149            }
150        }
151
152        // Populate storage prefix sets.
153        let mut storage_prefix_sets =
154            HashMap::with_capacity_and_hasher(self.storages.len(), Default::default());
155        for (hashed_address, hashed_storage) in &self.storages {
156            account_prefix_set.insert(Nibbles::unpack(hashed_address));
157            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
158        }
159
160        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
161    }
162
163    /// Create multiproof targets for this state.
164    pub fn multi_proof_targets(&self) -> MultiProofTargets {
165        // Pre-allocate minimum capacity for the targets.
166        let mut targets = MultiProofTargets::with_capacity(self.accounts.len());
167        for hashed_address in self.accounts.keys() {
168            targets.insert(*hashed_address, Default::default());
169        }
170        for (hashed_address, storage) in &self.storages {
171            targets.entry(*hashed_address).or_default().extend(storage.storage.keys().copied());
172        }
173        targets
174    }
175
176    /// Create multiproof targets difference for this state,
177    /// i.e., the targets that are in targets create from `self` but not in `excluded`.
178    ///
179    /// This method is preferred to first calling `Self::multi_proof_targets` and the calling
180    /// `MultiProofTargets::retain_difference`, because it does not over allocate the targets map.
181    pub fn multi_proof_targets_difference(
182        &self,
183        excluded: &MultiProofTargets,
184    ) -> MultiProofTargets {
185        let mut targets = MultiProofTargets::default();
186        for hashed_address in self.accounts.keys() {
187            if !excluded.contains_key(hashed_address) {
188                targets.insert(*hashed_address, Default::default());
189            }
190        }
191        for (hashed_address, storage) in &self.storages {
192            let maybe_excluded_storage = excluded.get(hashed_address);
193            let mut hashed_slots_targets = storage
194                .storage
195                .keys()
196                .filter(|slot| !maybe_excluded_storage.is_some_and(|f| f.contains(*slot)))
197                .peekable();
198            if hashed_slots_targets.peek().is_some() {
199                targets.entry(*hashed_address).or_default().extend(hashed_slots_targets);
200            }
201        }
202        targets
203    }
204
205    /// Partition the state update into two state updates:
206    /// - First with accounts and storages slots that are present in the provided targets.
207    /// - Second with all other.
208    ///
209    /// CAUTION: The state updates are expected to be applied in order, so that the storage wipes
210    /// are done correctly.
211    pub fn partition_by_targets(
212        mut self,
213        targets: &MultiProofTargets,
214        added_removed_keys: &MultiAddedRemovedKeys,
215    ) -> (Self, Self) {
216        let mut state_updates_not_in_targets = Self::default();
217
218        self.storages.retain(|&address, storage| {
219            let storage_added_removed_keys = added_removed_keys.get_storage(&address);
220
221            let (retain, storage_not_in_targets) = match targets.get(&address) {
222                Some(storage_in_targets) => {
223                    let mut storage_not_in_targets = HashedStorage::default();
224                    storage.storage.retain(|&slot, value| {
225                        if storage_in_targets.contains(&slot) &&
226                            !storage_added_removed_keys.is_some_and(|k| k.is_removed(&slot))
227                        {
228                            return true
229                        }
230
231                        storage_not_in_targets.storage.insert(slot, *value);
232                        false
233                    });
234
235                    // We do not check the wiped flag here, because targets only contain addresses
236                    // and storage slots. So if there are no storage slots left, the storage update
237                    // can be fully removed.
238                    let retain = !storage.storage.is_empty();
239
240                    // Since state updates are expected to be applied in order, we can only set the
241                    // wiped flag in the second storage update if the first storage update is empty
242                    // and will not be retained.
243                    if !retain {
244                        storage_not_in_targets.wiped = storage.wiped;
245                    }
246
247                    (
248                        retain,
249                        storage_not_in_targets.is_empty().not().then_some(storage_not_in_targets),
250                    )
251                }
252                None => (false, Some(core::mem::take(storage))),
253            };
254
255            if let Some(storage_not_in_targets) = storage_not_in_targets {
256                state_updates_not_in_targets.storages.insert(address, storage_not_in_targets);
257            }
258
259            retain
260        });
261        self.accounts.retain(|&address, account| {
262            if targets.contains_key(&address) {
263                return true
264            }
265
266            state_updates_not_in_targets.accounts.insert(address, *account);
267            false
268        });
269
270        (self, state_updates_not_in_targets)
271    }
272
273    /// Returns an iterator that yields chunks of the specified size.
274    ///
275    /// See [`ChunkedHashedPostState`] for more information.
276    pub fn chunks(self, size: usize) -> ChunkedHashedPostState {
277        ChunkedHashedPostState::new(self, size)
278    }
279
280    /// Extend this hashed post state with contents of another.
281    /// Entries in the second hashed post state take precedence.
282    pub fn extend(&mut self, other: Self) {
283        self.extend_inner(Cow::Owned(other));
284    }
285
286    /// Extend this hashed post state with contents of another.
287    /// Entries in the second hashed post state take precedence.
288    ///
289    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
290    pub fn extend_ref(&mut self, other: &Self) {
291        self.extend_inner(Cow::Borrowed(other));
292    }
293
294    fn extend_inner(&mut self, other: Cow<'_, Self>) {
295        self.accounts.extend(other.accounts.iter().map(|(&k, &v)| (k, v)));
296
297        self.storages.reserve(other.storages.len());
298        match other {
299            Cow::Borrowed(other) => {
300                self.extend_storages(other.storages.iter().map(|(k, v)| (*k, Cow::Borrowed(v))))
301            }
302            Cow::Owned(other) => {
303                self.extend_storages(other.storages.into_iter().map(|(k, v)| (k, Cow::Owned(v))))
304            }
305        }
306    }
307
308    fn extend_storages<'a>(
309        &mut self,
310        storages: impl IntoIterator<Item = (B256, Cow<'a, HashedStorage>)>,
311    ) {
312        for (hashed_address, storage) in storages {
313            match self.storages.entry(hashed_address) {
314                hash_map::Entry::Vacant(entry) => {
315                    entry.insert(storage.into_owned());
316                }
317                hash_map::Entry::Occupied(mut entry) => {
318                    entry.get_mut().extend(&storage);
319                }
320            }
321        }
322    }
323
324    /// Converts hashed post state into [`HashedPostStateSorted`].
325    pub fn into_sorted(self) -> HashedPostStateSorted {
326        let mut updated_accounts = Vec::new();
327        let mut destroyed_accounts = HashSet::default();
328        for (hashed_address, info) in self.accounts {
329            if let Some(info) = info {
330                updated_accounts.push((hashed_address, info));
331            } else {
332                destroyed_accounts.insert(hashed_address);
333            }
334        }
335        updated_accounts.sort_unstable_by_key(|(address, _)| *address);
336        let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts };
337
338        let storages = self
339            .storages
340            .into_iter()
341            .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
342            .collect();
343
344        HashedPostStateSorted { accounts, storages }
345    }
346
347    /// Converts hashed post state into [`HashedPostStateSorted`], but keeping the maps allocated by
348    /// draining.
349    ///
350    /// This effectively clears all the fields in the [`HashedPostStateSorted`].
351    ///
352    /// This allows us to reuse the allocated space. This allocates new space for the sorted hashed
353    /// post state, like `into_sorted`.
354    pub fn drain_into_sorted(&mut self) -> HashedPostStateSorted {
355        let mut updated_accounts = Vec::new();
356        let mut destroyed_accounts = HashSet::default();
357        for (hashed_address, info) in self.accounts.drain() {
358            if let Some(info) = info {
359                updated_accounts.push((hashed_address, info));
360            } else {
361                destroyed_accounts.insert(hashed_address);
362            }
363        }
364        updated_accounts.sort_unstable_by_key(|(address, _)| *address);
365        let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts };
366
367        let storages = self
368            .storages
369            .drain()
370            .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
371            .collect();
372
373        HashedPostStateSorted { accounts, storages }
374    }
375
376    /// Clears the account and storage maps of this `HashedPostState`.
377    pub fn clear(&mut self) {
378        self.accounts.clear();
379        self.storages.clear();
380    }
381}
382
383/// Representation of in-memory hashed storage.
384#[derive(PartialEq, Eq, Clone, Debug, Default)]
385#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
386pub struct HashedStorage {
387    /// Flag indicating whether the storage was wiped or not.
388    pub wiped: bool,
389    /// Mapping of hashed storage slot to storage value.
390    pub storage: B256Map<U256>,
391}
392
393impl HashedStorage {
394    /// Create new instance of [`HashedStorage`].
395    pub fn new(wiped: bool) -> Self {
396        Self { wiped, storage: HashMap::default() }
397    }
398
399    /// Check if self is empty.
400    pub fn is_empty(&self) -> bool {
401        !self.wiped && self.storage.is_empty()
402    }
403
404    /// Create new hashed storage from iterator.
405    pub fn from_iter(wiped: bool, iter: impl IntoIterator<Item = (B256, U256)>) -> Self {
406        Self { wiped, storage: HashMap::from_iter(iter) }
407    }
408
409    /// Create new hashed storage from account status and plain storage.
410    pub fn from_plain_storage<'a>(
411        status: AccountStatus,
412        storage: impl IntoIterator<Item = (&'a U256, &'a U256)>,
413    ) -> Self {
414        Self::from_iter(
415            status.was_destroyed(),
416            storage.into_iter().map(|(key, value)| (keccak256(B256::from(*key)), *value)),
417        )
418    }
419
420    /// Construct [`PrefixSetMut`] from hashed storage.
421    pub fn construct_prefix_set(&self) -> PrefixSetMut {
422        if self.wiped {
423            PrefixSetMut::all()
424        } else {
425            let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
426            for hashed_slot in self.storage.keys() {
427                prefix_set.insert(Nibbles::unpack(hashed_slot));
428            }
429            prefix_set
430        }
431    }
432
433    /// Extend hashed storage with contents of other.
434    /// The entries in second hashed storage take precedence.
435    pub fn extend(&mut self, other: &Self) {
436        if other.wiped {
437            self.wiped = true;
438            self.storage.clear();
439        }
440        self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v)));
441    }
442
443    /// Converts hashed storage into [`HashedStorageSorted`].
444    pub fn into_sorted(self) -> HashedStorageSorted {
445        let mut non_zero_valued_slots = Vec::new();
446        let mut zero_valued_slots = HashSet::default();
447        for (hashed_slot, value) in self.storage {
448            if value.is_zero() {
449                zero_valued_slots.insert(hashed_slot);
450            } else {
451                non_zero_valued_slots.push((hashed_slot, value));
452            }
453        }
454        non_zero_valued_slots.sort_unstable_by_key(|(key, _)| *key);
455
456        HashedStorageSorted { non_zero_valued_slots, zero_valued_slots, wiped: self.wiped }
457    }
458}
459
460/// Sorted hashed post state optimized for iterating during state trie calculation.
461#[derive(PartialEq, Eq, Clone, Default, Debug)]
462pub struct HashedPostStateSorted {
463    /// Updated state of accounts.
464    pub accounts: HashedAccountsSorted,
465    /// Map of hashed addresses to hashed storage.
466    pub storages: B256Map<HashedStorageSorted>,
467}
468
469impl HashedPostStateSorted {
470    /// Create new instance of [`HashedPostStateSorted`]
471    pub const fn new(
472        accounts: HashedAccountsSorted,
473        storages: B256Map<HashedStorageSorted>,
474    ) -> Self {
475        Self { accounts, storages }
476    }
477
478    /// Returns reference to hashed accounts.
479    pub const fn accounts(&self) -> &HashedAccountsSorted {
480        &self.accounts
481    }
482
483    /// Returns reference to hashed account storages.
484    pub const fn account_storages(&self) -> &B256Map<HashedStorageSorted> {
485        &self.storages
486    }
487}
488
489/// Sorted account state optimized for iterating during state trie calculation.
490#[derive(Clone, Eq, PartialEq, Default, Debug)]
491pub struct HashedAccountsSorted {
492    /// Sorted collection of hashed addresses and their account info.
493    pub accounts: Vec<(B256, Account)>,
494    /// Set of destroyed account keys.
495    pub destroyed_accounts: B256Set,
496}
497
498impl HashedAccountsSorted {
499    /// Returns a sorted iterator over updated accounts.
500    pub fn accounts_sorted(&self) -> impl Iterator<Item = (B256, Option<Account>)> {
501        self.accounts
502            .iter()
503            .map(|(address, account)| (*address, Some(*account)))
504            .chain(self.destroyed_accounts.iter().map(|address| (*address, None)))
505            .sorted_by_key(|entry| *entry.0)
506    }
507}
508
509/// Sorted hashed storage optimized for iterating during state trie calculation.
510#[derive(Clone, Eq, PartialEq, Debug)]
511pub struct HashedStorageSorted {
512    /// Sorted hashed storage slots with non-zero value.
513    pub non_zero_valued_slots: Vec<(B256, U256)>,
514    /// Slots that have been zero valued.
515    pub zero_valued_slots: B256Set,
516    /// Flag indicating whether the storage was wiped or not.
517    pub wiped: bool,
518}
519
520impl HashedStorageSorted {
521    /// Returns `true` if the account was wiped.
522    pub const fn is_wiped(&self) -> bool {
523        self.wiped
524    }
525
526    /// Returns a sorted iterator over updated storage slots.
527    pub fn storage_slots_sorted(&self) -> impl Iterator<Item = (B256, U256)> {
528        self.non_zero_valued_slots
529            .iter()
530            .map(|(hashed_slot, value)| (*hashed_slot, *value))
531            .chain(self.zero_valued_slots.iter().map(|hashed_slot| (*hashed_slot, U256::ZERO)))
532            .sorted_by_key(|entry| *entry.0)
533    }
534}
535
536/// An iterator that yields chunks of the state updates of at most `size` account and storage
537/// targets.
538///
539/// # Notes
540/// 1. Chunks are expected to be applied in order, because of storage wipes. If applied out of
541///    order, it's possible to wipe more storage than in the original state update.
542/// 2. For each account, chunks with storage updates come first, followed by account updates.
543#[derive(Debug)]
544pub struct ChunkedHashedPostState {
545    flattened: alloc::vec::IntoIter<(B256, FlattenedHashedPostStateItem)>,
546    size: usize,
547}
548
549#[derive(Debug)]
550enum FlattenedHashedPostStateItem {
551    Account(Option<Account>),
552    StorageWipe,
553    StorageUpdate { slot: B256, value: U256 },
554}
555
556impl ChunkedHashedPostState {
557    fn new(hashed_post_state: HashedPostState, size: usize) -> Self {
558        let flattened = hashed_post_state
559            .storages
560            .into_iter()
561            .flat_map(|(address, storage)| {
562                // Storage wipes should go first
563                Some((address, FlattenedHashedPostStateItem::StorageWipe))
564                    .filter(|_| storage.wiped)
565                    .into_iter()
566                    .chain(
567                        storage.storage.into_iter().sorted_unstable_by_key(|(slot, _)| *slot).map(
568                            move |(slot, value)| {
569                                (
570                                    address,
571                                    FlattenedHashedPostStateItem::StorageUpdate { slot, value },
572                                )
573                            },
574                        ),
575                    )
576            })
577            .chain(hashed_post_state.accounts.into_iter().map(|(address, account)| {
578                (address, FlattenedHashedPostStateItem::Account(account))
579            }))
580            // We need stable sort here to preserve the order for each address:
581            // 1. Storage wipes
582            // 2. Storage updates
583            // 3. Account update
584            .sorted_by_key(|(address, _)| *address);
585
586        Self { flattened, size }
587    }
588}
589
590impl Iterator for ChunkedHashedPostState {
591    type Item = HashedPostState;
592
593    fn next(&mut self) -> Option<Self::Item> {
594        let mut chunk = HashedPostState::default();
595
596        let mut current_size = 0;
597        while current_size < self.size {
598            let Some((address, item)) = self.flattened.next() else { break };
599
600            match item {
601                FlattenedHashedPostStateItem::Account(account) => {
602                    chunk.accounts.insert(address, account);
603                }
604                FlattenedHashedPostStateItem::StorageWipe => {
605                    chunk.storages.entry(address).or_default().wiped = true;
606                }
607                FlattenedHashedPostStateItem::StorageUpdate { slot, value } => {
608                    chunk.storages.entry(address).or_default().storage.insert(slot, value);
609                }
610            }
611
612            current_size += 1;
613        }
614
615        if chunk.is_empty() {
616            None
617        } else {
618            Some(chunk)
619        }
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626    use crate::KeccakKeyHasher;
627    use alloy_primitives::Bytes;
628    use revm_database::{states::StorageSlot, StorageWithOriginalValues};
629    use revm_state::{AccountInfo, Bytecode};
630
631    #[test]
632    fn hashed_state_wiped_extension() {
633        let hashed_address = B256::default();
634        let hashed_slot = B256::with_last_byte(64);
635        let hashed_slot2 = B256::with_last_byte(65);
636
637        // Initialize post state storage
638        let original_slot_value = U256::from(123);
639        let mut hashed_state = HashedPostState::default().with_storages([(
640            hashed_address,
641            HashedStorage::from_iter(
642                false,
643                [(hashed_slot, original_slot_value), (hashed_slot2, original_slot_value)],
644            ),
645        )]);
646
647        // Update single slot value
648        let updated_slot_value = U256::from(321);
649        let extension = HashedPostState::default().with_storages([(
650            hashed_address,
651            HashedStorage::from_iter(false, [(hashed_slot, updated_slot_value)]),
652        )]);
653        hashed_state.extend(extension);
654
655        let account_storage = hashed_state.storages.get(&hashed_address);
656        assert_eq!(
657            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
658            Some(&updated_slot_value)
659        );
660        assert_eq!(
661            account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
662            Some(&original_slot_value)
663        );
664        assert_eq!(account_storage.map(|st| st.wiped), Some(false));
665
666        // Wipe account storage
667        let wiped_extension =
668            HashedPostState::default().with_storages([(hashed_address, HashedStorage::new(true))]);
669        hashed_state.extend(wiped_extension);
670
671        let account_storage = hashed_state.storages.get(&hashed_address);
672        assert_eq!(account_storage.map(|st| st.storage.is_empty()), Some(true));
673        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
674
675        // Reinitialize single slot value
676        hashed_state.extend(HashedPostState::default().with_storages([(
677            hashed_address,
678            HashedStorage::from_iter(false, [(hashed_slot, original_slot_value)]),
679        )]));
680        let account_storage = hashed_state.storages.get(&hashed_address);
681        assert_eq!(
682            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
683            Some(&original_slot_value)
684        );
685        assert_eq!(account_storage.and_then(|st| st.storage.get(&hashed_slot2)), None);
686        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
687
688        // Reinitialize single slot value
689        hashed_state.extend(HashedPostState::default().with_storages([(
690            hashed_address,
691            HashedStorage::from_iter(false, [(hashed_slot2, updated_slot_value)]),
692        )]));
693        let account_storage = hashed_state.storages.get(&hashed_address);
694        assert_eq!(
695            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
696            Some(&original_slot_value)
697        );
698        assert_eq!(
699            account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
700            Some(&updated_slot_value)
701        );
702        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
703    }
704
705    #[test]
706    fn test_hashed_post_state_from_bundle_state() {
707        // Prepare a random Ethereum address as a key for the account.
708        let address = Address::random();
709
710        // Create a mock account info object.
711        let account_info = AccountInfo {
712            balance: U256::from(123),
713            nonce: 42,
714            code_hash: B256::random(),
715            code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2]))),
716        };
717
718        let mut storage = StorageWithOriginalValues::default();
719        storage.insert(
720            U256::from(1),
721            StorageSlot { present_value: U256::from(4), ..Default::default() },
722        );
723
724        // Create a `BundleAccount` struct to represent the account and its storage.
725        let account = BundleAccount {
726            status: AccountStatus::Changed,
727            info: Some(account_info.clone()),
728            storage,
729            original_info: None,
730        };
731
732        // Create a vector of tuples representing the bundle state.
733        let state = vec![(&address, &account)];
734
735        // Convert the bundle state into a hashed post state.
736        let hashed_state = HashedPostState::from_bundle_state::<KeccakKeyHasher>(state);
737
738        // Validate the hashed post state.
739        assert_eq!(hashed_state.accounts.len(), 1);
740        assert_eq!(hashed_state.storages.len(), 1);
741
742        // Validate the account info.
743        assert_eq!(
744            *hashed_state.accounts.get(&keccak256(address)).unwrap(),
745            Some(account_info.into())
746        );
747    }
748
749    #[test]
750    fn test_hashed_post_state_with_accounts() {
751        // Prepare random addresses and mock account info.
752        let address_1 = Address::random();
753        let address_2 = Address::random();
754
755        let account_info_1 = AccountInfo {
756            balance: U256::from(1000),
757            nonce: 1,
758            code_hash: B256::random(),
759            code: None,
760        };
761
762        // Create hashed accounts with addresses.
763        let account_1 = (keccak256(address_1), Some(account_info_1.into()));
764        let account_2 = (keccak256(address_2), None);
765
766        // Add accounts to the hashed post state.
767        let hashed_state = HashedPostState::default().with_accounts(vec![account_1, account_2]);
768
769        // Validate the hashed post state.
770        assert_eq!(hashed_state.accounts.len(), 2);
771        assert!(hashed_state.accounts.contains_key(&keccak256(address_1)));
772        assert!(hashed_state.accounts.contains_key(&keccak256(address_2)));
773    }
774
775    #[test]
776    fn test_hashed_post_state_with_storages() {
777        // Prepare random addresses and mock storage entries.
778        let address_1 = Address::random();
779        let address_2 = Address::random();
780
781        let storage_1 = (keccak256(address_1), HashedStorage::new(false));
782        let storage_2 = (keccak256(address_2), HashedStorage::new(true));
783
784        // Add storages to the hashed post state.
785        let hashed_state = HashedPostState::default().with_storages(vec![storage_1, storage_2]);
786
787        // Validate the hashed post state.
788        assert_eq!(hashed_state.storages.len(), 2);
789        assert!(hashed_state.storages.contains_key(&keccak256(address_1)));
790        assert!(hashed_state.storages.contains_key(&keccak256(address_2)));
791    }
792
793    #[test]
794    fn test_hashed_post_state_is_empty() {
795        // Create an empty hashed post state and validate it's empty.
796        let empty_state = HashedPostState::default();
797        assert!(empty_state.is_empty());
798
799        // Add an account and validate the state is no longer empty.
800        let non_empty_state = HashedPostState::default()
801            .with_accounts(vec![(keccak256(Address::random()), Some(Account::default()))]);
802        assert!(!non_empty_state.is_empty());
803    }
804
805    fn create_state_for_multi_proof_targets() -> HashedPostState {
806        let mut state = HashedPostState::default();
807
808        let addr1 = B256::random();
809        let addr2 = B256::random();
810        state.accounts.insert(addr1, Some(Default::default()));
811        state.accounts.insert(addr2, Some(Default::default()));
812
813        let mut storage = HashedStorage::default();
814        let slot1 = B256::random();
815        let slot2 = B256::random();
816        storage.storage.insert(slot1, U256::ZERO);
817        storage.storage.insert(slot2, U256::from(1));
818        state.storages.insert(addr1, storage);
819
820        state
821    }
822
823    #[test]
824    fn test_multi_proof_targets_difference_empty_state() {
825        let state = HashedPostState::default();
826        let excluded = MultiProofTargets::default();
827
828        let targets = state.multi_proof_targets_difference(&excluded);
829        assert!(targets.is_empty());
830    }
831
832    #[test]
833    fn test_multi_proof_targets_difference_new_account_targets() {
834        let state = create_state_for_multi_proof_targets();
835        let excluded = MultiProofTargets::default();
836
837        // should return all accounts as targets since excluded is empty
838        let targets = state.multi_proof_targets_difference(&excluded);
839        assert_eq!(targets.len(), state.accounts.len());
840        for addr in state.accounts.keys() {
841            assert!(targets.contains_key(addr));
842        }
843    }
844
845    #[test]
846    fn test_multi_proof_targets_difference_new_storage_targets() {
847        let state = create_state_for_multi_proof_targets();
848        let excluded = MultiProofTargets::default();
849
850        let targets = state.multi_proof_targets_difference(&excluded);
851
852        // verify storage slots are included for accounts with storage
853        for (addr, storage) in &state.storages {
854            assert!(targets.contains_key(addr));
855            let target_slots = &targets[addr];
856            assert_eq!(target_slots.len(), storage.storage.len());
857            for slot in storage.storage.keys() {
858                assert!(target_slots.contains(slot));
859            }
860        }
861    }
862
863    #[test]
864    fn test_multi_proof_targets_difference_filter_excluded_accounts() {
865        let state = create_state_for_multi_proof_targets();
866        let mut excluded = MultiProofTargets::default();
867
868        // select an account that has no storage updates
869        let excluded_addr = state
870            .accounts
871            .keys()
872            .find(|&&addr| !state.storages.contains_key(&addr))
873            .expect("Should have an account without storage");
874
875        // mark the account as excluded
876        excluded.insert(*excluded_addr, HashSet::default());
877
878        let targets = state.multi_proof_targets_difference(&excluded);
879
880        // should not include the already excluded account since it has no storage updates
881        assert!(!targets.contains_key(excluded_addr));
882        // other accounts should still be included
883        assert_eq!(targets.len(), state.accounts.len() - 1);
884    }
885
886    #[test]
887    fn test_multi_proof_targets_difference_filter_excluded_storage() {
888        let state = create_state_for_multi_proof_targets();
889        let mut excluded = MultiProofTargets::default();
890
891        // mark one storage slot as excluded
892        let (addr, storage) = state.storages.iter().next().unwrap();
893        let mut excluded_slots = HashSet::default();
894        let excluded_slot = *storage.storage.keys().next().unwrap();
895        excluded_slots.insert(excluded_slot);
896        excluded.insert(*addr, excluded_slots);
897
898        let targets = state.multi_proof_targets_difference(&excluded);
899
900        // should not include the excluded storage slot
901        let target_slots = &targets[addr];
902        assert!(!target_slots.contains(&excluded_slot));
903        assert_eq!(target_slots.len(), storage.storage.len() - 1);
904    }
905
906    #[test]
907    fn test_multi_proof_targets_difference_mixed_excluded_state() {
908        let mut state = HashedPostState::default();
909        let mut excluded = MultiProofTargets::default();
910
911        let addr1 = B256::random();
912        let addr2 = B256::random();
913        let slot1 = B256::random();
914        let slot2 = B256::random();
915
916        state.accounts.insert(addr1, Some(Default::default()));
917        state.accounts.insert(addr2, Some(Default::default()));
918
919        let mut storage = HashedStorage::default();
920        storage.storage.insert(slot1, U256::ZERO);
921        storage.storage.insert(slot2, U256::from(1));
922        state.storages.insert(addr1, storage);
923
924        let mut excluded_slots = HashSet::default();
925        excluded_slots.insert(slot1);
926        excluded.insert(addr1, excluded_slots);
927
928        let targets = state.multi_proof_targets_difference(&excluded);
929
930        assert!(targets.contains_key(&addr2));
931        assert!(!targets[&addr1].contains(&slot1));
932        assert!(targets[&addr1].contains(&slot2));
933    }
934
935    #[test]
936    fn test_multi_proof_targets_difference_unmodified_account_with_storage() {
937        let mut state = HashedPostState::default();
938        let excluded = MultiProofTargets::default();
939
940        let addr = B256::random();
941        let slot1 = B256::random();
942        let slot2 = B256::random();
943
944        // don't add the account to state.accounts (simulating unmodified account)
945        // but add storage updates for this account
946        let mut storage = HashedStorage::default();
947        storage.storage.insert(slot1, U256::from(1));
948        storage.storage.insert(slot2, U256::from(2));
949        state.storages.insert(addr, storage);
950
951        assert!(!state.accounts.contains_key(&addr));
952        assert!(!excluded.contains_key(&addr));
953
954        let targets = state.multi_proof_targets_difference(&excluded);
955
956        // verify that we still get the storage slots for the unmodified account
957        assert!(targets.contains_key(&addr));
958
959        let target_slots = &targets[&addr];
960        assert_eq!(target_slots.len(), 2);
961        assert!(target_slots.contains(&slot1));
962        assert!(target_slots.contains(&slot2));
963    }
964
965    #[test]
966    fn test_partition_by_targets() {
967        let addr1 = B256::random();
968        let addr2 = B256::random();
969        let slot1 = B256::random();
970        let slot2 = B256::random();
971
972        let state = HashedPostState {
973            accounts: B256Map::from_iter([
974                (addr1, Some(Default::default())),
975                (addr2, Some(Default::default())),
976            ]),
977            storages: B256Map::from_iter([(
978                addr1,
979                HashedStorage {
980                    wiped: true,
981                    storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
982                },
983            )]),
984        };
985        let targets = MultiProofTargets::from_iter([(addr1, HashSet::from_iter([slot1]))]);
986
987        let (with_targets, without_targets) =
988            state.partition_by_targets(&targets, &MultiAddedRemovedKeys::new());
989
990        assert_eq!(
991            with_targets,
992            HashedPostState {
993                accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
994                storages: B256Map::from_iter([(
995                    addr1,
996                    HashedStorage {
997                        wiped: true,
998                        storage: B256Map::from_iter([(slot1, U256::ZERO)])
999                    }
1000                )]),
1001            }
1002        );
1003        assert_eq!(
1004            without_targets,
1005            HashedPostState {
1006                accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1007                storages: B256Map::from_iter([(
1008                    addr1,
1009                    HashedStorage {
1010                        wiped: false,
1011                        storage: B256Map::from_iter([(slot2, U256::from(1))])
1012                    }
1013                )]),
1014            }
1015        );
1016    }
1017
1018    #[test]
1019    fn test_chunks() {
1020        let addr1 = B256::from([1; 32]);
1021        let addr2 = B256::from([2; 32]);
1022        let slot1 = B256::from([1; 32]);
1023        let slot2 = B256::from([2; 32]);
1024
1025        let state = HashedPostState {
1026            accounts: B256Map::from_iter([
1027                (addr1, Some(Default::default())),
1028                (addr2, Some(Default::default())),
1029            ]),
1030            storages: B256Map::from_iter([(
1031                addr2,
1032                HashedStorage {
1033                    wiped: true,
1034                    storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1035                },
1036            )]),
1037        };
1038
1039        let mut chunks = state.chunks(2);
1040        assert_eq!(
1041            chunks.next(),
1042            Some(HashedPostState {
1043                accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
1044                storages: B256Map::from_iter([(addr2, HashedStorage::new(true)),])
1045            })
1046        );
1047        assert_eq!(
1048            chunks.next(),
1049            Some(HashedPostState {
1050                accounts: B256Map::default(),
1051                storages: B256Map::from_iter([(
1052                    addr2,
1053                    HashedStorage {
1054                        wiped: false,
1055                        storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1056                    },
1057                )])
1058            })
1059        );
1060        assert_eq!(
1061            chunks.next(),
1062            Some(HashedPostState {
1063                accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1064                storages: B256Map::default()
1065            })
1066        );
1067        assert_eq!(chunks.next(), None);
1068    }
1069}