Skip to main content

reth_trie_common/
hashed_state.rs

1use core::ops::Not;
2
3use crate::{
4    added_removed_keys::MultiAddedRemovedKeys,
5    prefix_set::{PrefixSetMut, TriePrefixSetsMut},
6    utils::{extend_sorted_vec, kway_merge_disjoint_sorted, kway_merge_sorted},
7    KeyHasher, MultiProofTargets, Nibbles,
8};
9use alloc::{borrow::Cow, vec::Vec};
10use alloy_primitives::{
11    keccak256,
12    map::{hash_map, B256Map, HashMap, HashSet},
13    Address, B256, U256,
14};
15use itertools::Itertools;
16#[cfg(feature = "rayon")]
17pub use rayon::*;
18use reth_primitives_traits::Account;
19
20#[cfg(feature = "rayon")]
21use rayon::prelude::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
22
23use revm::database::{AccountStatus, BundleAccount};
24
25/// In-memory hashed state that stores account and storage changes with keccak256-hashed keys in
26/// hash maps.
27#[derive(PartialEq, Eq, Clone, Default, Debug)]
28#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
29pub struct HashedPostState {
30    /// Mapping of hashed address to account info, `None` if destroyed.
31    pub accounts: B256Map<Option<Account>>,
32    /// Mapping of hashed address to hashed storage.
33    pub storages: B256Map<HashedStorage>,
34}
35
36impl HashedPostState {
37    /// Create new instance of [`HashedPostState`].
38    pub fn with_capacity(capacity: usize) -> Self {
39        Self {
40            accounts: B256Map::with_capacity_and_hasher(capacity, Default::default()),
41            storages: B256Map::with_capacity_and_hasher(capacity, Default::default()),
42        }
43    }
44
45    /// Initialize [`HashedPostState`] from bundle state.
46    /// Hashes all changed accounts and storage entries that are currently stored in the bundle
47    /// state.
48    #[inline]
49    pub fn from_bundle_state<'a, KH: KeyHasher>(
50        state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
51    ) -> Self {
52        state
53            .into_iter()
54            .map(|(address, account)| {
55                let hashed_address = KH::hash_key(address);
56                let hashed_account = account.info.as_ref().map(Into::into);
57                let hashed_storage = HashedStorage::from_iter(
58                    false,
59                    account
60                        .storage
61                        .iter()
62                        .map(|(slot, value)| (keccak256(B256::from(*slot)), value.present_value)),
63                );
64
65                (
66                    hashed_address,
67                    hashed_account,
68                    (!hashed_storage.is_empty()).then_some(hashed_storage),
69                )
70            })
71            .collect()
72    }
73
74    /// Construct [`HashedPostState`] from a single [`HashedStorage`].
75    pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
76        Self {
77            accounts: HashMap::default(),
78            storages: HashMap::from_iter([(hashed_address, storage)]),
79        }
80    }
81
82    /// Set account entries on hashed state.
83    pub fn with_accounts(
84        mut self,
85        accounts: impl IntoIterator<Item = (B256, Option<Account>)>,
86    ) -> Self {
87        self.accounts = HashMap::from_iter(accounts);
88        self
89    }
90
91    /// Set storage entries on hashed state.
92    pub fn with_storages(
93        mut self,
94        storages: impl IntoIterator<Item = (B256, HashedStorage)>,
95    ) -> Self {
96        self.storages = HashMap::from_iter(storages);
97        self
98    }
99
100    /// Returns `true` if the hashed state is empty.
101    pub fn is_empty(&self) -> bool {
102        self.accounts.is_empty() && self.storages.is_empty()
103    }
104
105    /// Construct [`TriePrefixSetsMut`] from hashed post state.
106    /// The prefix sets contain the hashed account and storage keys that have been changed in the
107    /// post state.
108    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
109        // Populate account prefix set.
110        let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
111        let mut destroyed_accounts = HashSet::default();
112        for (hashed_address, account) in &self.accounts {
113            account_prefix_set.insert(Nibbles::unpack(hashed_address));
114
115            if account.is_none() {
116                destroyed_accounts.insert(*hashed_address);
117            }
118        }
119
120        // Populate storage prefix sets.
121        let mut storage_prefix_sets =
122            HashMap::with_capacity_and_hasher(self.storages.len(), Default::default());
123        for (hashed_address, hashed_storage) in &self.storages {
124            account_prefix_set.insert(Nibbles::unpack(hashed_address));
125            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
126        }
127
128        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
129    }
130
131    /// Create multiproof targets for this state.
132    pub fn multi_proof_targets(&self) -> MultiProofTargets {
133        // Pre-allocate minimum capacity for the targets.
134        let mut targets = MultiProofTargets::with_capacity(self.accounts.len());
135        for hashed_address in self.accounts.keys() {
136            targets.insert(*hashed_address, Default::default());
137        }
138        for (hashed_address, storage) in &self.storages {
139            targets.entry(*hashed_address).or_default().extend(storage.storage.keys().copied());
140        }
141        targets
142    }
143
144    /// Create multiproof targets difference for this state,
145    /// i.e., the targets that are in targets create from `self` but not in `excluded`.
146    ///
147    /// This method is preferred to first calling `Self::multi_proof_targets` and the calling
148    /// `MultiProofTargets::retain_difference`, because it does not over allocate the targets map.
149    pub fn multi_proof_targets_difference(
150        &self,
151        excluded: &MultiProofTargets,
152    ) -> MultiProofTargets {
153        let mut targets = MultiProofTargets::default();
154        for hashed_address in self.accounts.keys() {
155            if !excluded.contains_key(hashed_address) {
156                targets.insert(*hashed_address, Default::default());
157            }
158        }
159        for (hashed_address, storage) in &self.storages {
160            let maybe_excluded_storage = excluded.get(hashed_address);
161            let mut hashed_slots_targets = storage
162                .storage
163                .keys()
164                .filter(|slot| !maybe_excluded_storage.is_some_and(|f| f.contains(*slot)))
165                .peekable();
166            if hashed_slots_targets.peek().is_some() {
167                targets.entry(*hashed_address).or_default().extend(hashed_slots_targets);
168            }
169        }
170        targets
171    }
172
173    /// Partition the state update into two state updates:
174    /// - First with accounts and storages slots that are present in the provided targets.
175    /// - Second with all other.
176    ///
177    /// CAUTION: The state updates are expected to be applied in order, so that the storage wipes
178    /// are done correctly.
179    pub fn partition_by_targets(
180        mut self,
181        targets: &MultiProofTargets,
182        added_removed_keys: &MultiAddedRemovedKeys,
183    ) -> (Self, Self) {
184        let mut state_updates_not_in_targets = Self::default();
185
186        self.storages.retain(|&address, storage| {
187            let storage_added_removed_keys = added_removed_keys.get_storage(&address);
188
189            let (retain, storage_not_in_targets) = match targets.get(&address) {
190                Some(storage_in_targets) => {
191                    let mut storage_not_in_targets = HashedStorage::default();
192                    storage.storage.retain(|&slot, value| {
193                        if storage_in_targets.contains(&slot) &&
194                            !storage_added_removed_keys.is_some_and(|k| k.is_removed(&slot))
195                        {
196                            return true
197                        }
198
199                        storage_not_in_targets.storage.insert(slot, *value);
200                        false
201                    });
202
203                    // We do not check the wiped flag here, because targets only contain addresses
204                    // and storage slots. So if there are no storage slots left, the storage update
205                    // can be fully removed.
206                    let retain = !storage.storage.is_empty();
207
208                    // Since state updates are expected to be applied in order, we can only set the
209                    // wiped flag in the second storage update if the first storage update is empty
210                    // and will not be retained.
211                    if !retain {
212                        storage_not_in_targets.wiped = storage.wiped;
213                    }
214
215                    (
216                        retain,
217                        storage_not_in_targets.is_empty().not().then_some(storage_not_in_targets),
218                    )
219                }
220                None => (false, Some(core::mem::take(storage))),
221            };
222
223            if let Some(storage_not_in_targets) = storage_not_in_targets {
224                state_updates_not_in_targets.storages.insert(address, storage_not_in_targets);
225            }
226
227            retain
228        });
229        self.accounts.retain(|&address, account| {
230            if targets.contains_key(&address) {
231                return true
232            }
233
234            state_updates_not_in_targets.accounts.insert(address, *account);
235            false
236        });
237
238        (self, state_updates_not_in_targets)
239    }
240
241    /// Returns an iterator that yields chunks of the specified size.
242    ///
243    /// See [`ChunkedHashedPostState`] for more information.
244    pub fn chunks(self, size: usize) -> ChunkedHashedPostState {
245        ChunkedHashedPostState::new(self, size)
246    }
247
248    /// Returns the number of items that will be considered during chunking in `[Self::chunks]`.
249    pub fn chunking_length(&self) -> usize {
250        self.accounts.len() +
251            self.storages
252                .values()
253                .map(|storage| if storage.wiped { 1 } else { 0 } + storage.storage.len())
254                .sum::<usize>()
255    }
256
257    /// Extend this hashed post state with contents of another.
258    /// Entries in the second hashed post state take precedence.
259    pub fn extend(&mut self, other: Self) {
260        self.extend_inner(Cow::Owned(other));
261    }
262
263    /// Extend this hashed post state with contents of another.
264    /// Entries in the second hashed post state take precedence.
265    ///
266    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
267    pub fn extend_ref(&mut self, other: &Self) {
268        self.extend_inner(Cow::Borrowed(other));
269    }
270
271    fn extend_inner(&mut self, other: Cow<'_, Self>) {
272        self.accounts.extend(other.accounts.iter().map(|(&k, &v)| (k, v)));
273
274        self.storages.reserve(other.storages.len());
275        match other {
276            Cow::Borrowed(other) => {
277                self.extend_storages(other.storages.iter().map(|(k, v)| (*k, Cow::Borrowed(v))))
278            }
279            Cow::Owned(other) => {
280                self.extend_storages(other.storages.into_iter().map(|(k, v)| (k, Cow::Owned(v))))
281            }
282        }
283    }
284
285    fn extend_storages<'a>(
286        &mut self,
287        storages: impl IntoIterator<Item = (B256, Cow<'a, HashedStorage>)>,
288    ) {
289        for (hashed_address, storage) in storages {
290            match self.storages.entry(hashed_address) {
291                hash_map::Entry::Vacant(entry) => {
292                    entry.insert(storage.into_owned());
293                }
294                hash_map::Entry::Occupied(mut entry) => {
295                    entry.get_mut().extend(&storage);
296                }
297            }
298        }
299    }
300
301    /// Extend this hashed post state with sorted data, converting directly into the unsorted
302    /// `HashMap` representation. This is more efficient than first converting to `HashedPostState`
303    /// and then extending, as it avoids creating intermediate `HashMap` allocations.
304    pub fn extend_from_sorted(&mut self, sorted: &HashedPostStateSorted) {
305        // Reserve capacity for accounts
306        self.accounts.reserve(sorted.accounts.len());
307
308        // Insert accounts (Some = updated, None = destroyed)
309        for (address, account) in &sorted.accounts {
310            self.accounts.insert(*address, *account);
311        }
312
313        // Reserve capacity for storages
314        self.storages.reserve(sorted.storages.len());
315
316        // Extend storages
317        for (hashed_address, sorted_storage) in &sorted.storages {
318            match self.storages.entry(*hashed_address) {
319                hash_map::Entry::Vacant(entry) => {
320                    let mut new_storage = HashedStorage::new(false);
321                    new_storage.extend_from_sorted(sorted_storage);
322                    entry.insert(new_storage);
323                }
324                hash_map::Entry::Occupied(mut entry) => {
325                    entry.get_mut().extend_from_sorted(sorted_storage);
326                }
327            }
328        }
329    }
330
331    /// Converts hashed post state into [`HashedPostStateSorted`].
332    pub fn into_sorted(self) -> HashedPostStateSorted {
333        let mut accounts: Vec<_> = self.accounts.into_iter().collect();
334        accounts.sort_unstable_by_key(|(address, _)| *address);
335
336        let storages = self
337            .storages
338            .into_iter()
339            .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
340            .collect();
341
342        HashedPostStateSorted { accounts, storages }
343    }
344
345    /// Creates a sorted copy without consuming self.
346    /// More efficient than `.clone().into_sorted()` as it avoids cloning `HashMap` metadata.
347    pub fn clone_into_sorted(&self) -> HashedPostStateSorted {
348        let mut accounts: Vec<_> = self.accounts.iter().map(|(&k, &v)| (k, v)).collect();
349        accounts.sort_unstable_by_key(|(address, _)| *address);
350
351        let storages = self
352            .storages
353            .iter()
354            .map(|(&hashed_address, storage)| (hashed_address, storage.clone_into_sorted()))
355            .collect();
356
357        HashedPostStateSorted { accounts, storages }
358    }
359
360    /// Clears the account and storage maps of this `HashedPostState`.
361    pub fn clear(&mut self) {
362        self.accounts.clear();
363        self.storages.clear();
364    }
365}
366
367impl FromIterator<(B256, Option<Account>, Option<HashedStorage>)> for HashedPostState {
368    /// Constructs a [`HashedPostState`] from an iterator of tuples containing:
369    /// - Hashed address (B256)
370    /// - Optional account info (`None` indicates destroyed account)
371    /// - Optional hashed storage
372    ///
373    /// # Important
374    ///
375    /// - The iterator **assumes unique hashed addresses** (B256). If duplicate addresses are
376    ///   present, later entries will overwrite earlier ones for accounts, and storage will be
377    ///   merged.
378    /// - The [`HashedStorage`] **must not be empty** (as determined by
379    ///   [`HashedStorage::is_empty`]). Empty storage should be represented as `None` rather than
380    ///   `Some(empty_storage)`. This ensures the storage map only contains meaningful entries.
381    ///
382    /// Use `(!storage.is_empty()).then_some(storage)` to convert empty storage to `None`.
383    fn from_iter<T: IntoIterator<Item = (B256, Option<Account>, Option<HashedStorage>)>>(
384        iter: T,
385    ) -> Self {
386        let iter = iter.into_iter();
387        let (lower, _) = iter.size_hint();
388        let mut hashed_state = Self::with_capacity(lower);
389
390        for (hashed_address, info, hashed_storage) in iter {
391            hashed_state.accounts.insert(hashed_address, info);
392            if let Some(storage) = hashed_storage {
393                hashed_state.storages.insert(hashed_address, storage);
394            }
395        }
396
397        hashed_state
398    }
399}
400
401#[cfg(feature = "rayon")]
402impl FromParallelIterator<(B256, Option<Account>, Option<HashedStorage>)> for HashedPostState {
403    /// Parallel version of [`FromIterator`] for constructing [`HashedPostState`] from a parallel
404    /// iterator.
405    ///
406    /// See [`FromIterator::from_iter`] for details on the expected input format.
407    ///
408    /// # Important
409    ///
410    /// - The iterator **assumes unique hashed addresses** (B256). If duplicate addresses are
411    ///   present, later entries will overwrite earlier ones for accounts, and storage will be
412    ///   merged.
413    /// - The [`HashedStorage`] **must not be empty**. Empty storage should be `None`.
414    fn from_par_iter<I>(par_iter: I) -> Self
415    where
416        I: IntoParallelIterator<Item = (B256, Option<Account>, Option<HashedStorage>)>,
417    {
418        let vec: Vec<_> = par_iter.into_par_iter().collect();
419        vec.into_iter().collect()
420    }
421}
422
423/// Representation of in-memory hashed storage.
424#[derive(PartialEq, Eq, Clone, Debug, Default)]
425#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
426pub struct HashedStorage {
427    /// Flag indicating whether the storage was wiped or not.
428    pub wiped: bool,
429    /// Mapping of hashed storage slot to storage value.
430    pub storage: B256Map<U256>,
431}
432
433impl HashedStorage {
434    /// Create new instance of [`HashedStorage`].
435    pub fn new(wiped: bool) -> Self {
436        Self { wiped, storage: HashMap::default() }
437    }
438
439    /// Check if self is empty.
440    pub fn is_empty(&self) -> bool {
441        !self.wiped && self.storage.is_empty()
442    }
443
444    /// Create new hashed storage from iterator.
445    pub fn from_iter(wiped: bool, iter: impl IntoIterator<Item = (B256, U256)>) -> Self {
446        Self { wiped, storage: HashMap::from_iter(iter) }
447    }
448
449    /// Create new hashed storage from account status and plain storage.
450    pub fn from_plain_storage<'a>(
451        status: AccountStatus,
452        storage: impl IntoIterator<Item = (&'a U256, &'a U256)>,
453    ) -> Self {
454        Self::from_iter(
455            status.was_destroyed(),
456            storage.into_iter().map(|(key, value)| (keccak256(B256::from(*key)), *value)),
457        )
458    }
459
460    /// Construct [`PrefixSetMut`] from hashed storage.
461    pub fn construct_prefix_set(&self) -> PrefixSetMut {
462        if self.wiped {
463            PrefixSetMut::all()
464        } else {
465            let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
466            for hashed_slot in self.storage.keys() {
467                prefix_set.insert(Nibbles::unpack(hashed_slot));
468            }
469            prefix_set
470        }
471    }
472
473    /// Extend hashed storage with contents of other.
474    /// The entries in second hashed storage take precedence.
475    pub fn extend(&mut self, other: &Self) {
476        if other.wiped {
477            self.wiped = true;
478            self.storage.clear();
479        }
480        self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v)));
481    }
482
483    /// Extend hashed storage with sorted data, converting directly into the unsorted `HashMap`
484    /// representation. This is more efficient than first converting to `HashedStorage` and
485    /// then extending, as it avoids creating intermediate `HashMap` allocations.
486    pub fn extend_from_sorted(&mut self, sorted: &HashedStorageSorted) {
487        if sorted.wiped {
488            self.wiped = true;
489            self.storage.clear();
490        }
491
492        // Reserve capacity for all slots
493        self.storage.reserve(sorted.storage_slots.len());
494
495        // Insert all storage slots
496        for (slot, value) in &sorted.storage_slots {
497            self.storage.insert(*slot, *value);
498        }
499    }
500
501    /// Converts hashed storage into [`HashedStorageSorted`].
502    pub fn into_sorted(self) -> HashedStorageSorted {
503        let mut storage_slots: Vec<_> = self.storage.into_iter().collect();
504        storage_slots.sort_unstable_by_key(|(key, _)| *key);
505
506        HashedStorageSorted { storage_slots, wiped: self.wiped }
507    }
508
509    /// Creates a sorted copy without consuming self.
510    /// More efficient than `.clone().into_sorted()` as it avoids cloning `HashMap` metadata.
511    pub fn clone_into_sorted(&self) -> HashedStorageSorted {
512        let mut storage_slots: Vec<_> = self.storage.iter().map(|(&k, &v)| (k, v)).collect();
513        storage_slots.sort_unstable_by_key(|(key, _)| *key);
514
515        HashedStorageSorted { storage_slots, wiped: self.wiped }
516    }
517}
518
519/// Sorted hashed post state optimized for iterating during state trie calculation.
520#[derive(PartialEq, Eq, Clone, Default, Debug)]
521#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
522pub struct HashedPostStateSorted {
523    /// Sorted collection of account updates. `None` indicates a destroyed account.
524    pub accounts: Vec<(B256, Option<Account>)>,
525    /// Map of hashed addresses to their sorted storage updates.
526    pub storages: B256Map<HashedStorageSorted>,
527}
528
529impl HashedPostStateSorted {
530    /// Create new instance of [`HashedPostStateSorted`]
531    pub const fn new(
532        accounts: Vec<(B256, Option<Account>)>,
533        storages: B256Map<HashedStorageSorted>,
534    ) -> Self {
535        Self { accounts, storages }
536    }
537
538    /// Returns reference to hashed accounts.
539    pub const fn accounts(&self) -> &Vec<(B256, Option<Account>)> {
540        &self.accounts
541    }
542
543    /// Returns reference to hashed account storages.
544    pub const fn account_storages(&self) -> &B256Map<HashedStorageSorted> {
545        &self.storages
546    }
547
548    /// Returns `true` if there are no account or storage updates.
549    pub fn is_empty(&self) -> bool {
550        self.accounts.is_empty() && self.storages.is_empty()
551    }
552
553    /// Returns the total number of updates including all accounts and storage updates.
554    pub fn total_len(&self) -> usize {
555        self.accounts.len() + self.storages.values().map(|s| s.len()).sum::<usize>()
556    }
557
558    /// Construct [`TriePrefixSetsMut`] from hashed post state.
559    ///
560    /// The prefix sets contain the hashed account and storage keys that have been changed in the
561    /// post state.
562    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
563        let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
564        let mut destroyed_accounts = HashSet::default();
565        for (hashed_address, account) in &self.accounts {
566            account_prefix_set.insert(Nibbles::unpack(hashed_address));
567            if account.is_none() {
568                destroyed_accounts.insert(*hashed_address);
569            }
570        }
571
572        let mut storage_prefix_sets =
573            B256Map::with_capacity_and_hasher(self.storages.len(), Default::default());
574        for (hashed_address, hashed_storage) in &self.storages {
575            // Ensure account trie covers storage overlays even if account map is empty.
576            account_prefix_set.insert(Nibbles::unpack(hashed_address));
577
578            let prefix_set = if hashed_storage.wiped {
579                PrefixSetMut::all()
580            } else {
581                let mut prefix_set =
582                    PrefixSetMut::with_capacity(hashed_storage.storage_slots.len());
583                prefix_set.extend_keys(
584                    hashed_storage
585                        .storage_slots
586                        .iter()
587                        .map(|(hashed_slot, _)| Nibbles::unpack(hashed_slot)),
588                );
589                prefix_set
590            };
591
592            storage_prefix_sets.insert(*hashed_address, prefix_set);
593        }
594
595        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
596    }
597
598    /// Extends this state with contents of another sorted state.
599    /// Entries in `other` take precedence for duplicate keys.
600    ///
601    /// Sorts the accounts after extending. Sorts the storage after extending, for each account.
602    pub fn extend_ref_and_sort(&mut self, other: &Self) {
603        // Extend accounts
604        extend_sorted_vec(&mut self.accounts, &other.accounts);
605
606        // Extend storages
607        for (hashed_address, other_storage) in &other.storages {
608            self.storages
609                .entry(*hashed_address)
610                .and_modify(|existing| existing.extend_ref(other_storage))
611                .or_insert_with(|| other_storage.clone());
612        }
613    }
614
615    /// Batch-merge sorted hashed post states. Iterator yields **newest to oldest**.
616    ///
617    /// For small batches, uses `extend_ref_and_sort` loop.
618    /// For large batches, uses k-way merge for O(n log k) complexity.
619    pub fn merge_batch<T: AsRef<Self> + From<Self>>(iter: impl IntoIterator<Item = T>) -> T {
620        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
621        match items.len() {
622            0 => Self::default().into(),
623            1 => items.into_iter().next().expect("len == 1"),
624            _ => Self::merge_slice(&items).into(),
625        }
626    }
627
628    /// Batch-merge sorted hashed post states from a slice. Slice is **newest to oldest**.
629    ///
630    /// This variant takes a slice reference directly, avoiding iterator collection overhead.
631    /// For small batches, uses `extend_ref_and_sort` loop.
632    /// For large batches, uses k-way merge for O(n log k) complexity.
633    pub fn merge_slice<T: AsRef<Self>>(items: &[T]) -> Self {
634        const THRESHOLD: usize = 30;
635
636        let k = items.len();
637
638        if k == 0 {
639            return Self::default();
640        }
641        if k == 1 {
642            return items[0].as_ref().clone();
643        }
644
645        if k < THRESHOLD {
646            // Small k: extend loop, oldest-to-newest so newer overrides older.
647            let mut iter = items.iter().rev();
648            let mut acc = iter.next().expect("k > 0").as_ref().clone();
649            for next in iter {
650                acc.extend_ref_and_sort(next.as_ref());
651            }
652            return acc;
653        }
654
655        // Large k: k-way merge.
656        let accounts = kway_merge_sorted(items.iter().map(|i| i.as_ref().accounts.as_slice()));
657
658        struct StorageAcc<'a> {
659            wiped: bool,
660            sealed: bool,
661            slices: Vec<&'a [(B256, U256)]>,
662        }
663
664        let mut acc: B256Map<StorageAcc<'_>> = B256Map::default();
665
666        for item in items {
667            for (addr, storage) in &item.as_ref().storages {
668                let entry = acc.entry(*addr).or_insert_with(|| StorageAcc {
669                    wiped: false,
670                    sealed: false,
671                    slices: Vec::new(),
672                });
673
674                if entry.sealed {
675                    continue;
676                }
677
678                entry.slices.push(storage.storage_slots.as_slice());
679                if storage.wiped {
680                    entry.wiped = true;
681                    entry.sealed = true;
682                }
683            }
684        }
685
686        let storages = acc
687            .into_iter()
688            .map(|(addr, entry)| {
689                let storage_slots = kway_merge_sorted(entry.slices);
690                (addr, HashedStorageSorted { wiped: entry.wiped, storage_slots })
691            })
692            .collect();
693
694        Self { accounts, storages }
695    }
696
697    /// Merges the batch and removes overlapping keys whose mask values all differ from the merged
698    /// batch value.
699    ///
700    /// Account keys are masked at the top level, while storage entries are masked at the slot
701    /// level. For duplicate keys in the batch, later items take precedence over earlier ones. An
702    /// overlapping entry is retained if any mask value is equal to the merged batch value. The
703    /// order of the mask does not matter. An empty mask merges the batch without filtering.
704    ///
705    /// # Panics
706    ///
707    /// Panics if any batch or mask entry wipes an entire storage.
708    pub fn disjointed_merge_batch<'a>(batch: &[&'a Self], mask: &[&'a Self]) -> Self {
709        let account_count = batch.iter().map(|item| item.accounts.len()).sum();
710        let mut accounts = Vec::with_capacity(account_count);
711        accounts.extend(kway_merge_disjoint_sorted(
712            batch.iter().rev().map(|item| item.accounts.as_slice()),
713            mask.iter().map(|item| item.accounts.as_slice()),
714        ));
715
716        struct StorageAcc<'a> {
717            slot_count: usize,
718            slices: Vec<&'a [(B256, U256)]>,
719        }
720
721        #[derive(Default)]
722        struct StorageMaskAcc<'a> {
723            slices: Vec<&'a [(B256, U256)]>,
724        }
725
726        let mut storages = B256Map::with_capacity_and_hasher(
727            batch.iter().map(|item| item.storages.len()).sum(),
728            Default::default(),
729        );
730
731        for item in batch.iter().rev() {
732            for (hashed_address, storage) in &item.storages {
733                assert!(
734                    !storage.wiped,
735                    "storage wipes are not supported by disjointed_merge_batch"
736                );
737                let entry = storages
738                    .entry(*hashed_address)
739                    .or_insert_with(|| StorageAcc { slot_count: 0, slices: Vec::new() });
740                entry.slices.push(storage.storage_slots.as_slice());
741                entry.slot_count += storage.storage_slots.len();
742            }
743        }
744
745        let mut storage_masks: B256Map<StorageMaskAcc<'a>> = B256Map::with_capacity_and_hasher(
746            mask.iter().map(|item| item.storages.len()).sum(),
747            Default::default(),
748        );
749        for item in mask {
750            for (hashed_address, storage) in &item.storages {
751                assert!(
752                    !storage.wiped,
753                    "storage wipes are not supported by disjointed_merge_batch"
754                );
755                let entry = storage_masks.entry(*hashed_address).or_default();
756                entry.slices.push(storage.storage_slots.as_slice());
757            }
758        }
759
760        let storages = storages
761            .into_iter()
762            .filter_map(|(hashed_address, entry)| {
763                let slot_count = entry.slot_count;
764                let storage_slots = match storage_masks.get(&hashed_address) {
765                    Some(mask_entry) => {
766                        let mut storage_slots = Vec::with_capacity(slot_count);
767                        storage_slots.extend(kway_merge_disjoint_sorted(
768                            entry.slices,
769                            mask_entry.slices.iter().copied(),
770                        ));
771                        storage_slots
772                    }
773                    None => kway_merge_sorted(entry.slices),
774                };
775
776                (!storage_slots.is_empty() || mask.is_empty()).then_some((
777                    hashed_address,
778                    HashedStorageSorted { wiped: false, storage_slots },
779                ))
780            })
781            .collect();
782
783        Self { accounts, storages }
784    }
785
786    /// Clears all accounts and storage data.
787    pub fn clear(&mut self) {
788        self.accounts.clear();
789        self.storages.clear();
790    }
791}
792
793impl AsRef<Self> for HashedPostStateSorted {
794    fn as_ref(&self) -> &Self {
795        self
796    }
797}
798
799/// Sorted hashed storage optimized for iterating during state trie calculation.
800#[derive(Clone, Eq, PartialEq, Debug, Default)]
801#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
802pub struct HashedStorageSorted {
803    /// Sorted collection of updated storage slots. [`U256::ZERO`] indicates a deleted value.
804    pub storage_slots: Vec<(B256, U256)>,
805    /// Flag indicating whether the storage was wiped or not.
806    pub wiped: bool,
807}
808
809impl HashedStorageSorted {
810    /// Returns `true` if the account was wiped.
811    pub const fn is_wiped(&self) -> bool {
812        self.wiped
813    }
814
815    /// Returns reference to updated storage slots.
816    pub fn storage_slots_ref(&self) -> &[(B256, U256)] {
817        &self.storage_slots
818    }
819
820    /// Returns the total number of storage slot updates.
821    pub const fn len(&self) -> usize {
822        self.storage_slots.len()
823    }
824
825    /// Returns `true` if there are no storage slot updates.
826    pub const fn is_empty(&self) -> bool {
827        self.storage_slots.is_empty()
828    }
829
830    /// Extends the storage slots updates with another set of sorted updates.
831    ///
832    /// If `other` is marked as deleted, this will be marked as deleted and all slots cleared.
833    /// Otherwise, nodes are merged with `other`'s values taking precedence for duplicates.
834    pub fn extend_ref(&mut self, other: &Self) {
835        if other.wiped {
836            // If other is wiped, clear everything and copy from other
837            self.wiped = true;
838            self.storage_slots.clear();
839            self.storage_slots.extend(other.storage_slots.iter().copied());
840            return;
841        }
842
843        // Extend the sorted non-zero valued slots
844        extend_sorted_vec(&mut self.storage_slots, &other.storage_slots);
845    }
846
847    /// Batch-merge sorted hashed storage. Iterator yields **newest to oldest**.
848    /// If any update is wiped, prior data is discarded.
849    pub fn merge_batch<'a>(updates: impl IntoIterator<Item = &'a Self>) -> Self {
850        let updates: Vec<_> = updates.into_iter().collect();
851        if updates.is_empty() {
852            return Self::default();
853        }
854
855        let wipe_idx = updates.iter().position(|u| u.wiped);
856        let relevant = wipe_idx.map_or(&updates[..], |idx| &updates[..=idx]);
857        let storage_slots = kway_merge_sorted(relevant.iter().map(|u| u.storage_slots.as_slice()));
858
859        Self { wiped: wipe_idx.is_some(), storage_slots }
860    }
861}
862
863impl From<HashedStorageSorted> for HashedStorage {
864    fn from(sorted: HashedStorageSorted) -> Self {
865        let mut storage = B256Map::default();
866
867        // Add all storage slots (including zero-valued ones which indicate deletion)
868        for (slot, value) in sorted.storage_slots {
869            storage.insert(slot, value);
870        }
871
872        Self { wiped: sorted.wiped, storage }
873    }
874}
875
876impl From<HashedPostStateSorted> for HashedPostState {
877    fn from(sorted: HashedPostStateSorted) -> Self {
878        let mut accounts =
879            B256Map::with_capacity_and_hasher(sorted.accounts.len(), Default::default());
880
881        // Add all accounts (Some for updated, None for destroyed)
882        for (address, account) in sorted.accounts {
883            accounts.insert(address, account);
884        }
885
886        // Convert storages
887        let storages = sorted
888            .storages
889            .into_iter()
890            .map(|(address, storage)| (address, storage.into()))
891            .collect();
892
893        Self { accounts, storages }
894    }
895}
896
897/// An iterator that yields chunks of the state updates of at most `size` account and storage
898/// targets.
899///
900/// # Notes
901/// 1. Chunks are expected to be applied in order, because of storage wipes. If applied out of
902///    order, it's possible to wipe more storage than in the original state update.
903/// 2. For each account, chunks with storage updates come first, followed by account updates.
904#[derive(Debug)]
905pub struct ChunkedHashedPostState {
906    flattened: alloc::vec::IntoIter<(B256, FlattenedHashedPostStateItem)>,
907    size: usize,
908}
909
910/// Order discriminant for sorting flattened state items.
911/// Ordering: `StorageWipe` < `StorageUpdate` (by slot) < `Account`
912#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
913enum FlattenedStateOrder {
914    StorageWipe,
915    StorageUpdate(B256),
916    Account,
917}
918
919#[derive(Debug)]
920enum FlattenedHashedPostStateItem {
921    Account(Option<Account>),
922    StorageWipe,
923    StorageUpdate { slot: B256, value: U256 },
924}
925
926impl FlattenedHashedPostStateItem {
927    const fn order(&self) -> FlattenedStateOrder {
928        match self {
929            Self::StorageWipe => FlattenedStateOrder::StorageWipe,
930            Self::StorageUpdate { slot, .. } => FlattenedStateOrder::StorageUpdate(*slot),
931            Self::Account(_) => FlattenedStateOrder::Account,
932        }
933    }
934}
935
936impl ChunkedHashedPostState {
937    fn new(hashed_post_state: HashedPostState, size: usize) -> Self {
938        let flattened = hashed_post_state
939            .storages
940            .into_iter()
941            .flat_map(|(address, storage)| {
942                storage
943                    .wiped
944                    .then_some((address, FlattenedHashedPostStateItem::StorageWipe))
945                    .into_iter()
946                    .chain(storage.storage.into_iter().map(move |(slot, value)| {
947                        (address, FlattenedHashedPostStateItem::StorageUpdate { slot, value })
948                    }))
949            })
950            .chain(hashed_post_state.accounts.into_iter().map(|(address, account)| {
951                (address, FlattenedHashedPostStateItem::Account(account))
952            }))
953            // Sort by address, then by item order to ensure correct application sequence:
954            // 1. Storage wipes (must come first to clear storage)
955            // 2. Storage updates (sorted by slot for determinism)
956            // 3. Account updates (can be applied last)
957            .sorted_unstable_by_key(|(address, item)| (*address, item.order()));
958
959        Self { flattened, size }
960    }
961}
962
963impl Iterator for ChunkedHashedPostState {
964    type Item = HashedPostState;
965
966    fn next(&mut self) -> Option<Self::Item> {
967        let mut chunk = HashedPostState::default();
968
969        let mut current_size = 0;
970        while current_size < self.size {
971            let Some((address, item)) = self.flattened.next() else { break };
972
973            match item {
974                FlattenedHashedPostStateItem::Account(account) => {
975                    chunk.accounts.insert(address, account);
976                }
977                FlattenedHashedPostStateItem::StorageWipe => {
978                    chunk.storages.entry(address).or_default().wiped = true;
979                }
980                FlattenedHashedPostStateItem::StorageUpdate { slot, value } => {
981                    chunk.storages.entry(address).or_default().storage.insert(slot, value);
982                }
983            }
984
985            current_size += 1;
986        }
987
988        if chunk.is_empty() {
989            None
990        } else {
991            Some(chunk)
992        }
993    }
994}
995
996#[cfg(test)]
997mod tests {
998    use super::*;
999    use crate::KeccakKeyHasher;
1000    use alloy_primitives::Bytes;
1001    use revm::{
1002        database::{states::StorageSlot, StorageWithOriginalValues},
1003        state::{AccountInfo, Bytecode},
1004    };
1005
1006    fn bundle_hashed_storage(account: &BundleAccount) -> Option<HashedStorage> {
1007        let address = Address::ZERO;
1008        let mut state =
1009            HashedPostState::from_bundle_state::<KeccakKeyHasher>([(&address, account)]);
1010        state.storages.remove(&keccak256(address))
1011    }
1012
1013    fn changed_storage(original: U256, present: U256) -> StorageWithOriginalValues {
1014        core::iter::once((U256::from(1), StorageSlot::new_changed(original, present))).collect()
1015    }
1016
1017    #[test]
1018    fn hashed_state_wiped_extension() {
1019        let hashed_address = B256::default();
1020        let hashed_slot = B256::with_last_byte(64);
1021        let hashed_slot2 = B256::with_last_byte(65);
1022
1023        // Initialize post state storage
1024        let original_slot_value = U256::from(123);
1025        let mut hashed_state = HashedPostState::default().with_storages([(
1026            hashed_address,
1027            HashedStorage::from_iter(
1028                false,
1029                [(hashed_slot, original_slot_value), (hashed_slot2, original_slot_value)],
1030            ),
1031        )]);
1032
1033        // Update single slot value
1034        let updated_slot_value = U256::from(321);
1035        let extension = HashedPostState::default().with_storages([(
1036            hashed_address,
1037            HashedStorage::from_iter(false, [(hashed_slot, updated_slot_value)]),
1038        )]);
1039        hashed_state.extend(extension);
1040
1041        let account_storage = hashed_state.storages.get(&hashed_address);
1042        assert_eq!(
1043            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
1044            Some(&updated_slot_value)
1045        );
1046        assert_eq!(
1047            account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
1048            Some(&original_slot_value)
1049        );
1050        assert_eq!(account_storage.map(|st| st.wiped), Some(false));
1051
1052        // Wipe account storage
1053        let wiped_extension =
1054            HashedPostState::default().with_storages([(hashed_address, HashedStorage::new(true))]);
1055        hashed_state.extend(wiped_extension);
1056
1057        let account_storage = hashed_state.storages.get(&hashed_address);
1058        assert_eq!(account_storage.map(|st| st.storage.is_empty()), Some(true));
1059        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
1060
1061        // Reinitialize single slot value
1062        hashed_state.extend(HashedPostState::default().with_storages([(
1063            hashed_address,
1064            HashedStorage::from_iter(false, [(hashed_slot, original_slot_value)]),
1065        )]));
1066        let account_storage = hashed_state.storages.get(&hashed_address);
1067        assert_eq!(
1068            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
1069            Some(&original_slot_value)
1070        );
1071        assert_eq!(account_storage.and_then(|st| st.storage.get(&hashed_slot2)), None);
1072        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
1073
1074        // Reinitialize single slot value
1075        hashed_state.extend(HashedPostState::default().with_storages([(
1076            hashed_address,
1077            HashedStorage::from_iter(false, [(hashed_slot2, updated_slot_value)]),
1078        )]));
1079        let account_storage = hashed_state.storages.get(&hashed_address);
1080        assert_eq!(
1081            account_storage.and_then(|st| st.storage.get(&hashed_slot)),
1082            Some(&original_slot_value)
1083        );
1084        assert_eq!(
1085            account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
1086            Some(&updated_slot_value)
1087        );
1088        assert_eq!(account_storage.map(|st| st.wiped), Some(true));
1089    }
1090
1091    #[test]
1092    fn test_hashed_post_state_from_bundle_state() {
1093        // Prepare a random Ethereum address as a key for the account.
1094        let address = Address::random();
1095
1096        // Create a mock account info object.
1097        let account_info = AccountInfo {
1098            balance: U256::from(123),
1099            nonce: 42,
1100            code_hash: B256::random(),
1101            code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2]))),
1102            account_id: None,
1103        };
1104
1105        let mut storage = StorageWithOriginalValues::default();
1106        storage.insert(
1107            U256::from(1),
1108            StorageSlot { present_value: U256::from(4), ..Default::default() },
1109        );
1110
1111        // Create a `BundleAccount` struct to represent the account and its storage.
1112        let account = BundleAccount {
1113            status: AccountStatus::Changed,
1114            info: Some(account_info.clone()),
1115            storage,
1116            original_info: None,
1117        };
1118
1119        // Create a vector of tuples representing the bundle state.
1120        let state = vec![(&address, &account)];
1121
1122        // Convert the bundle state into a hashed post state.
1123        let hashed_state = HashedPostState::from_bundle_state::<KeccakKeyHasher>(state);
1124
1125        // Validate the hashed post state.
1126        assert_eq!(hashed_state.accounts.len(), 1);
1127        assert_eq!(hashed_state.storages.len(), 1);
1128
1129        // Validate the account info.
1130        assert_eq!(
1131            *hashed_state.accounts.get(&keccak256(address)).unwrap(),
1132            Some(account_info.into())
1133        );
1134    }
1135
1136    #[test]
1137    fn destroyed_prefunded_account_without_storage_emits_no_storage() {
1138        let original_info = AccountInfo { balance: U256::from(1), ..Default::default() };
1139        let account = BundleAccount::new(
1140            Some(original_info),
1141            None,
1142            StorageWithOriginalValues::default(),
1143            AccountStatus::Destroyed,
1144        );
1145
1146        assert!(bundle_hashed_storage(&account).is_none());
1147    }
1148
1149    #[test]
1150    fn destroyed_accounts_emit_zero_storage_changes_without_wipe() {
1151        let existing_contract =
1152            AccountInfo { code_hash: B256::repeat_byte(0x01), ..Default::default() };
1153        let legacy_empty_account = AccountInfo::default();
1154        let prefunded_account = AccountInfo { balance: U256::from(1), ..Default::default() };
1155
1156        for original_info in
1157            [Some(existing_contract), Some(legacy_empty_account), Some(prefunded_account), None]
1158        {
1159            let account = BundleAccount::new(
1160                original_info,
1161                None,
1162                changed_storage(U256::from(2), U256::ZERO),
1163                AccountStatus::Destroyed,
1164            );
1165
1166            let storage = bundle_hashed_storage(&account).unwrap();
1167            let hashed_slot = keccak256(B256::from(U256::from(1)));
1168            assert!(!storage.wiped);
1169            assert_eq!(storage.storage[&hashed_slot], U256::ZERO);
1170        }
1171    }
1172
1173    #[test]
1174    fn destroyed_recreated_accounts_preserve_storage_without_wipe() {
1175        let value = U256::from(2);
1176        let new_account = BundleAccount::new(
1177            None,
1178            Some(AccountInfo::default()),
1179            changed_storage(U256::ZERO, value),
1180            AccountStatus::DestroyedChanged,
1181        );
1182        let original_info =
1183            AccountInfo { code_hash: B256::repeat_byte(0x01), ..Default::default() };
1184        let existing_account = BundleAccount::new(
1185            Some(original_info),
1186            Some(AccountInfo::default()),
1187            changed_storage(U256::ZERO, value),
1188            AccountStatus::DestroyedChanged,
1189        );
1190
1191        let new_storage = bundle_hashed_storage(&new_account).unwrap();
1192        let existing_storage = bundle_hashed_storage(&existing_account).unwrap();
1193        let hashed_slot = keccak256(B256::from(U256::from(1)));
1194
1195        assert!(!new_storage.wiped);
1196        assert!(!existing_storage.wiped);
1197        assert_eq!(new_storage.storage[&hashed_slot], value);
1198        assert_eq!(existing_storage.storage[&hashed_slot], value);
1199    }
1200
1201    #[test]
1202    fn test_hashed_post_state_with_accounts() {
1203        // Prepare random addresses and mock account info.
1204        let address_1 = Address::random();
1205        let address_2 = Address::random();
1206
1207        let account_info_1 = AccountInfo {
1208            balance: U256::from(1000),
1209            nonce: 1,
1210            code_hash: B256::random(),
1211            code: None,
1212            account_id: None,
1213        };
1214
1215        // Create hashed accounts with addresses.
1216        let account_1 = (keccak256(address_1), Some(account_info_1.into()));
1217        let account_2 = (keccak256(address_2), None);
1218
1219        // Add accounts to the hashed post state.
1220        let hashed_state = HashedPostState::default().with_accounts(vec![account_1, account_2]);
1221
1222        // Validate the hashed post state.
1223        assert_eq!(hashed_state.accounts.len(), 2);
1224        assert!(hashed_state.accounts.contains_key(&keccak256(address_1)));
1225        assert!(hashed_state.accounts.contains_key(&keccak256(address_2)));
1226    }
1227
1228    #[test]
1229    fn test_hashed_post_state_with_storages() {
1230        // Prepare random addresses and mock storage entries.
1231        let address_1 = Address::random();
1232        let address_2 = Address::random();
1233
1234        let storage_1 = (keccak256(address_1), HashedStorage::new(false));
1235        let storage_2 = (keccak256(address_2), HashedStorage::new(true));
1236
1237        // Add storages to the hashed post state.
1238        let hashed_state = HashedPostState::default().with_storages(vec![storage_1, storage_2]);
1239
1240        // Validate the hashed post state.
1241        assert_eq!(hashed_state.storages.len(), 2);
1242        assert!(hashed_state.storages.contains_key(&keccak256(address_1)));
1243        assert!(hashed_state.storages.contains_key(&keccak256(address_2)));
1244    }
1245
1246    #[test]
1247    fn test_hashed_post_state_is_empty() {
1248        // Create an empty hashed post state and validate it's empty.
1249        let empty_state = HashedPostState::default();
1250        assert!(empty_state.is_empty());
1251
1252        // Add an account and validate the state is no longer empty.
1253        let non_empty_state = HashedPostState::default()
1254            .with_accounts(vec![(keccak256(Address::random()), Some(Account::default()))]);
1255        assert!(!non_empty_state.is_empty());
1256    }
1257
1258    fn create_state_for_multi_proof_targets() -> HashedPostState {
1259        let mut state = HashedPostState::default();
1260
1261        let addr1 = B256::random();
1262        let addr2 = B256::random();
1263        state.accounts.insert(addr1, Some(Default::default()));
1264        state.accounts.insert(addr2, Some(Default::default()));
1265
1266        let mut storage = HashedStorage::default();
1267        let slot1 = B256::random();
1268        let slot2 = B256::random();
1269        storage.storage.insert(slot1, U256::ZERO);
1270        storage.storage.insert(slot2, U256::from(1));
1271        state.storages.insert(addr1, storage);
1272
1273        state
1274    }
1275
1276    #[test]
1277    fn test_multi_proof_targets_difference_empty_state() {
1278        let state = HashedPostState::default();
1279        let excluded = MultiProofTargets::default();
1280
1281        let targets = state.multi_proof_targets_difference(&excluded);
1282        assert!(targets.is_empty());
1283    }
1284
1285    #[test]
1286    fn test_multi_proof_targets_difference_new_account_targets() {
1287        let state = create_state_for_multi_proof_targets();
1288        let excluded = MultiProofTargets::default();
1289
1290        // should return all accounts as targets since excluded is empty
1291        let targets = state.multi_proof_targets_difference(&excluded);
1292        assert_eq!(targets.len(), state.accounts.len());
1293        for addr in state.accounts.keys() {
1294            assert!(targets.contains_key(addr));
1295        }
1296    }
1297
1298    #[test]
1299    fn test_multi_proof_targets_difference_new_storage_targets() {
1300        let state = create_state_for_multi_proof_targets();
1301        let excluded = MultiProofTargets::default();
1302
1303        let targets = state.multi_proof_targets_difference(&excluded);
1304
1305        // verify storage slots are included for accounts with storage
1306        for (addr, storage) in &state.storages {
1307            assert!(targets.contains_key(addr));
1308            let target_slots = &targets[addr];
1309            assert_eq!(target_slots.len(), storage.storage.len());
1310            for slot in storage.storage.keys() {
1311                assert!(target_slots.contains(slot));
1312            }
1313        }
1314    }
1315
1316    #[test]
1317    fn test_multi_proof_targets_difference_filter_excluded_accounts() {
1318        let state = create_state_for_multi_proof_targets();
1319        let mut excluded = MultiProofTargets::default();
1320
1321        // select an account that has no storage updates
1322        let excluded_addr = state
1323            .accounts
1324            .keys()
1325            .find(|&&addr| !state.storages.contains_key(&addr))
1326            .expect("Should have an account without storage");
1327
1328        // mark the account as excluded
1329        excluded.insert(*excluded_addr, HashSet::default());
1330
1331        let targets = state.multi_proof_targets_difference(&excluded);
1332
1333        // should not include the already excluded account since it has no storage updates
1334        assert!(!targets.contains_key(excluded_addr));
1335        // other accounts should still be included
1336        assert_eq!(targets.len(), state.accounts.len() - 1);
1337    }
1338
1339    #[test]
1340    fn test_multi_proof_targets_difference_filter_excluded_storage() {
1341        let state = create_state_for_multi_proof_targets();
1342        let mut excluded = MultiProofTargets::default();
1343
1344        // mark one storage slot as excluded
1345        let (addr, storage) = state.storages.iter().next().unwrap();
1346        let mut excluded_slots = HashSet::default();
1347        let excluded_slot = *storage.storage.keys().next().unwrap();
1348        excluded_slots.insert(excluded_slot);
1349        excluded.insert(*addr, excluded_slots);
1350
1351        let targets = state.multi_proof_targets_difference(&excluded);
1352
1353        // should not include the excluded storage slot
1354        let target_slots = &targets[addr];
1355        assert!(!target_slots.contains(&excluded_slot));
1356        assert_eq!(target_slots.len(), storage.storage.len() - 1);
1357    }
1358
1359    #[test]
1360    fn test_multi_proof_targets_difference_mixed_excluded_state() {
1361        let mut state = HashedPostState::default();
1362        let mut excluded = MultiProofTargets::default();
1363
1364        let addr1 = B256::random();
1365        let addr2 = B256::random();
1366        let slot1 = B256::random();
1367        let slot2 = B256::random();
1368
1369        state.accounts.insert(addr1, Some(Default::default()));
1370        state.accounts.insert(addr2, Some(Default::default()));
1371
1372        let mut storage = HashedStorage::default();
1373        storage.storage.insert(slot1, U256::ZERO);
1374        storage.storage.insert(slot2, U256::from(1));
1375        state.storages.insert(addr1, storage);
1376
1377        let mut excluded_slots = HashSet::default();
1378        excluded_slots.insert(slot1);
1379        excluded.insert(addr1, excluded_slots);
1380
1381        let targets = state.multi_proof_targets_difference(&excluded);
1382
1383        assert!(targets.contains_key(&addr2));
1384        assert!(!targets[&addr1].contains(&slot1));
1385        assert!(targets[&addr1].contains(&slot2));
1386    }
1387
1388    #[test]
1389    fn test_multi_proof_targets_difference_unmodified_account_with_storage() {
1390        let mut state = HashedPostState::default();
1391        let excluded = MultiProofTargets::default();
1392
1393        let addr = B256::random();
1394        let slot1 = B256::random();
1395        let slot2 = B256::random();
1396
1397        // don't add the account to state.accounts (simulating unmodified account)
1398        // but add storage updates for this account
1399        let mut storage = HashedStorage::default();
1400        storage.storage.insert(slot1, U256::from(1));
1401        storage.storage.insert(slot2, U256::from(2));
1402        state.storages.insert(addr, storage);
1403
1404        assert!(!state.accounts.contains_key(&addr));
1405        assert!(!excluded.contains_key(&addr));
1406
1407        let targets = state.multi_proof_targets_difference(&excluded);
1408
1409        // verify that we still get the storage slots for the unmodified account
1410        assert!(targets.contains_key(&addr));
1411
1412        let target_slots = &targets[&addr];
1413        assert_eq!(target_slots.len(), 2);
1414        assert!(target_slots.contains(&slot1));
1415        assert!(target_slots.contains(&slot2));
1416    }
1417
1418    #[test]
1419    fn test_partition_by_targets() {
1420        let addr1 = B256::random();
1421        let addr2 = B256::random();
1422        let slot1 = B256::random();
1423        let slot2 = B256::random();
1424
1425        let state = HashedPostState {
1426            accounts: B256Map::from_iter([
1427                (addr1, Some(Default::default())),
1428                (addr2, Some(Default::default())),
1429            ]),
1430            storages: B256Map::from_iter([(
1431                addr1,
1432                HashedStorage {
1433                    wiped: true,
1434                    storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1435                },
1436            )]),
1437        };
1438        let targets = MultiProofTargets::from_iter([(addr1, HashSet::from_iter([slot1]))]);
1439
1440        let (with_targets, without_targets) =
1441            state.partition_by_targets(&targets, &MultiAddedRemovedKeys::new());
1442
1443        assert_eq!(
1444            with_targets,
1445            HashedPostState {
1446                accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
1447                storages: B256Map::from_iter([(
1448                    addr1,
1449                    HashedStorage {
1450                        wiped: true,
1451                        storage: B256Map::from_iter([(slot1, U256::ZERO)])
1452                    }
1453                )]),
1454            }
1455        );
1456        assert_eq!(
1457            without_targets,
1458            HashedPostState {
1459                accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1460                storages: B256Map::from_iter([(
1461                    addr1,
1462                    HashedStorage {
1463                        wiped: false,
1464                        storage: B256Map::from_iter([(slot2, U256::from(1))])
1465                    }
1466                )]),
1467            }
1468        );
1469    }
1470
1471    #[test]
1472    fn test_chunks() {
1473        let addr1 = B256::from([1; 32]);
1474        let addr2 = B256::from([2; 32]);
1475        let slot1 = B256::from([1; 32]);
1476        let slot2 = B256::from([2; 32]);
1477
1478        let state = HashedPostState {
1479            accounts: B256Map::from_iter([
1480                (addr1, Some(Default::default())),
1481                (addr2, Some(Default::default())),
1482            ]),
1483            storages: B256Map::from_iter([(
1484                addr2,
1485                HashedStorage {
1486                    wiped: true,
1487                    storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1488                },
1489            )]),
1490        };
1491
1492        let mut chunks = state.chunks(2);
1493        assert_eq!(
1494            chunks.next(),
1495            Some(HashedPostState {
1496                accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
1497                storages: B256Map::from_iter([(addr2, HashedStorage::new(true)),])
1498            })
1499        );
1500        assert_eq!(
1501            chunks.next(),
1502            Some(HashedPostState {
1503                accounts: B256Map::default(),
1504                storages: B256Map::from_iter([(
1505                    addr2,
1506                    HashedStorage {
1507                        wiped: false,
1508                        storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1509                    },
1510                )])
1511            })
1512        );
1513        assert_eq!(
1514            chunks.next(),
1515            Some(HashedPostState {
1516                accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1517                storages: B256Map::default()
1518            })
1519        );
1520        assert_eq!(chunks.next(), None);
1521    }
1522
1523    #[test]
1524    fn test_chunks_ordering_guarantee() {
1525        // Test that chunks preserve the ordering: wipe -> storage updates -> account
1526        // Use chunk size of 1 to verify each item comes out in the correct order
1527        let addr = B256::from([1; 32]);
1528        let slot1 = B256::from([1; 32]);
1529        let slot2 = B256::from([2; 32]);
1530
1531        let state = HashedPostState {
1532            accounts: B256Map::from_iter([(addr, Some(Default::default()))]),
1533            storages: B256Map::from_iter([(
1534                addr,
1535                HashedStorage {
1536                    wiped: true,
1537                    storage: B256Map::from_iter([(slot1, U256::from(1)), (slot2, U256::from(2))]),
1538                },
1539            )]),
1540        };
1541
1542        let chunks: Vec<_> = state.chunks(1).collect();
1543
1544        // Should have 4 chunks: 1 wipe + 2 storage updates + 1 account
1545        assert_eq!(chunks.len(), 4);
1546
1547        // First chunk must be the storage wipe
1548        assert!(chunks[0].accounts.is_empty());
1549        assert_eq!(chunks[0].storages.len(), 1);
1550        assert!(chunks[0].storages.get(&addr).unwrap().wiped);
1551        assert!(chunks[0].storages.get(&addr).unwrap().storage.is_empty());
1552
1553        // Next two chunks must be storage updates (order between them doesn't matter)
1554        assert!(chunks[1].accounts.is_empty());
1555        assert!(!chunks[1].storages.get(&addr).unwrap().wiped);
1556        assert_eq!(chunks[1].storages.get(&addr).unwrap().storage.len(), 1);
1557
1558        assert!(chunks[2].accounts.is_empty());
1559        assert!(!chunks[2].storages.get(&addr).unwrap().wiped);
1560        assert_eq!(chunks[2].storages.get(&addr).unwrap().storage.len(), 1);
1561
1562        // Last chunk must be the account update
1563        assert_eq!(chunks[3].accounts.len(), 1);
1564        assert!(chunks[3].accounts.contains_key(&addr));
1565        assert!(chunks[3].storages.is_empty());
1566    }
1567
1568    #[test]
1569    fn test_hashed_post_state_sorted_extend_ref() {
1570        // Test extending accounts
1571        let mut state1 = HashedPostStateSorted {
1572            accounts: vec![
1573                (B256::from([1; 32]), Some(Account::default())),
1574                (B256::from([3; 32]), Some(Account::default())),
1575                (B256::from([5; 32]), None),
1576            ],
1577            storages: B256Map::default(),
1578        };
1579
1580        let state2 = HashedPostStateSorted {
1581            accounts: vec![
1582                (B256::from([2; 32]), Some(Account::default())),
1583                (B256::from([3; 32]), Some(Account { nonce: 1, ..Default::default() })), /* Override */
1584                (B256::from([4; 32]), Some(Account::default())),
1585                (B256::from([6; 32]), None),
1586            ],
1587            storages: B256Map::default(),
1588        };
1589
1590        state1.extend_ref_and_sort(&state2);
1591
1592        // Check accounts are merged and sorted
1593        assert_eq!(state1.accounts.len(), 6);
1594        assert_eq!(state1.accounts[0].0, B256::from([1; 32]));
1595        assert_eq!(state1.accounts[1].0, B256::from([2; 32]));
1596        assert_eq!(state1.accounts[2].0, B256::from([3; 32]));
1597        assert_eq!(state1.accounts[2].1.unwrap().nonce, 1); // Should have state2's value
1598        assert_eq!(state1.accounts[3].0, B256::from([4; 32]));
1599        assert_eq!(state1.accounts[4].0, B256::from([5; 32]));
1600        assert_eq!(state1.accounts[4].1, None);
1601        assert_eq!(state1.accounts[5].0, B256::from([6; 32]));
1602        assert_eq!(state1.accounts[5].1, None);
1603    }
1604
1605    #[test]
1606    fn test_hashed_storage_sorted_extend_ref() {
1607        // Test normal extension
1608        let mut storage1 = HashedStorageSorted {
1609            storage_slots: vec![
1610                (B256::from([1; 32]), U256::from(10)),
1611                (B256::from([3; 32]), U256::from(30)),
1612                (B256::from([5; 32]), U256::ZERO),
1613            ],
1614            wiped: false,
1615        };
1616
1617        let storage2 = HashedStorageSorted {
1618            storage_slots: vec![
1619                (B256::from([2; 32]), U256::from(20)),
1620                (B256::from([3; 32]), U256::from(300)), // Override
1621                (B256::from([4; 32]), U256::from(40)),
1622                (B256::from([6; 32]), U256::ZERO),
1623            ],
1624            wiped: false,
1625        };
1626
1627        storage1.extend_ref(&storage2);
1628
1629        assert_eq!(storage1.storage_slots.len(), 6);
1630        assert_eq!(storage1.storage_slots[0].0, B256::from([1; 32]));
1631        assert_eq!(storage1.storage_slots[0].1, U256::from(10));
1632        assert_eq!(storage1.storage_slots[1].0, B256::from([2; 32]));
1633        assert_eq!(storage1.storage_slots[1].1, U256::from(20));
1634        assert_eq!(storage1.storage_slots[2].0, B256::from([3; 32]));
1635        assert_eq!(storage1.storage_slots[2].1, U256::from(300)); // Should have storage2's value
1636        assert_eq!(storage1.storage_slots[3].0, B256::from([4; 32]));
1637        assert_eq!(storage1.storage_slots[3].1, U256::from(40));
1638        assert_eq!(storage1.storage_slots[4].0, B256::from([5; 32]));
1639        assert_eq!(storage1.storage_slots[4].1, U256::ZERO);
1640        assert_eq!(storage1.storage_slots[5].0, B256::from([6; 32]));
1641        assert_eq!(storage1.storage_slots[5].1, U256::ZERO);
1642        assert!(!storage1.wiped);
1643
1644        // Test wiped storage
1645        let mut storage3 = HashedStorageSorted {
1646            storage_slots: vec![
1647                (B256::from([1; 32]), U256::from(10)),
1648                (B256::from([2; 32]), U256::ZERO),
1649            ],
1650            wiped: false,
1651        };
1652
1653        let storage4 = HashedStorageSorted {
1654            storage_slots: vec![
1655                (B256::from([3; 32]), U256::from(30)),
1656                (B256::from([4; 32]), U256::ZERO),
1657            ],
1658            wiped: true,
1659        };
1660
1661        storage3.extend_ref(&storage4);
1662
1663        assert!(storage3.wiped);
1664        // When wiped, should only have storage4's values
1665        assert_eq!(storage3.storage_slots.len(), 2);
1666        assert_eq!(storage3.storage_slots[0].0, B256::from([3; 32]));
1667        assert_eq!(storage3.storage_slots[0].1, U256::from(30));
1668        assert_eq!(storage3.storage_slots[1].0, B256::from([4; 32]));
1669        assert_eq!(storage3.storage_slots[1].1, U256::ZERO);
1670    }
1671
1672    /// Test extending with sorted accounts merges correctly into `HashMap`
1673    #[test]
1674    fn test_hashed_post_state_extend_from_sorted_with_accounts() {
1675        let addr1 = B256::random();
1676        let addr2 = B256::random();
1677
1678        let mut state = HashedPostState::default();
1679        state.accounts.insert(addr1, Some(Default::default()));
1680
1681        let mut sorted_state = HashedPostStateSorted::default();
1682        sorted_state.accounts.push((addr2, Some(Default::default())));
1683
1684        state.extend_from_sorted(&sorted_state);
1685
1686        assert_eq!(state.accounts.len(), 2);
1687        assert!(state.accounts.contains_key(&addr1));
1688        assert!(state.accounts.contains_key(&addr2));
1689    }
1690
1691    /// Test destroyed accounts (None values) are inserted correctly
1692    #[test]
1693    fn test_hashed_post_state_extend_from_sorted_with_destroyed_accounts() {
1694        let addr1 = B256::random();
1695
1696        let mut state = HashedPostState::default();
1697
1698        let mut sorted_state = HashedPostStateSorted::default();
1699        sorted_state.accounts.push((addr1, None));
1700
1701        state.extend_from_sorted(&sorted_state);
1702
1703        assert!(state.accounts.contains_key(&addr1));
1704        assert_eq!(state.accounts.get(&addr1), Some(&None));
1705    }
1706
1707    #[test]
1708    fn test_hashed_post_state_sorted_disjointed_merge_batch() {
1709        fn account(nonce: u64) -> Account {
1710            Account { nonce, balance: U256::ZERO, bytecode_hash: None }
1711        }
1712
1713        let kept_account = B256::with_last_byte(1);
1714        let removed_account = B256::with_last_byte(2);
1715        let kept_storage = B256::with_last_byte(3);
1716        let slot1 = B256::with_last_byte(11);
1717        let slot2 = B256::with_last_byte(12);
1718
1719        let older = HashedPostStateSorted::new(
1720            vec![(kept_account, Some(account(1))), (removed_account, Some(account(10)))],
1721            B256Map::from_iter([(
1722                kept_storage,
1723                HashedStorageSorted { wiped: false, storage_slots: vec![(slot1, U256::from(1))] },
1724            )]),
1725        );
1726
1727        let newer = HashedPostStateSorted::new(
1728            vec![(kept_account, Some(account(2)))],
1729            B256Map::from_iter([(
1730                kept_storage,
1731                HashedStorageSorted {
1732                    wiped: false,
1733                    storage_slots: vec![(slot1, U256::from(3)), (slot2, U256::from(4))],
1734                },
1735            )]),
1736        );
1737
1738        let remove_a = HashedPostStateSorted::new(
1739            vec![(removed_account, None)],
1740            B256Map::from_iter([(
1741                kept_storage,
1742                HashedStorageSorted { wiped: false, storage_slots: vec![(slot2, U256::ZERO)] },
1743            )]),
1744        );
1745
1746        let remove_b = HashedPostStateSorted::new(
1747            vec![(B256::with_last_byte(255), Some(account(99)))],
1748            B256Map::default(),
1749        );
1750
1751        let result = HashedPostStateSorted::disjointed_merge_batch(
1752            &[&older, &newer],
1753            &[&remove_b, &remove_a],
1754        );
1755
1756        assert_eq!(result.accounts, vec![(kept_account, Some(account(2)))]);
1757        assert_eq!(result.storages.len(), 1);
1758        assert_eq!(
1759            result.storages.get(&kept_storage),
1760            Some(&HashedStorageSorted {
1761                wiped: false,
1762                storage_slots: vec![(slot1, U256::from(3))],
1763            })
1764        );
1765    }
1766
1767    #[test]
1768    fn test_hashed_post_state_sorted_disjointed_merge_batch_empty_mask_merges_batch() {
1769        let address = B256::with_last_byte(1);
1770        let storage = B256::with_last_byte(2);
1771        let slot = B256::with_last_byte(3);
1772        let empty_storage = B256::with_last_byte(4);
1773        let older = HashedPostStateSorted::new(
1774            vec![(address, Some(Account { nonce: 1, ..Default::default() }))],
1775            B256Map::from_iter([
1776                (
1777                    storage,
1778                    HashedStorageSorted {
1779                        wiped: false,
1780                        storage_slots: vec![(slot, U256::from(1))],
1781                    },
1782                ),
1783                (empty_storage, HashedStorageSorted::default()),
1784            ]),
1785        );
1786        let newer = HashedPostStateSorted::new(
1787            vec![(address, Some(Account { nonce: 2, ..Default::default() }))],
1788            B256Map::from_iter([(
1789                storage,
1790                HashedStorageSorted { wiped: false, storage_slots: vec![(slot, U256::from(2))] },
1791            )]),
1792        );
1793        let expected = HashedPostStateSorted::merge_batch(vec![newer.clone(), older.clone()]);
1794
1795        let result = HashedPostStateSorted::disjointed_merge_batch(&[&older, &newer], &[]);
1796
1797        assert_eq!(result, expected);
1798    }
1799
1800    #[test]
1801    fn test_hashed_post_state_sorted_disjointed_merge_batch_removes_overlapping_batch_key() {
1802        fn account(nonce: u64) -> Account {
1803            Account { nonce, balance: U256::ZERO, bytecode_hash: None }
1804        }
1805
1806        let overlapping_account = B256::with_last_byte(21);
1807
1808        let older = HashedPostStateSorted::new(
1809            vec![(overlapping_account, Some(account(1)))],
1810            B256Map::default(),
1811        );
1812
1813        let newer = HashedPostStateSorted::new(
1814            vec![(overlapping_account, Some(account(2)))],
1815            B256Map::default(),
1816        );
1817
1818        let remove =
1819            HashedPostStateSorted::new(vec![(overlapping_account, None)], B256Map::default());
1820
1821        let result = HashedPostStateSorted::disjointed_merge_batch(&[&older, &newer], &[&remove]);
1822
1823        assert!(result.accounts.is_empty());
1824    }
1825
1826    #[test]
1827    fn test_hashed_post_state_sorted_disjointed_merge_batch_keeps_equal_overlaps() {
1828        fn account(nonce: u64) -> Account {
1829            Account { nonce, balance: U256::ZERO, bytecode_hash: None }
1830        }
1831
1832        let address = B256::with_last_byte(21);
1833        let deleted_address = B256::with_last_byte(24);
1834        let storage = B256::with_last_byte(22);
1835        let deleted_storage = B256::with_last_byte(25);
1836        let slot = B256::with_last_byte(23);
1837        let deleted_slot = B256::with_last_byte(26);
1838        let batch = HashedPostStateSorted::new(
1839            vec![(address, Some(account(1))), (deleted_address, None)],
1840            B256Map::from_iter([
1841                (
1842                    storage,
1843                    HashedStorageSorted {
1844                        wiped: false,
1845                        storage_slots: vec![(slot, U256::from(1))],
1846                    },
1847                ),
1848                (
1849                    deleted_storage,
1850                    HashedStorageSorted {
1851                        wiped: false,
1852                        storage_slots: vec![(deleted_slot, U256::ZERO)],
1853                    },
1854                ),
1855            ]),
1856        );
1857        let different_mask = HashedPostStateSorted::new(
1858            vec![(address, Some(account(2))), (deleted_address, Some(account(3)))],
1859            B256Map::from_iter([
1860                (
1861                    storage,
1862                    HashedStorageSorted {
1863                        wiped: false,
1864                        storage_slots: vec![(slot, U256::from(2))],
1865                    },
1866                ),
1867                (
1868                    deleted_storage,
1869                    HashedStorageSorted {
1870                        wiped: false,
1871                        storage_slots: vec![(deleted_slot, U256::from(3))],
1872                    },
1873                ),
1874            ]),
1875        );
1876        let equal_mask = batch.clone();
1877
1878        let result = HashedPostStateSorted::disjointed_merge_batch(
1879            &[&batch],
1880            &[&different_mask, &equal_mask],
1881        );
1882        let reversed = HashedPostStateSorted::disjointed_merge_batch(
1883            &[&batch],
1884            &[&equal_mask, &different_mask],
1885        );
1886
1887        assert_eq!(result, batch);
1888        assert_eq!(reversed, result);
1889    }
1890
1891    #[test]
1892    fn test_hashed_post_state_sorted_disjointed_merge_batch_ignores_empty_storage_mask() {
1893        let storage = B256::with_last_byte(31);
1894        let slot = B256::with_last_byte(32);
1895
1896        let batch = HashedPostStateSorted::new(
1897            vec![],
1898            B256Map::from_iter([(
1899                storage,
1900                HashedStorageSorted { wiped: false, storage_slots: vec![(slot, U256::from(1))] },
1901            )]),
1902        );
1903        let mask = HashedPostStateSorted::new(
1904            vec![],
1905            B256Map::from_iter([(
1906                storage,
1907                HashedStorageSorted { wiped: false, storage_slots: vec![] },
1908            )]),
1909        );
1910
1911        let result = HashedPostStateSorted::disjointed_merge_batch(&[&batch], &[&mask]);
1912
1913        assert_eq!(
1914            result.storages.get(&storage),
1915            Some(&HashedStorageSorted { wiped: false, storage_slots: vec![(slot, U256::from(1))] })
1916        );
1917    }
1918
1919    /// Test non-wiped storage merges both zero and non-zero valued slots
1920    #[test]
1921    fn test_hashed_storage_extend_from_sorted_non_wiped() {
1922        let slot1 = B256::random();
1923        let slot2 = B256::random();
1924        let slot3 = B256::random();
1925
1926        let mut storage = HashedStorage::from_iter(false, [(slot1, U256::from(100))]);
1927
1928        let sorted = HashedStorageSorted {
1929            storage_slots: vec![(slot2, U256::from(200)), (slot3, U256::ZERO)],
1930            wiped: false,
1931        };
1932
1933        storage.extend_from_sorted(&sorted);
1934
1935        assert!(!storage.wiped);
1936        assert_eq!(storage.storage.len(), 3);
1937        assert_eq!(storage.storage.get(&slot1), Some(&U256::from(100)));
1938        assert_eq!(storage.storage.get(&slot2), Some(&U256::from(200)));
1939        assert_eq!(storage.storage.get(&slot3), Some(&U256::ZERO));
1940    }
1941
1942    /// Test wiped=true clears existing storage and only keeps new slots (critical edge case)
1943    #[test]
1944    fn test_hashed_storage_extend_from_sorted_wiped() {
1945        let slot1 = B256::random();
1946        let slot2 = B256::random();
1947
1948        let mut storage = HashedStorage::from_iter(false, [(slot1, U256::from(100))]);
1949
1950        let sorted =
1951            HashedStorageSorted { storage_slots: vec![(slot2, U256::from(200))], wiped: true };
1952
1953        storage.extend_from_sorted(&sorted);
1954
1955        assert!(storage.wiped);
1956        // After wipe, old storage should be cleared and only new storage remains
1957        assert_eq!(storage.storage.len(), 1);
1958        assert_eq!(storage.storage.get(&slot2), Some(&U256::from(200)));
1959    }
1960
1961    #[test]
1962    fn test_hashed_post_state_chunking_length() {
1963        let addr1 = B256::from([1; 32]);
1964        let addr2 = B256::from([2; 32]);
1965        let addr3 = B256::from([3; 32]);
1966        let addr4 = B256::from([4; 32]);
1967        let slot1 = B256::from([1; 32]);
1968        let slot2 = B256::from([2; 32]);
1969        let slot3 = B256::from([3; 32]);
1970
1971        let state = HashedPostState {
1972            accounts: B256Map::from_iter([(addr1, None), (addr2, None), (addr4, None)]),
1973            storages: B256Map::from_iter([
1974                (
1975                    addr1,
1976                    HashedStorage {
1977                        wiped: false,
1978                        storage: B256Map::from_iter([
1979                            (slot1, U256::ZERO),
1980                            (slot2, U256::ZERO),
1981                            (slot3, U256::ZERO),
1982                        ]),
1983                    },
1984                ),
1985                (
1986                    addr2,
1987                    HashedStorage {
1988                        wiped: true,
1989                        storage: B256Map::from_iter([
1990                            (slot1, U256::ZERO),
1991                            (slot2, U256::ZERO),
1992                            (slot3, U256::ZERO),
1993                        ]),
1994                    },
1995                ),
1996                (
1997                    addr3,
1998                    HashedStorage {
1999                        wiped: false,
2000                        storage: B256Map::from_iter([
2001                            (slot1, U256::ZERO),
2002                            (slot2, U256::ZERO),
2003                            (slot3, U256::ZERO),
2004                        ]),
2005                    },
2006                ),
2007            ]),
2008        };
2009
2010        let chunking_length = state.chunking_length();
2011        for size in 1..=state.clone().chunks(1).count() {
2012            let chunk_count = state.clone().chunks(size).count();
2013            let expected_count = chunking_length.div_ceil(size);
2014            assert_eq!(
2015                chunk_count, expected_count,
2016                "chunking_length: {}, size: {}",
2017                chunking_length, size
2018            );
2019        }
2020    }
2021
2022    #[test]
2023    fn test_clone_into_sorted_equivalence() {
2024        let addr1 = B256::from([1; 32]);
2025        let addr2 = B256::from([2; 32]);
2026        let addr3 = B256::from([3; 32]);
2027        let slot1 = B256::from([1; 32]);
2028        let slot2 = B256::from([2; 32]);
2029        let slot3 = B256::from([3; 32]);
2030
2031        let state = HashedPostState {
2032            accounts: B256Map::from_iter([
2033                (addr1, Some(Account { nonce: 1, balance: U256::from(100), bytecode_hash: None })),
2034                (addr2, None),
2035                (addr3, Some(Account::default())),
2036            ]),
2037            storages: B256Map::from_iter([
2038                (
2039                    addr1,
2040                    HashedStorage {
2041                        wiped: false,
2042                        storage: B256Map::from_iter([
2043                            (slot1, U256::from(10)),
2044                            (slot2, U256::from(20)),
2045                        ]),
2046                    },
2047                ),
2048                (
2049                    addr2,
2050                    HashedStorage {
2051                        wiped: true,
2052                        storage: B256Map::from_iter([(slot3, U256::ZERO)]),
2053                    },
2054                ),
2055            ]),
2056        };
2057
2058        // clone_into_sorted should produce the same result as clone().into_sorted()
2059        let sorted_via_clone = state.clone().into_sorted();
2060        let sorted_via_clone_into = state.clone_into_sorted();
2061
2062        assert_eq!(sorted_via_clone, sorted_via_clone_into);
2063
2064        // Verify the original state is not consumed
2065        assert_eq!(state.accounts.len(), 3);
2066        assert_eq!(state.storages.len(), 2);
2067    }
2068
2069    #[test]
2070    fn test_hashed_storage_clone_into_sorted_equivalence() {
2071        let slot1 = B256::from([1; 32]);
2072        let slot2 = B256::from([2; 32]);
2073        let slot3 = B256::from([3; 32]);
2074
2075        let storage = HashedStorage {
2076            wiped: true,
2077            storage: B256Map::from_iter([
2078                (slot1, U256::from(100)),
2079                (slot2, U256::ZERO),
2080                (slot3, U256::from(300)),
2081            ]),
2082        };
2083
2084        // clone_into_sorted should produce the same result as clone().into_sorted()
2085        let sorted_via_clone = storage.clone().into_sorted();
2086        let sorted_via_clone_into = storage.clone_into_sorted();
2087
2088        assert_eq!(sorted_via_clone, sorted_via_clone_into);
2089
2090        // Verify the original storage is not consumed
2091        assert_eq!(storage.storage.len(), 3);
2092        assert!(storage.wiped);
2093    }
2094}
2095
2096/// Bincode-compatible hashed state type serde implementations.
2097#[cfg(feature = "serde-bincode-compat")]
2098pub mod serde_bincode_compat {
2099    use super::Account;
2100    use alloc::borrow::Cow;
2101    use alloy_primitives::{map::B256Map, B256, U256};
2102    use serde::{Deserialize, Deserializer, Serialize, Serializer};
2103    use serde_with::{DeserializeAs, SerializeAs};
2104
2105    /// Bincode-compatible [`super::HashedPostState`] serde implementation.
2106    ///
2107    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
2108    /// ```rust
2109    /// use reth_trie_common::{serde_bincode_compat, HashedPostState};
2110    /// use serde::{Deserialize, Serialize};
2111    /// use serde_with::serde_as;
2112    ///
2113    /// #[serde_as]
2114    /// #[derive(Serialize, Deserialize)]
2115    /// struct Data {
2116    ///     #[serde_as(as = "serde_bincode_compat::hashed_state::HashedPostState")]
2117    ///     hashed_state: HashedPostState,
2118    /// }
2119    /// ```
2120    #[derive(Debug, Serialize, Deserialize)]
2121    pub struct HashedPostState<'a> {
2122        accounts: Cow<'a, B256Map<Option<Account>>>,
2123        storages: B256Map<HashedStorage<'a>>,
2124    }
2125
2126    impl<'a> From<&'a super::HashedPostState> for HashedPostState<'a> {
2127        fn from(value: &'a super::HashedPostState) -> Self {
2128            Self {
2129                accounts: Cow::Borrowed(&value.accounts),
2130                storages: value.storages.iter().map(|(k, v)| (*k, v.into())).collect(),
2131            }
2132        }
2133    }
2134
2135    impl<'a> From<HashedPostState<'a>> for super::HashedPostState {
2136        fn from(value: HashedPostState<'a>) -> Self {
2137            Self {
2138                accounts: value.accounts.into_owned(),
2139                storages: value.storages.into_iter().map(|(k, v)| (k, v.into())).collect(),
2140            }
2141        }
2142    }
2143
2144    impl SerializeAs<super::HashedPostState> for HashedPostState<'_> {
2145        fn serialize_as<S>(
2146            source: &super::HashedPostState,
2147            serializer: S,
2148        ) -> Result<S::Ok, S::Error>
2149        where
2150            S: Serializer,
2151        {
2152            HashedPostState::from(source).serialize(serializer)
2153        }
2154    }
2155
2156    impl<'de> DeserializeAs<'de, super::HashedPostState> for HashedPostState<'de> {
2157        fn deserialize_as<D>(deserializer: D) -> Result<super::HashedPostState, D::Error>
2158        where
2159            D: Deserializer<'de>,
2160        {
2161            HashedPostState::deserialize(deserializer).map(Into::into)
2162        }
2163    }
2164
2165    /// Bincode-compatible [`super::HashedStorage`] serde implementation.
2166    ///
2167    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
2168    /// ```rust
2169    /// use reth_trie_common::{serde_bincode_compat, HashedStorage};
2170    /// use serde::{Deserialize, Serialize};
2171    /// use serde_with::serde_as;
2172    ///
2173    /// #[serde_as]
2174    /// #[derive(Serialize, Deserialize)]
2175    /// struct Data {
2176    ///     #[serde_as(as = "serde_bincode_compat::hashed_state::HashedStorage")]
2177    ///     hashed_storage: HashedStorage,
2178    /// }
2179    /// ```
2180    #[derive(Debug, Serialize, Deserialize)]
2181    pub struct HashedStorage<'a> {
2182        wiped: bool,
2183        storage: Cow<'a, B256Map<U256>>,
2184    }
2185
2186    impl<'a> From<&'a super::HashedStorage> for HashedStorage<'a> {
2187        fn from(value: &'a super::HashedStorage) -> Self {
2188            Self { wiped: value.wiped, storage: Cow::Borrowed(&value.storage) }
2189        }
2190    }
2191
2192    impl<'a> From<HashedStorage<'a>> for super::HashedStorage {
2193        fn from(value: HashedStorage<'a>) -> Self {
2194            Self { wiped: value.wiped, storage: value.storage.into_owned() }
2195        }
2196    }
2197
2198    impl SerializeAs<super::HashedStorage> for HashedStorage<'_> {
2199        fn serialize_as<S>(source: &super::HashedStorage, serializer: S) -> Result<S::Ok, S::Error>
2200        where
2201            S: Serializer,
2202        {
2203            HashedStorage::from(source).serialize(serializer)
2204        }
2205    }
2206
2207    impl<'de> DeserializeAs<'de, super::HashedStorage> for HashedStorage<'de> {
2208        fn deserialize_as<D>(deserializer: D) -> Result<super::HashedStorage, D::Error>
2209        where
2210            D: Deserializer<'de>,
2211        {
2212            HashedStorage::deserialize(deserializer).map(Into::into)
2213        }
2214    }
2215
2216    /// Bincode-compatible [`super::HashedPostStateSorted`] serde implementation.
2217    ///
2218    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
2219    /// ```rust
2220    /// use reth_trie_common::{serde_bincode_compat, HashedPostStateSorted};
2221    /// use serde::{Deserialize, Serialize};
2222    /// use serde_with::serde_as;
2223    ///
2224    /// #[serde_as]
2225    /// #[derive(Serialize, Deserialize)]
2226    /// struct Data {
2227    ///     #[serde_as(as = "serde_bincode_compat::hashed_state::HashedPostStateSorted")]
2228    ///     hashed_state: HashedPostStateSorted,
2229    /// }
2230    /// ```
2231    #[derive(Debug, Serialize, Deserialize)]
2232    pub struct HashedPostStateSorted<'a> {
2233        accounts: Cow<'a, [(B256, Option<Account>)]>,
2234        storages: B256Map<HashedStorageSorted<'a>>,
2235    }
2236
2237    impl<'a> From<&'a super::HashedPostStateSorted> for HashedPostStateSorted<'a> {
2238        fn from(value: &'a super::HashedPostStateSorted) -> Self {
2239            Self {
2240                accounts: Cow::Borrowed(&value.accounts),
2241                storages: value.storages.iter().map(|(k, v)| (*k, v.into())).collect(),
2242            }
2243        }
2244    }
2245
2246    impl<'a> From<HashedPostStateSorted<'a>> for super::HashedPostStateSorted {
2247        fn from(value: HashedPostStateSorted<'a>) -> Self {
2248            Self {
2249                accounts: value.accounts.into_owned(),
2250                storages: value.storages.into_iter().map(|(k, v)| (k, v.into())).collect(),
2251            }
2252        }
2253    }
2254
2255    impl SerializeAs<super::HashedPostStateSorted> for HashedPostStateSorted<'_> {
2256        fn serialize_as<S>(
2257            source: &super::HashedPostStateSorted,
2258            serializer: S,
2259        ) -> Result<S::Ok, S::Error>
2260        where
2261            S: Serializer,
2262        {
2263            HashedPostStateSorted::from(source).serialize(serializer)
2264        }
2265    }
2266
2267    impl<'de> DeserializeAs<'de, super::HashedPostStateSorted> for HashedPostStateSorted<'de> {
2268        fn deserialize_as<D>(deserializer: D) -> Result<super::HashedPostStateSorted, D::Error>
2269        where
2270            D: Deserializer<'de>,
2271        {
2272            HashedPostStateSorted::deserialize(deserializer).map(Into::into)
2273        }
2274    }
2275
2276    /// Bincode-compatible [`super::HashedStorageSorted`] serde implementation.
2277    ///
2278    /// Intended to use with the [`serde_with::serde_as`] macro in the following way:
2279    /// ```rust
2280    /// use reth_trie_common::{serde_bincode_compat, HashedStorageSorted};
2281    /// use serde::{Deserialize, Serialize};
2282    /// use serde_with::serde_as;
2283    ///
2284    /// #[serde_as]
2285    /// #[derive(Serialize, Deserialize)]
2286    /// struct Data {
2287    ///     #[serde_as(as = "serde_bincode_compat::hashed_state::HashedStorageSorted")]
2288    ///     hashed_storage: HashedStorageSorted,
2289    /// }
2290    /// ```
2291    #[derive(Debug, Serialize, Deserialize)]
2292    pub struct HashedStorageSorted<'a> {
2293        storage_slots: Cow<'a, [(B256, U256)]>,
2294        wiped: bool,
2295    }
2296
2297    impl<'a> From<&'a super::HashedStorageSorted> for HashedStorageSorted<'a> {
2298        fn from(value: &'a super::HashedStorageSorted) -> Self {
2299            Self { storage_slots: Cow::Borrowed(&value.storage_slots), wiped: value.wiped }
2300        }
2301    }
2302
2303    impl<'a> From<HashedStorageSorted<'a>> for super::HashedStorageSorted {
2304        fn from(value: HashedStorageSorted<'a>) -> Self {
2305            Self { storage_slots: value.storage_slots.into_owned(), wiped: value.wiped }
2306        }
2307    }
2308
2309    impl SerializeAs<super::HashedStorageSorted> for HashedStorageSorted<'_> {
2310        fn serialize_as<S>(
2311            source: &super::HashedStorageSorted,
2312            serializer: S,
2313        ) -> Result<S::Ok, S::Error>
2314        where
2315            S: Serializer,
2316        {
2317            HashedStorageSorted::from(source).serialize(serializer)
2318        }
2319    }
2320
2321    impl<'de> DeserializeAs<'de, super::HashedStorageSorted> for HashedStorageSorted<'de> {
2322        fn deserialize_as<D>(deserializer: D) -> Result<super::HashedStorageSorted, D::Error>
2323        where
2324            D: Deserializer<'de>,
2325        {
2326            HashedStorageSorted::deserialize(deserializer).map(Into::into)
2327        }
2328    }
2329
2330    #[cfg(test)]
2331    mod tests {
2332        use crate::{
2333            hashed_state::{
2334                HashedPostState, HashedPostStateSorted, HashedStorage, HashedStorageSorted,
2335            },
2336            serde_bincode_compat,
2337        };
2338        use alloy_primitives::{B256, U256};
2339        use reth_primitives_traits::Account;
2340        use serde::{Deserialize, Serialize};
2341        use serde_with::serde_as;
2342
2343        #[test]
2344        fn test_hashed_post_state_bincode_roundtrip() {
2345            #[serde_as]
2346            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
2347            struct Data {
2348                #[serde_as(as = "serde_bincode_compat::hashed_state::HashedPostState")]
2349                hashed_state: HashedPostState,
2350            }
2351
2352            let mut data = Data { hashed_state: HashedPostState::default() };
2353            let encoded = bincode::serialize(&data).unwrap();
2354            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2355            assert_eq!(decoded, data);
2356
2357            data.hashed_state.accounts.insert(B256::random(), Some(Account::default()));
2358            let encoded = bincode::serialize(&data).unwrap();
2359            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2360            assert_eq!(decoded, data);
2361
2362            data.hashed_state.storages.insert(B256::random(), HashedStorage::default());
2363            let encoded = bincode::serialize(&data).unwrap();
2364            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2365            assert_eq!(decoded, data);
2366        }
2367
2368        #[test]
2369        fn test_hashed_storage_bincode_roundtrip() {
2370            #[serde_as]
2371            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
2372            struct Data {
2373                #[serde_as(as = "serde_bincode_compat::hashed_state::HashedStorage")]
2374                hashed_storage: HashedStorage,
2375            }
2376
2377            let mut data = Data { hashed_storage: HashedStorage::default() };
2378            let encoded = bincode::serialize(&data).unwrap();
2379            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2380            assert_eq!(decoded, data);
2381
2382            data.hashed_storage.wiped = true;
2383            let encoded = bincode::serialize(&data).unwrap();
2384            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2385            assert_eq!(decoded, data);
2386
2387            data.hashed_storage.storage.insert(B256::random(), U256::from(1));
2388            let encoded = bincode::serialize(&data).unwrap();
2389            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2390            assert_eq!(decoded, data);
2391        }
2392
2393        #[test]
2394        fn test_hashed_post_state_sorted_bincode_roundtrip() {
2395            #[serde_as]
2396            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
2397            struct Data {
2398                #[serde_as(as = "serde_bincode_compat::hashed_state::HashedPostStateSorted")]
2399                hashed_state: HashedPostStateSorted,
2400            }
2401
2402            let mut data = Data { hashed_state: HashedPostStateSorted::default() };
2403            let encoded = bincode::serialize(&data).unwrap();
2404            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2405            assert_eq!(decoded, data);
2406
2407            data.hashed_state.accounts.push((B256::random(), Some(Account::default())));
2408            data.hashed_state
2409                .accounts
2410                .push((B256::random(), Some(Account { nonce: 1, ..Default::default() })));
2411            let encoded = bincode::serialize(&data).unwrap();
2412            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2413            assert_eq!(decoded, data);
2414
2415            data.hashed_state.storages.insert(
2416                B256::random(),
2417                HashedStorageSorted {
2418                    storage_slots: vec![(B256::from([1; 32]), U256::from(10))],
2419                    wiped: false,
2420                },
2421            );
2422            let encoded = bincode::serialize(&data).unwrap();
2423            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2424            assert_eq!(decoded, data);
2425        }
2426
2427        #[test]
2428        fn test_hashed_storage_sorted_bincode_roundtrip() {
2429            #[serde_as]
2430            #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
2431            struct Data {
2432                #[serde_as(as = "serde_bincode_compat::hashed_state::HashedStorageSorted")]
2433                hashed_storage: HashedStorageSorted,
2434            }
2435
2436            let mut data = Data {
2437                hashed_storage: HashedStorageSorted { storage_slots: Vec::new(), wiped: false },
2438            };
2439            let encoded = bincode::serialize(&data).unwrap();
2440            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2441            assert_eq!(decoded, data);
2442
2443            data.hashed_storage.wiped = true;
2444            let encoded = bincode::serialize(&data).unwrap();
2445            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2446            assert_eq!(decoded, data);
2447
2448            data.hashed_storage.storage_slots.push((B256::random(), U256::from(1)));
2449            let encoded = bincode::serialize(&data).unwrap();
2450            let decoded: Data = bincode::deserialize(&encoded).unwrap();
2451            assert_eq!(decoded, data);
2452        }
2453    }
2454}