1use super::{
2 changes::{PlainStorageChangeset, StateChangeset},
3 reverts::{AccountInfoRevert, Reverts},
4 AccountRevert, AccountStatus, BundleAccount, PlainStateReverts, RevertToSlot, StorageSlot,
5 TransitionState,
6};
7use bytecode::Bytecode;
8use core::{mem, ops::RangeInclusive};
9use primitives::{
10 hash_map::Entry, Address, HashMap, HashSet, StorageKey, StorageValue, B256, KECCAK_EMPTY,
11};
12use state::AccountInfo;
13use std::{
14 collections::{BTreeMap, BTreeSet},
15 vec::Vec,
16};
17
18#[derive(Debug)]
20pub struct BundleBuilder {
21 states: HashSet<Address>,
22 state_original: HashMap<Address, AccountInfo>,
23 state_present: HashMap<Address, AccountInfo>,
24 state_storage: HashMap<Address, HashMap<StorageKey, (StorageValue, StorageValue)>>,
25
26 reverts: BTreeSet<(u64, Address)>,
27 revert_range: RangeInclusive<u64>,
28 revert_account: HashMap<(u64, Address), Option<Option<AccountInfo>>>,
29 revert_storage: HashMap<(u64, Address), Vec<(StorageKey, StorageValue)>>,
30
31 contracts: HashMap<B256, Bytecode>,
32}
33
34#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
36pub enum OriginalValuesKnown {
37 Yes,
42 No,
48}
49impl OriginalValuesKnown {
50 pub fn is_not_known(&self) -> bool {
52 matches!(self, Self::No)
53 }
54}
55
56impl Default for BundleBuilder {
57 fn default() -> Self {
58 BundleBuilder {
59 states: HashSet::default(),
60 state_original: HashMap::default(),
61 state_present: HashMap::default(),
62 state_storage: HashMap::default(),
63 reverts: BTreeSet::new(),
64 revert_range: 0..=0,
65 revert_account: HashMap::default(),
66 revert_storage: HashMap::default(),
67 contracts: HashMap::default(),
68 }
69 }
70}
71
72impl BundleBuilder {
73 pub fn new(revert_range: RangeInclusive<u64>) -> Self {
77 BundleBuilder {
78 revert_range,
79 ..Default::default()
80 }
81 }
82
83 pub fn apply<F>(self, f: F) -> Self
85 where
86 F: FnOnce(Self) -> Self,
87 {
88 f(self)
89 }
90
91 pub fn apply_mut<F>(&mut self, f: F) -> &mut Self
93 where
94 F: FnOnce(&mut Self),
95 {
96 f(self);
97 self
98 }
99
100 pub fn state_address(mut self, address: Address) -> Self {
102 self.set_state_address(address);
103 self
104 }
105
106 pub fn state_original_account_info(mut self, address: Address, original: AccountInfo) -> Self {
108 self.set_state_original_account_info(address, original);
109 self
110 }
111
112 pub fn state_present_account_info(mut self, address: Address, present: AccountInfo) -> Self {
114 self.set_state_present_account_info(address, present);
115 self
116 }
117
118 pub fn state_storage(
120 mut self,
121 address: Address,
122 storage: HashMap<StorageKey, (StorageValue, StorageValue)>,
123 ) -> Self {
124 self.set_state_storage(address, storage);
125 self
126 }
127
128 pub fn revert_address(mut self, block_number: u64, address: Address) -> Self {
133 self.set_revert_address(block_number, address);
134 self
135 }
136
137 pub fn revert_account_info(
142 mut self,
143 block_number: u64,
144 address: Address,
145 account: Option<Option<AccountInfo>>,
146 ) -> Self {
147 self.set_revert_account_info(block_number, address, account);
148 self
149 }
150
151 pub fn revert_storage(
156 mut self,
157 block_number: u64,
158 address: Address,
159 storage: Vec<(StorageKey, StorageValue)>,
160 ) -> Self {
161 self.set_revert_storage(block_number, address, storage);
162 self
163 }
164
165 pub fn contract(mut self, address: B256, bytecode: Bytecode) -> Self {
167 self.set_contract(address, bytecode);
168 self
169 }
170
171 pub fn set_state_address(&mut self, address: Address) -> &mut Self {
173 self.states.insert(address);
174 self
175 }
176
177 pub fn set_state_original_account_info(
179 &mut self,
180 address: Address,
181 original: AccountInfo,
182 ) -> &mut Self {
183 self.states.insert(address);
184 self.state_original.insert(address, original);
185 self
186 }
187
188 pub fn set_state_present_account_info(
190 &mut self,
191 address: Address,
192 present: AccountInfo,
193 ) -> &mut Self {
194 self.states.insert(address);
195 self.state_present.insert(address, present);
196 self
197 }
198
199 pub fn set_state_storage(
201 &mut self,
202 address: Address,
203 storage: HashMap<StorageKey, (StorageValue, StorageValue)>,
204 ) -> &mut Self {
205 self.states.insert(address);
206 self.state_storage.insert(address, storage);
207 self
208 }
209
210 pub fn set_revert_address(&mut self, block_number: u64, address: Address) -> &mut Self {
212 self.reverts.insert((block_number, address));
213 self
214 }
215
216 pub fn set_revert_account_info(
218 &mut self,
219 block_number: u64,
220 address: Address,
221 account: Option<Option<AccountInfo>>,
222 ) -> &mut Self {
223 self.reverts.insert((block_number, address));
224 self.revert_account.insert((block_number, address), account);
225 self
226 }
227
228 pub fn set_revert_storage(
230 &mut self,
231 block_number: u64,
232 address: Address,
233 storage: Vec<(StorageKey, StorageValue)>,
234 ) -> &mut Self {
235 self.reverts.insert((block_number, address));
236 self.revert_storage.insert((block_number, address), storage);
237 self
238 }
239
240 pub fn set_contract(&mut self, address: B256, bytecode: Bytecode) -> &mut Self {
242 self.contracts.insert(address, bytecode);
243 self
244 }
245
246 pub fn build(mut self) -> BundleState {
248 let mut state_size = 0;
249 let state = self
250 .states
251 .into_iter()
252 .map(|address| {
253 let storage = self
254 .state_storage
255 .remove(&address)
256 .map(|s| {
257 s.into_iter()
258 .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
259 .collect()
260 })
261 .unwrap_or_default();
262 let bundle_account = BundleAccount::new(
263 self.state_original.remove(&address),
264 self.state_present.remove(&address),
265 storage,
266 AccountStatus::Changed,
267 );
268 state_size += bundle_account.size_hint();
269 (address, bundle_account)
270 })
271 .collect();
272
273 let mut reverts_size = 0;
274 let mut reverts_map = BTreeMap::new();
275 for block_number in self.revert_range {
276 reverts_map.insert(block_number, Vec::new());
277 }
278 self.reverts
279 .into_iter()
280 .for_each(|(block_number, address)| {
281 let account = match self
282 .revert_account
283 .remove(&(block_number, address))
284 .unwrap_or_default()
285 {
286 Some(Some(account)) => AccountInfoRevert::RevertTo(account),
287 Some(None) => AccountInfoRevert::DeleteIt,
288 None => AccountInfoRevert::DoNothing,
289 };
290 let storage = self
291 .revert_storage
292 .remove(&(block_number, address))
293 .map(|s| {
294 s.into_iter()
295 .map(|(k, v)| (k, RevertToSlot::Some(v)))
296 .collect()
297 })
298 .unwrap_or_default();
299 let account_revert = AccountRevert {
300 account,
301 storage,
302 previous_status: AccountStatus::Changed,
303 wipe_storage: false,
304 };
305
306 if let Some(vec) = reverts_map.get_mut(&block_number) {
307 reverts_size += account_revert.size_hint();
308 vec.push((address, account_revert));
309 }
310 });
311
312 BundleState {
313 state,
314 contracts: self.contracts,
315 reverts: Reverts::new(reverts_map.into_values().collect()),
316 state_size,
317 reverts_size,
318 }
319 }
320
321 pub fn get_states(&self) -> &HashSet<Address> {
323 &self.states
324 }
325
326 pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
328 &mut self.states
329 }
330
331 pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
333 &mut self.state_original
334 }
335
336 pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
338 &mut self.state_present
339 }
340
341 pub fn get_state_storage_mut(
343 &mut self,
344 ) -> &mut HashMap<Address, HashMap<StorageKey, (StorageValue, StorageValue)>> {
345 &mut self.state_storage
346 }
347
348 pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> {
350 &mut self.reverts
351 }
352
353 pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive<u64> {
355 &mut self.revert_range
356 }
357
358 pub fn get_revert_account_mut(
360 &mut self,
361 ) -> &mut HashMap<(u64, Address), Option<Option<AccountInfo>>> {
362 &mut self.revert_account
363 }
364
365 pub fn get_revert_storage_mut(
367 &mut self,
368 ) -> &mut HashMap<(u64, Address), Vec<(StorageKey, StorageValue)>> {
369 &mut self.revert_storage
370 }
371
372 pub fn get_contracts_mut(&mut self) -> &mut HashMap<B256, Bytecode> {
374 &mut self.contracts
375 }
376}
377
378#[derive(Debug)]
380pub enum BundleRetention {
381 PlainState,
383 Reverts,
385}
386
387impl BundleRetention {
388 pub fn includes_reverts(&self) -> bool {
390 matches!(self, Self::Reverts)
391 }
392}
393
394#[derive(Default, Clone, Debug, PartialEq, Eq)]
404#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
405pub struct BundleState {
406 pub state: HashMap<Address, BundleAccount>,
408 pub contracts: HashMap<B256, Bytecode>,
410 pub reverts: Reverts,
416 pub state_size: usize,
418 pub reverts_size: usize,
420}
421
422impl BundleState {
423 pub fn builder(revert_range: RangeInclusive<u64>) -> BundleBuilder {
425 BundleBuilder::new(revert_range)
426 }
427
428 pub fn new(
430 state: impl IntoIterator<
431 Item = (
432 Address,
433 Option<AccountInfo>,
434 Option<AccountInfo>,
435 HashMap<StorageKey, (StorageValue, StorageValue)>,
436 ),
437 >,
438 reverts: impl IntoIterator<
439 Item = impl IntoIterator<
440 Item = (
441 Address,
442 Option<Option<AccountInfo>>,
443 impl IntoIterator<Item = (StorageKey, StorageValue)>,
444 ),
445 >,
446 >,
447 contracts: impl IntoIterator<Item = (B256, Bytecode)>,
448 ) -> Self {
449 let mut state_size = 0;
451 let state = state
452 .into_iter()
453 .map(|(address, original, present, storage)| {
454 let account = BundleAccount::new(
455 original,
456 present,
457 storage
458 .into_iter()
459 .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
460 .collect(),
461 AccountStatus::Changed,
462 );
463 state_size += account.size_hint();
464 (address, account)
465 })
466 .collect();
467
468 let mut reverts_size = 0;
470 let reverts = reverts
471 .into_iter()
472 .map(|block_reverts| {
473 block_reverts
474 .into_iter()
475 .map(|(address, account, storage)| {
476 let account = match account {
477 Some(Some(account)) => AccountInfoRevert::RevertTo(account),
478 Some(None) => AccountInfoRevert::DeleteIt,
479 None => AccountInfoRevert::DoNothing,
480 };
481 let revert = AccountRevert {
482 account,
483 storage: storage
484 .into_iter()
485 .map(|(k, v)| (k, RevertToSlot::Some(v)))
486 .collect(),
487 previous_status: AccountStatus::Changed,
488 wipe_storage: false,
489 };
490 reverts_size += revert.size_hint();
491 (address, revert)
492 })
493 .collect::<Vec<_>>()
494 })
495 .collect::<Vec<_>>();
496
497 Self {
498 state,
499 contracts: contracts.into_iter().collect(),
500 reverts: Reverts::new(reverts),
501 state_size,
502 reverts_size,
503 }
504 }
505
506 pub fn size_hint(&self) -> usize {
511 self.state_size + self.reverts_size + self.contracts.len()
512 }
513
514 pub fn state(&self) -> &HashMap<Address, BundleAccount> {
516 &self.state
517 }
518
519 pub fn is_empty(&self) -> bool {
521 self.len() == 0
522 }
523
524 pub fn len(&self) -> usize {
526 self.state.len()
527 }
528
529 pub fn account(&self, address: &Address) -> Option<&BundleAccount> {
531 self.state.get(address)
532 }
533
534 pub fn bytecode(&self, hash: &B256) -> Option<Bytecode> {
536 self.contracts.get(hash).cloned()
537 }
538
539 pub fn apply_transitions_and_create_reverts(
545 &mut self,
546 transitions: TransitionState,
547 retention: BundleRetention,
548 ) {
549 let include_reverts = retention.includes_reverts();
550 let reverts_capacity = if include_reverts {
552 transitions.transitions.len()
553 } else {
554 0
555 };
556 let mut reverts = Vec::with_capacity(reverts_capacity);
557
558 for (address, transition) in transitions.transitions.into_iter() {
559 if let Some((hash, new_bytecode)) = transition.has_new_contract() {
561 self.contracts.insert(hash, new_bytecode.clone());
562 }
563 let revert = match self.state.entry(address) {
565 Entry::Occupied(mut entry) => {
566 let entry = entry.get_mut();
567 self.state_size -= entry.size_hint();
568 let revert = entry.update_and_create_revert(transition);
570 self.state_size += entry.size_hint();
572 revert
573 }
574 Entry::Vacant(entry) => {
575 let present_bundle = transition.present_bundle_account();
577 let revert = transition.create_revert();
578 if revert.is_some() {
579 self.state_size += present_bundle.size_hint();
580 entry.insert(present_bundle);
581 }
582 revert
583 }
584 };
585
586 if let Some(revert) = revert.filter(|_| include_reverts) {
588 self.reverts_size += revert.size_hint();
589 reverts.push((address, revert));
590 }
591 }
592
593 self.reverts.push(reverts);
594 }
595
596 pub fn to_plain_state(&self, is_value_known: OriginalValuesKnown) -> StateChangeset {
599 let state_len = self.state.len();
601 let mut accounts = Vec::with_capacity(state_len);
602 let mut storage = Vec::with_capacity(state_len);
603
604 for (address, account) in &self.state {
605 let was_destroyed = account.was_destroyed();
607 if is_value_known.is_not_known() || account.is_info_changed() {
608 let info = account.info.as_ref().map(AccountInfo::copy_without_code);
609 accounts.push((*address, info));
610 }
611
612 let mut account_storage_changed = Vec::with_capacity(account.storage.len());
617
618 for (&key, &slot) in account.storage.iter() {
619 let destroyed_and_not_zero = was_destroyed && !slot.present_value.is_zero();
622
623 let not_destroyed_and_changed = !was_destroyed && slot.is_changed();
626
627 if is_value_known.is_not_known()
628 || destroyed_and_not_zero
629 || not_destroyed_and_changed
630 {
631 account_storage_changed.push((key, slot.present_value));
632 }
633 }
634
635 if !account_storage_changed.is_empty() || was_destroyed {
636 storage.push(PlainStorageChangeset {
638 address: *address,
639 wipe_storage: was_destroyed,
640 storage: account_storage_changed,
641 });
642 }
643 }
644
645 let contracts = self
646 .contracts
647 .iter()
648 .filter(|(b, _)| **b != KECCAK_EMPTY)
650 .map(|(b, code)| (*b, code.clone()))
651 .collect::<Vec<_>>();
652 StateChangeset {
653 accounts,
654 storage,
655 contracts,
656 }
657 }
658
659 pub fn to_plain_state_and_reverts(
662 &self,
663 is_value_known: OriginalValuesKnown,
664 ) -> (StateChangeset, PlainStateReverts) {
665 (
666 self.to_plain_state(is_value_known),
667 self.reverts.to_plain_state_reverts(),
668 )
669 }
670
671 #[deprecated = "Use `to_plain_state_and_reverts` instead"]
674 pub fn into_plain_state_and_reverts(
675 self,
676 is_value_known: OriginalValuesKnown,
677 ) -> (StateChangeset, PlainStateReverts) {
678 self.to_plain_state_and_reverts(is_value_known)
679 }
680
681 pub fn extend_state(&mut self, other_state: HashMap<Address, BundleAccount>) {
685 for (address, other_account) in other_state {
686 match self.state.entry(address) {
687 Entry::Occupied(mut entry) => {
688 let this = entry.get_mut();
689 self.state_size -= this.size_hint();
690
691 if other_account.was_destroyed() {
694 this.storage = other_account.storage;
695 } else {
696 for (key, storage_slot) in other_account.storage {
698 this.storage
700 .entry(key)
701 .or_insert(storage_slot)
702 .present_value = storage_slot.present_value;
703 }
704 }
705 this.info = other_account.info;
706 this.status.transition(other_account.status);
707
708 self.state_size += this.size_hint();
710 }
711 Entry::Vacant(entry) => {
712 self.state_size += other_account.size_hint();
714 entry.insert(other_account);
715 }
716 }
717 }
718 }
719 pub fn extend(&mut self, mut other: Self) {
730 for (address, revert) in other.reverts.iter_mut().flatten() {
733 if revert.wipe_storage {
734 if let Some(this_account) = self.state.get_mut(address) {
737 for (key, value) in this_account.storage.drain() {
740 revert
741 .storage
742 .entry(key)
743 .or_insert(RevertToSlot::Some(value.present_value));
744 }
745
746 if this_account.was_destroyed() {
748 revert.wipe_storage = false;
749 }
750 }
751 }
752
753 self.reverts_size += revert.size_hint();
755 }
756 self.extend_state(other.state);
758 self.contracts.extend(other.contracts);
760 self.reverts.extend(other.reverts);
762 }
763
764 pub fn take_n_reverts(&mut self, reverts_to_take: usize) -> Reverts {
766 if reverts_to_take > self.reverts.len() {
768 return self.take_all_reverts();
769 }
770 let (detach, this) = self.reverts.split_at(reverts_to_take);
771 let detached_reverts = Reverts::new(detach.to_vec());
772 self.reverts_size = this
773 .iter()
774 .flatten()
775 .fold(0, |acc, (_, revert)| acc + revert.size_hint());
776 self.reverts = Reverts::new(this.to_vec());
777 detached_reverts
778 }
779
780 pub fn take_all_reverts(&mut self) -> Reverts {
782 self.reverts_size = 0;
783 mem::take(&mut self.reverts)
784 }
785
786 pub fn revert_latest(&mut self) -> bool {
792 if let Some(reverts) = self.reverts.pop() {
794 for (address, revert_account) in reverts.into_iter() {
795 self.reverts_size -= revert_account.size_hint();
796 match self.state.entry(address) {
797 Entry::Occupied(mut entry) => {
798 let account = entry.get_mut();
799 self.state_size -= account.size_hint();
800 if account.revert(revert_account) {
801 entry.remove();
802 } else {
803 self.state_size += account.size_hint();
804 }
805 }
806 Entry::Vacant(entry) => {
807 let mut account = BundleAccount::new(
810 None,
811 None,
812 HashMap::default(),
813 AccountStatus::LoadedNotExisting,
814 );
815 if !account.revert(revert_account) {
816 self.state_size += account.size_hint();
817 entry.insert(account);
818 }
819 }
820 }
821 }
822 return true;
823 }
824
825 false
826 }
827
828 pub fn revert(&mut self, mut num_transitions: usize) {
832 if num_transitions == 0 {
833 return;
834 }
835
836 while self.revert_latest() {
837 num_transitions -= 1;
838 if num_transitions == 0 {
839 break;
841 }
842 }
843 }
844
845 pub fn prepend_state(&mut self, mut other: BundleState) {
851 let this_bundle = mem::take(self);
853 other.extend_state(this_bundle.state);
855 other.contracts.extend(this_bundle.contracts);
857 mem::swap(self, &mut other)
859 }
860}
861
862#[cfg(test)]
863mod tests {
864 use super::*;
865 use crate::{StorageWithOriginalValues, TransitionAccount};
866 use primitives::U256;
867
868 #[test]
869 fn transition_states() {
870 let address = Address::new([0x01; 20]);
872 let acc1 = AccountInfo {
873 balance: U256::from(10),
874 nonce: 1,
875 code_hash: KECCAK_EMPTY,
876 code: None,
877 };
878
879 let mut bundle_state = BundleState::default();
880
881 let transition = TransitionAccount {
884 info: Some(acc1),
885 status: AccountStatus::InMemoryChange,
886 previous_info: None,
887 previous_status: AccountStatus::LoadedNotExisting,
888 storage: StorageWithOriginalValues::default(),
889 storage_was_destroyed: false,
890 };
891
892 bundle_state.apply_transitions_and_create_reverts(
894 TransitionState::single(address, transition.clone()),
895 BundleRetention::Reverts,
896 );
897 }
898
899 const fn account1() -> Address {
900 Address::new([0x60; 20])
901 }
902
903 const fn account2() -> Address {
904 Address::new([0x61; 20])
905 }
906
907 fn slot1() -> StorageKey {
908 StorageKey::from(5)
909 }
910
911 fn slot2() -> StorageKey {
912 StorageKey::from(7)
913 }
914
915 fn test_bundle1() -> BundleState {
917 BundleState::new(
919 vec![
920 (
921 account1(),
922 None,
923 Some(AccountInfo {
924 nonce: 1,
925 balance: U256::from(10),
926 code_hash: KECCAK_EMPTY,
927 code: None,
928 }),
929 HashMap::from_iter([
930 (slot1(), (StorageValue::from(0), StorageValue::from(10))),
931 (slot2(), (StorageValue::from(0), StorageValue::from(15))),
932 ]),
933 ),
934 (
935 account2(),
936 None,
937 Some(AccountInfo {
938 nonce: 1,
939 balance: U256::from(10),
940 code_hash: KECCAK_EMPTY,
941 code: None,
942 }),
943 HashMap::default(),
944 ),
945 ],
946 vec![vec![
947 (
948 account1(),
949 Some(None),
950 vec![
951 (slot1(), StorageValue::from(0)),
952 (slot2(), StorageValue::from(0)),
953 ],
954 ),
955 (account2(), Some(None), vec![]),
956 ]],
957 vec![],
958 )
959 }
960
961 fn test_bundle2() -> BundleState {
963 BundleState::new(
965 vec![(
966 account1(),
967 None,
968 Some(AccountInfo {
969 nonce: 3,
970 balance: U256::from(20),
971 code_hash: KECCAK_EMPTY,
972 code: None,
973 }),
974 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
975 )],
976 vec![vec![(
977 account1(),
978 Some(Some(AccountInfo {
979 nonce: 1,
980 balance: U256::from(10),
981 code_hash: KECCAK_EMPTY,
982 code: None,
983 })),
984 vec![(slot1(), StorageValue::from(10))],
985 )]],
986 vec![],
987 )
988 }
989
990 fn test_bundle3() -> BundleState {
992 BundleState::builder(0..=0)
993 .state_present_account_info(
994 account1(),
995 AccountInfo {
996 nonce: 1,
997 balance: U256::from(10),
998 code_hash: KECCAK_EMPTY,
999 code: None,
1000 },
1001 )
1002 .state_storage(
1003 account1(),
1004 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(10)))]),
1005 )
1006 .state_address(account2())
1007 .state_present_account_info(
1008 account2(),
1009 AccountInfo {
1010 nonce: 1,
1011 balance: U256::from(10),
1012 code_hash: KECCAK_EMPTY,
1013 code: None,
1014 },
1015 )
1016 .revert_address(0, account1())
1017 .revert_account_info(0, account1(), Some(None))
1018 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(0))])
1019 .revert_account_info(0, account2(), Some(None))
1020 .build()
1021 }
1022
1023 fn test_bundle4() -> BundleState {
1025 BundleState::builder(0..=0)
1026 .state_present_account_info(
1027 account1(),
1028 AccountInfo {
1029 nonce: 3,
1030 balance: U256::from(20),
1031 code_hash: KECCAK_EMPTY,
1032 code: None,
1033 },
1034 )
1035 .state_storage(
1036 account1(),
1037 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
1038 )
1039 .revert_address(0, account1())
1040 .revert_account_info(
1041 0,
1042 account1(),
1043 Some(Some(AccountInfo {
1044 nonce: 1,
1045 balance: U256::from(10),
1046 code_hash: KECCAK_EMPTY,
1047 code: None,
1048 })),
1049 )
1050 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1051 .build()
1052 }
1053
1054 fn sanity_path(bundle1: BundleState, bundle2: BundleState) {
1055 let mut extended = bundle1.clone();
1056 extended.extend(bundle2.clone());
1057
1058 let mut reverted = extended.clone();
1059 reverted.revert(0);
1061 assert_eq!(reverted, extended);
1062
1063 reverted.revert(1);
1065 assert_eq!(reverted, bundle1);
1066
1067 reverted.revert(1);
1069 assert_eq!(reverted, BundleState::default());
1070
1071 let mut reverted = extended.clone();
1072
1073 reverted.revert(10);
1075 assert_eq!(reverted, BundleState::default());
1076 }
1077
1078 #[test]
1079 fn extend_on_destroyed_values() {
1080 let base_bundle1 = test_bundle1();
1081 let base_bundle2 = test_bundle2();
1082
1083 let mut b1 = base_bundle1.clone();
1088 let mut b2 = base_bundle2.clone();
1089 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1090 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1091 b1.extend(b2);
1092 assert_eq!(
1093 b1.state.get_mut(&account1()).unwrap().status,
1094 AccountStatus::DestroyedChanged
1095 );
1096
1097 let mut b1 = base_bundle1.clone();
1102 let mut b2 = base_bundle2.clone();
1103 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1104 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1105 b2.reverts[0][0].1.wipe_storage = true;
1106 b1.extend(b2);
1107 assert_eq!(
1108 b1.state.get_mut(&account1()).unwrap().status,
1109 AccountStatus::Destroyed
1110 );
1111
1112 let mut revert1 = base_bundle2.reverts[0][0].clone();
1115 revert1.1.wipe_storage = true;
1116 revert1
1117 .1
1118 .storage
1119 .insert(slot2(), RevertToSlot::Some(StorageValue::from(15)));
1120
1121 assert_eq!(
1122 b1.reverts.as_ref(),
1123 vec![base_bundle1.reverts[0].clone(), vec![revert1]],
1124 );
1125
1126 let mut b1 = base_bundle1.clone();
1132 let mut b2 = base_bundle2.clone();
1133 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::InMemoryChange;
1134 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1135 b1.extend(b2);
1136 assert_eq!(
1137 b1.state.get_mut(&account1()).unwrap().status,
1138 AccountStatus::InMemoryChange
1139 );
1140 }
1141
1142 #[test]
1143 fn test_sanity_path() {
1144 sanity_path(test_bundle1(), test_bundle2());
1145 sanity_path(test_bundle3(), test_bundle4());
1146 }
1147
1148 #[test]
1149 fn test_multi_reverts_with_delete() {
1150 let mut state = BundleBuilder::new(0..=3)
1151 .revert_address(0, account1())
1152 .revert_account_info(2, account1(), Some(Some(AccountInfo::default())))
1153 .revert_account_info(3, account1(), Some(None))
1154 .build();
1155
1156 state.revert_latest();
1157 assert_eq!(state.state.get(&account1()), None);
1159
1160 state.revert_latest();
1161 assert_eq!(
1163 state.state.get(&account1()),
1164 Some(&BundleAccount::new(
1165 None,
1166 Some(AccountInfo::default()),
1167 HashMap::default(),
1168 AccountStatus::Changed
1169 ))
1170 );
1171 }
1172
1173 #[test]
1174 fn test_revert_capacity() {
1175 let state = BundleState::builder(0..=3)
1176 .revert_address(0, account1())
1177 .revert_address(2, account2())
1178 .revert_account_info(0, account1(), Some(None))
1179 .revert_account_info(2, account2(), None)
1180 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1181 .build();
1182
1183 assert_eq!(state.reverts.len(), 4);
1184 assert_eq!(state.reverts[1], vec![]);
1185 assert_eq!(state.reverts[3], vec![]);
1186 assert_eq!(state.reverts[0].len(), 1);
1187 assert_eq!(state.reverts[2].len(), 1);
1188
1189 let (addr1, revert1) = &state.reverts[0][0];
1190 assert_eq!(addr1, &account1());
1191 assert_eq!(revert1.account, AccountInfoRevert::DeleteIt);
1192
1193 let (addr2, revert2) = &state.reverts[2][0];
1194 assert_eq!(addr2, &account2());
1195 assert_eq!(revert2.account, AccountInfoRevert::DoNothing);
1196 }
1197
1198 #[test]
1199 fn take_reverts() {
1200 let bundle1 = test_bundle1();
1201 let bundle2 = test_bundle2();
1202
1203 let mut extended = bundle1.clone();
1204 extended.extend(bundle2.clone());
1205
1206 assert_eq!(extended.reverts.len(), 2);
1208
1209 let mut extended2 = extended.clone();
1211 assert_eq!(extended2.take_n_reverts(100), extended.reverts);
1212
1213 let mut extended2 = extended.clone();
1215 assert_eq!(extended2.take_all_reverts(), extended.reverts);
1216
1217 let taken_reverts = extended.take_n_reverts(0);
1219 assert_eq!(taken_reverts, Reverts::default());
1220 assert_eq!(extended.reverts.len(), 2);
1221
1222 let taken_reverts = extended.take_n_reverts(1);
1224 assert_eq!(taken_reverts, bundle1.reverts);
1225
1226 let taken_reverts = extended.take_n_reverts(1);
1228 assert_eq!(taken_reverts, bundle2.reverts);
1229 }
1230
1231 #[test]
1232 fn prepend_state() {
1233 let address1 = account1();
1234 let address2 = account2();
1235
1236 let account1 = AccountInfo {
1237 nonce: 1,
1238 ..Default::default()
1239 };
1240 let account1_changed = AccountInfo {
1241 nonce: 1,
1242 ..Default::default()
1243 };
1244 let account2 = AccountInfo {
1245 nonce: 1,
1246 ..Default::default()
1247 };
1248
1249 let present_state = BundleState::builder(2..=2)
1250 .state_present_account_info(address1, account1_changed.clone())
1251 .build();
1252 assert_eq!(present_state.reverts.len(), 1);
1253 let previous_state = BundleState::builder(1..=1)
1254 .state_present_account_info(address1, account1)
1255 .state_present_account_info(address2, account2.clone())
1256 .build();
1257 assert_eq!(previous_state.reverts.len(), 1);
1258
1259 let mut test = present_state;
1260
1261 test.prepend_state(previous_state);
1262
1263 assert_eq!(test.state.len(), 2);
1264 assert_eq!(test.reverts.len(), 1);
1266 assert_eq!(
1268 test.state.get(&address1).unwrap().info,
1269 Some(account1_changed)
1270 );
1271 assert_eq!(test.state.get(&address2).unwrap().info, Some(account2));
1273 }
1274
1275 #[test]
1276 fn test_getters() {
1277 let mut builder = BundleBuilder::new(0..=3);
1278
1279 assert!(builder.get_states().is_empty());
1281 builder.get_states_mut().insert(account1());
1282 assert!(builder.get_states().contains(&account1()));
1283
1284 assert!(builder.get_state_original_mut().is_empty());
1286 builder
1287 .get_state_original_mut()
1288 .insert(account1(), AccountInfo::default());
1289 assert!(builder.get_state_original_mut().contains_key(&account1()));
1290
1291 assert!(builder.get_state_present_mut().is_empty());
1293 builder
1294 .get_state_present_mut()
1295 .insert(account1(), AccountInfo::default());
1296 assert!(builder.get_state_present_mut().contains_key(&account1()));
1297
1298 assert!(builder.get_state_storage_mut().is_empty());
1300 builder
1301 .get_state_storage_mut()
1302 .insert(account1(), HashMap::default());
1303 assert!(builder.get_state_storage_mut().contains_key(&account1()));
1304
1305 assert!(builder.get_reverts_mut().is_empty());
1307 builder.get_reverts_mut().insert((0, account1()));
1308 assert!(builder.get_reverts_mut().contains(&(0, account1())));
1309
1310 assert_eq!(builder.get_revert_range_mut().clone(), 0..=3);
1312
1313 assert!(builder.get_revert_account_mut().is_empty());
1315 builder
1316 .get_revert_account_mut()
1317 .insert((0, account1()), Some(None));
1318 assert!(builder
1319 .get_revert_account_mut()
1320 .contains_key(&(0, account1())));
1321
1322 assert!(builder.get_revert_storage_mut().is_empty());
1324 builder
1325 .get_revert_storage_mut()
1326 .insert((0, account1()), vec![(slot1(), StorageValue::from(0))]);
1327 assert!(builder
1328 .get_revert_storage_mut()
1329 .contains_key(&(0, account1())));
1330
1331 assert!(builder.get_contracts_mut().is_empty());
1333 builder
1334 .get_contracts_mut()
1335 .insert(B256::default(), Bytecode::default());
1336 assert!(builder.get_contracts_mut().contains_key(&B256::default()));
1337 }
1338}