1use core::ops::Not;
2
3use crate::{
4 added_removed_keys::MultiAddedRemovedKeys,
5 prefix_set::{PrefixSetMut, TriePrefixSetsMut},
6 KeyHasher, MultiProofTargets, Nibbles,
7};
8use alloc::{borrow::Cow, vec::Vec};
9use alloy_primitives::{
10 keccak256,
11 map::{hash_map, B256Map, B256Set, HashMap, HashSet},
12 Address, B256, U256,
13};
14use itertools::Itertools;
15#[cfg(feature = "rayon")]
16pub use rayon::*;
17use reth_primitives_traits::Account;
18
19#[cfg(feature = "rayon")]
20use rayon::prelude::{IntoParallelIterator, ParallelIterator};
21
22use revm_database::{AccountStatus, BundleAccount};
23
24#[derive(PartialEq, Eq, Clone, Default, Debug)]
26#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
27pub struct HashedPostState {
28 pub accounts: B256Map<Option<Account>>,
30 pub storages: B256Map<HashedStorage>,
32}
33
34impl HashedPostState {
35 pub fn with_capacity(capacity: usize) -> Self {
37 Self {
38 accounts: B256Map::with_capacity_and_hasher(capacity, Default::default()),
39 storages: B256Map::with_capacity_and_hasher(capacity, Default::default()),
40 }
41 }
42
43 #[inline]
47 #[cfg(feature = "rayon")]
48 pub fn from_bundle_state<'a, KH: KeyHasher>(
49 state: impl IntoParallelIterator<Item = (&'a Address, &'a BundleAccount)>,
50 ) -> Self {
51 let hashed = state
52 .into_par_iter()
53 .map(|(address, account)| {
54 let hashed_address = KH::hash_key(address);
55 let hashed_account = account.info.as_ref().map(Into::into);
56 let hashed_storage = HashedStorage::from_plain_storage(
57 account.status,
58 account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
59 );
60 (hashed_address, (hashed_account, hashed_storage))
61 })
62 .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
63
64 let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
65 let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
66 for (address, (account, storage)) in hashed {
67 accounts.insert(address, account);
68 if !storage.is_empty() {
69 storages.insert(address, storage);
70 }
71 }
72 Self { accounts, storages }
73 }
74
75 #[cfg(not(feature = "rayon"))]
79 pub fn from_bundle_state<'a, KH: KeyHasher>(
80 state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
81 ) -> Self {
82 let hashed = state
83 .into_iter()
84 .map(|(address, account)| {
85 let hashed_address = KH::hash_key(address);
86 let hashed_account = account.info.as_ref().map(Into::into);
87 let hashed_storage = HashedStorage::from_plain_storage(
88 account.status,
89 account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
90 );
91 (hashed_address, (hashed_account, hashed_storage))
92 })
93 .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
94
95 let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
96 let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
97 for (address, (account, storage)) in hashed {
98 accounts.insert(address, account);
99 if !storage.is_empty() {
100 storages.insert(address, storage);
101 }
102 }
103 Self { accounts, storages }
104 }
105
106 pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
108 Self {
109 accounts: HashMap::default(),
110 storages: HashMap::from_iter([(hashed_address, storage)]),
111 }
112 }
113
114 pub fn with_accounts(
116 mut self,
117 accounts: impl IntoIterator<Item = (B256, Option<Account>)>,
118 ) -> Self {
119 self.accounts = HashMap::from_iter(accounts);
120 self
121 }
122
123 pub fn with_storages(
125 mut self,
126 storages: impl IntoIterator<Item = (B256, HashedStorage)>,
127 ) -> Self {
128 self.storages = HashMap::from_iter(storages);
129 self
130 }
131
132 pub fn is_empty(&self) -> bool {
134 self.accounts.is_empty() && self.storages.is_empty()
135 }
136
137 pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
141 let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
143 let mut destroyed_accounts = HashSet::default();
144 for (hashed_address, account) in &self.accounts {
145 account_prefix_set.insert(Nibbles::unpack(hashed_address));
146
147 if account.is_none() {
148 destroyed_accounts.insert(*hashed_address);
149 }
150 }
151
152 let mut storage_prefix_sets =
154 HashMap::with_capacity_and_hasher(self.storages.len(), Default::default());
155 for (hashed_address, hashed_storage) in &self.storages {
156 account_prefix_set.insert(Nibbles::unpack(hashed_address));
157 storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
158 }
159
160 TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
161 }
162
163 pub fn multi_proof_targets(&self) -> MultiProofTargets {
165 let mut targets = MultiProofTargets::with_capacity(self.accounts.len());
167 for hashed_address in self.accounts.keys() {
168 targets.insert(*hashed_address, Default::default());
169 }
170 for (hashed_address, storage) in &self.storages {
171 targets.entry(*hashed_address).or_default().extend(storage.storage.keys().copied());
172 }
173 targets
174 }
175
176 pub fn multi_proof_targets_difference(
182 &self,
183 excluded: &MultiProofTargets,
184 ) -> MultiProofTargets {
185 let mut targets = MultiProofTargets::default();
186 for hashed_address in self.accounts.keys() {
187 if !excluded.contains_key(hashed_address) {
188 targets.insert(*hashed_address, Default::default());
189 }
190 }
191 for (hashed_address, storage) in &self.storages {
192 let maybe_excluded_storage = excluded.get(hashed_address);
193 let mut hashed_slots_targets = storage
194 .storage
195 .keys()
196 .filter(|slot| !maybe_excluded_storage.is_some_and(|f| f.contains(*slot)))
197 .peekable();
198 if hashed_slots_targets.peek().is_some() {
199 targets.entry(*hashed_address).or_default().extend(hashed_slots_targets);
200 }
201 }
202 targets
203 }
204
205 pub fn partition_by_targets(
212 mut self,
213 targets: &MultiProofTargets,
214 added_removed_keys: &MultiAddedRemovedKeys,
215 ) -> (Self, Self) {
216 let mut state_updates_not_in_targets = Self::default();
217
218 self.storages.retain(|&address, storage| {
219 let storage_added_removed_keys = added_removed_keys.get_storage(&address);
220
221 let (retain, storage_not_in_targets) = match targets.get(&address) {
222 Some(storage_in_targets) => {
223 let mut storage_not_in_targets = HashedStorage::default();
224 storage.storage.retain(|&slot, value| {
225 if storage_in_targets.contains(&slot) &&
226 !storage_added_removed_keys.is_some_and(|k| k.is_removed(&slot))
227 {
228 return true
229 }
230
231 storage_not_in_targets.storage.insert(slot, *value);
232 false
233 });
234
235 let retain = !storage.storage.is_empty();
239
240 if !retain {
244 storage_not_in_targets.wiped = storage.wiped;
245 }
246
247 (
248 retain,
249 storage_not_in_targets.is_empty().not().then_some(storage_not_in_targets),
250 )
251 }
252 None => (false, Some(core::mem::take(storage))),
253 };
254
255 if let Some(storage_not_in_targets) = storage_not_in_targets {
256 state_updates_not_in_targets.storages.insert(address, storage_not_in_targets);
257 }
258
259 retain
260 });
261 self.accounts.retain(|&address, account| {
262 if targets.contains_key(&address) {
263 return true
264 }
265
266 state_updates_not_in_targets.accounts.insert(address, *account);
267 false
268 });
269
270 (self, state_updates_not_in_targets)
271 }
272
273 pub fn chunks(self, size: usize) -> ChunkedHashedPostState {
277 ChunkedHashedPostState::new(self, size)
278 }
279
280 pub fn extend(&mut self, other: Self) {
283 self.extend_inner(Cow::Owned(other));
284 }
285
286 pub fn extend_ref(&mut self, other: &Self) {
291 self.extend_inner(Cow::Borrowed(other));
292 }
293
294 fn extend_inner(&mut self, other: Cow<'_, Self>) {
295 self.accounts.extend(other.accounts.iter().map(|(&k, &v)| (k, v)));
296
297 self.storages.reserve(other.storages.len());
298 match other {
299 Cow::Borrowed(other) => {
300 self.extend_storages(other.storages.iter().map(|(k, v)| (*k, Cow::Borrowed(v))))
301 }
302 Cow::Owned(other) => {
303 self.extend_storages(other.storages.into_iter().map(|(k, v)| (k, Cow::Owned(v))))
304 }
305 }
306 }
307
308 fn extend_storages<'a>(
309 &mut self,
310 storages: impl IntoIterator<Item = (B256, Cow<'a, HashedStorage>)>,
311 ) {
312 for (hashed_address, storage) in storages {
313 match self.storages.entry(hashed_address) {
314 hash_map::Entry::Vacant(entry) => {
315 entry.insert(storage.into_owned());
316 }
317 hash_map::Entry::Occupied(mut entry) => {
318 entry.get_mut().extend(&storage);
319 }
320 }
321 }
322 }
323
324 pub fn into_sorted(self) -> HashedPostStateSorted {
326 let mut updated_accounts = Vec::new();
327 let mut destroyed_accounts = HashSet::default();
328 for (hashed_address, info) in self.accounts {
329 if let Some(info) = info {
330 updated_accounts.push((hashed_address, info));
331 } else {
332 destroyed_accounts.insert(hashed_address);
333 }
334 }
335 updated_accounts.sort_unstable_by_key(|(address, _)| *address);
336 let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts };
337
338 let storages = self
339 .storages
340 .into_iter()
341 .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
342 .collect();
343
344 HashedPostStateSorted { accounts, storages }
345 }
346
347 pub fn drain_into_sorted(&mut self) -> HashedPostStateSorted {
355 let mut updated_accounts = Vec::new();
356 let mut destroyed_accounts = HashSet::default();
357 for (hashed_address, info) in self.accounts.drain() {
358 if let Some(info) = info {
359 updated_accounts.push((hashed_address, info));
360 } else {
361 destroyed_accounts.insert(hashed_address);
362 }
363 }
364 updated_accounts.sort_unstable_by_key(|(address, _)| *address);
365 let accounts = HashedAccountsSorted { accounts: updated_accounts, destroyed_accounts };
366
367 let storages = self
368 .storages
369 .drain()
370 .map(|(hashed_address, storage)| (hashed_address, storage.into_sorted()))
371 .collect();
372
373 HashedPostStateSorted { accounts, storages }
374 }
375
376 pub fn clear(&mut self) {
378 self.accounts.clear();
379 self.storages.clear();
380 }
381}
382
383#[derive(PartialEq, Eq, Clone, Debug, Default)]
385#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
386pub struct HashedStorage {
387 pub wiped: bool,
389 pub storage: B256Map<U256>,
391}
392
393impl HashedStorage {
394 pub fn new(wiped: bool) -> Self {
396 Self { wiped, storage: HashMap::default() }
397 }
398
399 pub fn is_empty(&self) -> bool {
401 !self.wiped && self.storage.is_empty()
402 }
403
404 pub fn from_iter(wiped: bool, iter: impl IntoIterator<Item = (B256, U256)>) -> Self {
406 Self { wiped, storage: HashMap::from_iter(iter) }
407 }
408
409 pub fn from_plain_storage<'a>(
411 status: AccountStatus,
412 storage: impl IntoIterator<Item = (&'a U256, &'a U256)>,
413 ) -> Self {
414 Self::from_iter(
415 status.was_destroyed(),
416 storage.into_iter().map(|(key, value)| (keccak256(B256::from(*key)), *value)),
417 )
418 }
419
420 pub fn construct_prefix_set(&self) -> PrefixSetMut {
422 if self.wiped {
423 PrefixSetMut::all()
424 } else {
425 let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
426 for hashed_slot in self.storage.keys() {
427 prefix_set.insert(Nibbles::unpack(hashed_slot));
428 }
429 prefix_set
430 }
431 }
432
433 pub fn extend(&mut self, other: &Self) {
436 if other.wiped {
437 self.wiped = true;
438 self.storage.clear();
439 }
440 self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v)));
441 }
442
443 pub fn into_sorted(self) -> HashedStorageSorted {
445 let mut non_zero_valued_slots = Vec::new();
446 let mut zero_valued_slots = HashSet::default();
447 for (hashed_slot, value) in self.storage {
448 if value.is_zero() {
449 zero_valued_slots.insert(hashed_slot);
450 } else {
451 non_zero_valued_slots.push((hashed_slot, value));
452 }
453 }
454 non_zero_valued_slots.sort_unstable_by_key(|(key, _)| *key);
455
456 HashedStorageSorted { non_zero_valued_slots, zero_valued_slots, wiped: self.wiped }
457 }
458}
459
460#[derive(PartialEq, Eq, Clone, Default, Debug)]
462pub struct HashedPostStateSorted {
463 pub accounts: HashedAccountsSorted,
465 pub storages: B256Map<HashedStorageSorted>,
467}
468
469impl HashedPostStateSorted {
470 pub const fn new(
472 accounts: HashedAccountsSorted,
473 storages: B256Map<HashedStorageSorted>,
474 ) -> Self {
475 Self { accounts, storages }
476 }
477
478 pub const fn accounts(&self) -> &HashedAccountsSorted {
480 &self.accounts
481 }
482
483 pub const fn account_storages(&self) -> &B256Map<HashedStorageSorted> {
485 &self.storages
486 }
487}
488
489#[derive(Clone, Eq, PartialEq, Default, Debug)]
491pub struct HashedAccountsSorted {
492 pub accounts: Vec<(B256, Account)>,
494 pub destroyed_accounts: B256Set,
496}
497
498impl HashedAccountsSorted {
499 pub fn accounts_sorted(&self) -> impl Iterator<Item = (B256, Option<Account>)> {
501 self.accounts
502 .iter()
503 .map(|(address, account)| (*address, Some(*account)))
504 .chain(self.destroyed_accounts.iter().map(|address| (*address, None)))
505 .sorted_by_key(|entry| *entry.0)
506 }
507}
508
509#[derive(Clone, Eq, PartialEq, Debug)]
511pub struct HashedStorageSorted {
512 pub non_zero_valued_slots: Vec<(B256, U256)>,
514 pub zero_valued_slots: B256Set,
516 pub wiped: bool,
518}
519
520impl HashedStorageSorted {
521 pub const fn is_wiped(&self) -> bool {
523 self.wiped
524 }
525
526 pub fn storage_slots_sorted(&self) -> impl Iterator<Item = (B256, U256)> {
528 self.non_zero_valued_slots
529 .iter()
530 .map(|(hashed_slot, value)| (*hashed_slot, *value))
531 .chain(self.zero_valued_slots.iter().map(|hashed_slot| (*hashed_slot, U256::ZERO)))
532 .sorted_by_key(|entry| *entry.0)
533 }
534}
535
536#[derive(Debug)]
544pub struct ChunkedHashedPostState {
545 flattened: alloc::vec::IntoIter<(B256, FlattenedHashedPostStateItem)>,
546 size: usize,
547}
548
549#[derive(Debug)]
550enum FlattenedHashedPostStateItem {
551 Account(Option<Account>),
552 StorageWipe,
553 StorageUpdate { slot: B256, value: U256 },
554}
555
556impl ChunkedHashedPostState {
557 fn new(hashed_post_state: HashedPostState, size: usize) -> Self {
558 let flattened = hashed_post_state
559 .storages
560 .into_iter()
561 .flat_map(|(address, storage)| {
562 Some((address, FlattenedHashedPostStateItem::StorageWipe))
564 .filter(|_| storage.wiped)
565 .into_iter()
566 .chain(
567 storage.storage.into_iter().sorted_unstable_by_key(|(slot, _)| *slot).map(
568 move |(slot, value)| {
569 (
570 address,
571 FlattenedHashedPostStateItem::StorageUpdate { slot, value },
572 )
573 },
574 ),
575 )
576 })
577 .chain(hashed_post_state.accounts.into_iter().map(|(address, account)| {
578 (address, FlattenedHashedPostStateItem::Account(account))
579 }))
580 .sorted_by_key(|(address, _)| *address);
585
586 Self { flattened, size }
587 }
588}
589
590impl Iterator for ChunkedHashedPostState {
591 type Item = HashedPostState;
592
593 fn next(&mut self) -> Option<Self::Item> {
594 let mut chunk = HashedPostState::default();
595
596 let mut current_size = 0;
597 while current_size < self.size {
598 let Some((address, item)) = self.flattened.next() else { break };
599
600 match item {
601 FlattenedHashedPostStateItem::Account(account) => {
602 chunk.accounts.insert(address, account);
603 }
604 FlattenedHashedPostStateItem::StorageWipe => {
605 chunk.storages.entry(address).or_default().wiped = true;
606 }
607 FlattenedHashedPostStateItem::StorageUpdate { slot, value } => {
608 chunk.storages.entry(address).or_default().storage.insert(slot, value);
609 }
610 }
611
612 current_size += 1;
613 }
614
615 if chunk.is_empty() {
616 None
617 } else {
618 Some(chunk)
619 }
620 }
621}
622
623#[cfg(test)]
624mod tests {
625 use super::*;
626 use crate::KeccakKeyHasher;
627 use alloy_primitives::Bytes;
628 use revm_database::{states::StorageSlot, StorageWithOriginalValues};
629 use revm_state::{AccountInfo, Bytecode};
630
631 #[test]
632 fn hashed_state_wiped_extension() {
633 let hashed_address = B256::default();
634 let hashed_slot = B256::with_last_byte(64);
635 let hashed_slot2 = B256::with_last_byte(65);
636
637 let original_slot_value = U256::from(123);
639 let mut hashed_state = HashedPostState::default().with_storages([(
640 hashed_address,
641 HashedStorage::from_iter(
642 false,
643 [(hashed_slot, original_slot_value), (hashed_slot2, original_slot_value)],
644 ),
645 )]);
646
647 let updated_slot_value = U256::from(321);
649 let extension = HashedPostState::default().with_storages([(
650 hashed_address,
651 HashedStorage::from_iter(false, [(hashed_slot, updated_slot_value)]),
652 )]);
653 hashed_state.extend(extension);
654
655 let account_storage = hashed_state.storages.get(&hashed_address);
656 assert_eq!(
657 account_storage.and_then(|st| st.storage.get(&hashed_slot)),
658 Some(&updated_slot_value)
659 );
660 assert_eq!(
661 account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
662 Some(&original_slot_value)
663 );
664 assert_eq!(account_storage.map(|st| st.wiped), Some(false));
665
666 let wiped_extension =
668 HashedPostState::default().with_storages([(hashed_address, HashedStorage::new(true))]);
669 hashed_state.extend(wiped_extension);
670
671 let account_storage = hashed_state.storages.get(&hashed_address);
672 assert_eq!(account_storage.map(|st| st.storage.is_empty()), Some(true));
673 assert_eq!(account_storage.map(|st| st.wiped), Some(true));
674
675 hashed_state.extend(HashedPostState::default().with_storages([(
677 hashed_address,
678 HashedStorage::from_iter(false, [(hashed_slot, original_slot_value)]),
679 )]));
680 let account_storage = hashed_state.storages.get(&hashed_address);
681 assert_eq!(
682 account_storage.and_then(|st| st.storage.get(&hashed_slot)),
683 Some(&original_slot_value)
684 );
685 assert_eq!(account_storage.and_then(|st| st.storage.get(&hashed_slot2)), None);
686 assert_eq!(account_storage.map(|st| st.wiped), Some(true));
687
688 hashed_state.extend(HashedPostState::default().with_storages([(
690 hashed_address,
691 HashedStorage::from_iter(false, [(hashed_slot2, updated_slot_value)]),
692 )]));
693 let account_storage = hashed_state.storages.get(&hashed_address);
694 assert_eq!(
695 account_storage.and_then(|st| st.storage.get(&hashed_slot)),
696 Some(&original_slot_value)
697 );
698 assert_eq!(
699 account_storage.and_then(|st| st.storage.get(&hashed_slot2)),
700 Some(&updated_slot_value)
701 );
702 assert_eq!(account_storage.map(|st| st.wiped), Some(true));
703 }
704
705 #[test]
706 fn test_hashed_post_state_from_bundle_state() {
707 let address = Address::random();
709
710 let account_info = AccountInfo {
712 balance: U256::from(123),
713 nonce: 42,
714 code_hash: B256::random(),
715 code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2]))),
716 };
717
718 let mut storage = StorageWithOriginalValues::default();
719 storage.insert(
720 U256::from(1),
721 StorageSlot { present_value: U256::from(4), ..Default::default() },
722 );
723
724 let account = BundleAccount {
726 status: AccountStatus::Changed,
727 info: Some(account_info.clone()),
728 storage,
729 original_info: None,
730 };
731
732 let state = vec![(&address, &account)];
734
735 let hashed_state = HashedPostState::from_bundle_state::<KeccakKeyHasher>(state);
737
738 assert_eq!(hashed_state.accounts.len(), 1);
740 assert_eq!(hashed_state.storages.len(), 1);
741
742 assert_eq!(
744 *hashed_state.accounts.get(&keccak256(address)).unwrap(),
745 Some(account_info.into())
746 );
747 }
748
749 #[test]
750 fn test_hashed_post_state_with_accounts() {
751 let address_1 = Address::random();
753 let address_2 = Address::random();
754
755 let account_info_1 = AccountInfo {
756 balance: U256::from(1000),
757 nonce: 1,
758 code_hash: B256::random(),
759 code: None,
760 };
761
762 let account_1 = (keccak256(address_1), Some(account_info_1.into()));
764 let account_2 = (keccak256(address_2), None);
765
766 let hashed_state = HashedPostState::default().with_accounts(vec![account_1, account_2]);
768
769 assert_eq!(hashed_state.accounts.len(), 2);
771 assert!(hashed_state.accounts.contains_key(&keccak256(address_1)));
772 assert!(hashed_state.accounts.contains_key(&keccak256(address_2)));
773 }
774
775 #[test]
776 fn test_hashed_post_state_with_storages() {
777 let address_1 = Address::random();
779 let address_2 = Address::random();
780
781 let storage_1 = (keccak256(address_1), HashedStorage::new(false));
782 let storage_2 = (keccak256(address_2), HashedStorage::new(true));
783
784 let hashed_state = HashedPostState::default().with_storages(vec![storage_1, storage_2]);
786
787 assert_eq!(hashed_state.storages.len(), 2);
789 assert!(hashed_state.storages.contains_key(&keccak256(address_1)));
790 assert!(hashed_state.storages.contains_key(&keccak256(address_2)));
791 }
792
793 #[test]
794 fn test_hashed_post_state_is_empty() {
795 let empty_state = HashedPostState::default();
797 assert!(empty_state.is_empty());
798
799 let non_empty_state = HashedPostState::default()
801 .with_accounts(vec![(keccak256(Address::random()), Some(Account::default()))]);
802 assert!(!non_empty_state.is_empty());
803 }
804
805 fn create_state_for_multi_proof_targets() -> HashedPostState {
806 let mut state = HashedPostState::default();
807
808 let addr1 = B256::random();
809 let addr2 = B256::random();
810 state.accounts.insert(addr1, Some(Default::default()));
811 state.accounts.insert(addr2, Some(Default::default()));
812
813 let mut storage = HashedStorage::default();
814 let slot1 = B256::random();
815 let slot2 = B256::random();
816 storage.storage.insert(slot1, U256::ZERO);
817 storage.storage.insert(slot2, U256::from(1));
818 state.storages.insert(addr1, storage);
819
820 state
821 }
822
823 #[test]
824 fn test_multi_proof_targets_difference_empty_state() {
825 let state = HashedPostState::default();
826 let excluded = MultiProofTargets::default();
827
828 let targets = state.multi_proof_targets_difference(&excluded);
829 assert!(targets.is_empty());
830 }
831
832 #[test]
833 fn test_multi_proof_targets_difference_new_account_targets() {
834 let state = create_state_for_multi_proof_targets();
835 let excluded = MultiProofTargets::default();
836
837 let targets = state.multi_proof_targets_difference(&excluded);
839 assert_eq!(targets.len(), state.accounts.len());
840 for addr in state.accounts.keys() {
841 assert!(targets.contains_key(addr));
842 }
843 }
844
845 #[test]
846 fn test_multi_proof_targets_difference_new_storage_targets() {
847 let state = create_state_for_multi_proof_targets();
848 let excluded = MultiProofTargets::default();
849
850 let targets = state.multi_proof_targets_difference(&excluded);
851
852 for (addr, storage) in &state.storages {
854 assert!(targets.contains_key(addr));
855 let target_slots = &targets[addr];
856 assert_eq!(target_slots.len(), storage.storage.len());
857 for slot in storage.storage.keys() {
858 assert!(target_slots.contains(slot));
859 }
860 }
861 }
862
863 #[test]
864 fn test_multi_proof_targets_difference_filter_excluded_accounts() {
865 let state = create_state_for_multi_proof_targets();
866 let mut excluded = MultiProofTargets::default();
867
868 let excluded_addr = state
870 .accounts
871 .keys()
872 .find(|&&addr| !state.storages.contains_key(&addr))
873 .expect("Should have an account without storage");
874
875 excluded.insert(*excluded_addr, HashSet::default());
877
878 let targets = state.multi_proof_targets_difference(&excluded);
879
880 assert!(!targets.contains_key(excluded_addr));
882 assert_eq!(targets.len(), state.accounts.len() - 1);
884 }
885
886 #[test]
887 fn test_multi_proof_targets_difference_filter_excluded_storage() {
888 let state = create_state_for_multi_proof_targets();
889 let mut excluded = MultiProofTargets::default();
890
891 let (addr, storage) = state.storages.iter().next().unwrap();
893 let mut excluded_slots = HashSet::default();
894 let excluded_slot = *storage.storage.keys().next().unwrap();
895 excluded_slots.insert(excluded_slot);
896 excluded.insert(*addr, excluded_slots);
897
898 let targets = state.multi_proof_targets_difference(&excluded);
899
900 let target_slots = &targets[addr];
902 assert!(!target_slots.contains(&excluded_slot));
903 assert_eq!(target_slots.len(), storage.storage.len() - 1);
904 }
905
906 #[test]
907 fn test_multi_proof_targets_difference_mixed_excluded_state() {
908 let mut state = HashedPostState::default();
909 let mut excluded = MultiProofTargets::default();
910
911 let addr1 = B256::random();
912 let addr2 = B256::random();
913 let slot1 = B256::random();
914 let slot2 = B256::random();
915
916 state.accounts.insert(addr1, Some(Default::default()));
917 state.accounts.insert(addr2, Some(Default::default()));
918
919 let mut storage = HashedStorage::default();
920 storage.storage.insert(slot1, U256::ZERO);
921 storage.storage.insert(slot2, U256::from(1));
922 state.storages.insert(addr1, storage);
923
924 let mut excluded_slots = HashSet::default();
925 excluded_slots.insert(slot1);
926 excluded.insert(addr1, excluded_slots);
927
928 let targets = state.multi_proof_targets_difference(&excluded);
929
930 assert!(targets.contains_key(&addr2));
931 assert!(!targets[&addr1].contains(&slot1));
932 assert!(targets[&addr1].contains(&slot2));
933 }
934
935 #[test]
936 fn test_multi_proof_targets_difference_unmodified_account_with_storage() {
937 let mut state = HashedPostState::default();
938 let excluded = MultiProofTargets::default();
939
940 let addr = B256::random();
941 let slot1 = B256::random();
942 let slot2 = B256::random();
943
944 let mut storage = HashedStorage::default();
947 storage.storage.insert(slot1, U256::from(1));
948 storage.storage.insert(slot2, U256::from(2));
949 state.storages.insert(addr, storage);
950
951 assert!(!state.accounts.contains_key(&addr));
952 assert!(!excluded.contains_key(&addr));
953
954 let targets = state.multi_proof_targets_difference(&excluded);
955
956 assert!(targets.contains_key(&addr));
958
959 let target_slots = &targets[&addr];
960 assert_eq!(target_slots.len(), 2);
961 assert!(target_slots.contains(&slot1));
962 assert!(target_slots.contains(&slot2));
963 }
964
965 #[test]
966 fn test_partition_by_targets() {
967 let addr1 = B256::random();
968 let addr2 = B256::random();
969 let slot1 = B256::random();
970 let slot2 = B256::random();
971
972 let state = HashedPostState {
973 accounts: B256Map::from_iter([
974 (addr1, Some(Default::default())),
975 (addr2, Some(Default::default())),
976 ]),
977 storages: B256Map::from_iter([(
978 addr1,
979 HashedStorage {
980 wiped: true,
981 storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
982 },
983 )]),
984 };
985 let targets = MultiProofTargets::from_iter([(addr1, HashSet::from_iter([slot1]))]);
986
987 let (with_targets, without_targets) =
988 state.partition_by_targets(&targets, &MultiAddedRemovedKeys::new());
989
990 assert_eq!(
991 with_targets,
992 HashedPostState {
993 accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
994 storages: B256Map::from_iter([(
995 addr1,
996 HashedStorage {
997 wiped: true,
998 storage: B256Map::from_iter([(slot1, U256::ZERO)])
999 }
1000 )]),
1001 }
1002 );
1003 assert_eq!(
1004 without_targets,
1005 HashedPostState {
1006 accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1007 storages: B256Map::from_iter([(
1008 addr1,
1009 HashedStorage {
1010 wiped: false,
1011 storage: B256Map::from_iter([(slot2, U256::from(1))])
1012 }
1013 )]),
1014 }
1015 );
1016 }
1017
1018 #[test]
1019 fn test_chunks() {
1020 let addr1 = B256::from([1; 32]);
1021 let addr2 = B256::from([2; 32]);
1022 let slot1 = B256::from([1; 32]);
1023 let slot2 = B256::from([2; 32]);
1024
1025 let state = HashedPostState {
1026 accounts: B256Map::from_iter([
1027 (addr1, Some(Default::default())),
1028 (addr2, Some(Default::default())),
1029 ]),
1030 storages: B256Map::from_iter([(
1031 addr2,
1032 HashedStorage {
1033 wiped: true,
1034 storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1035 },
1036 )]),
1037 };
1038
1039 let mut chunks = state.chunks(2);
1040 assert_eq!(
1041 chunks.next(),
1042 Some(HashedPostState {
1043 accounts: B256Map::from_iter([(addr1, Some(Default::default()))]),
1044 storages: B256Map::from_iter([(addr2, HashedStorage::new(true)),])
1045 })
1046 );
1047 assert_eq!(
1048 chunks.next(),
1049 Some(HashedPostState {
1050 accounts: B256Map::default(),
1051 storages: B256Map::from_iter([(
1052 addr2,
1053 HashedStorage {
1054 wiped: false,
1055 storage: B256Map::from_iter([(slot1, U256::ZERO), (slot2, U256::from(1))]),
1056 },
1057 )])
1058 })
1059 );
1060 assert_eq!(
1061 chunks.next(),
1062 Some(HashedPostState {
1063 accounts: B256Map::from_iter([(addr2, Some(Default::default()))]),
1064 storages: B256Map::default()
1065 })
1066 );
1067 assert_eq!(chunks.next(), None);
1068 }
1069}