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::{hash_map::Entry, Address, HashMap, HashSet, B256, KECCAK_EMPTY, U256};
10use state::AccountInfo;
11use std::{
12 collections::{BTreeMap, BTreeSet},
13 vec::Vec,
14};
15
16#[derive(Debug)]
18pub struct BundleBuilder {
19 states: HashSet<Address>,
20 state_original: HashMap<Address, AccountInfo>,
21 state_present: HashMap<Address, AccountInfo>,
22 state_storage: HashMap<Address, HashMap<U256, (U256, U256)>>,
23
24 reverts: BTreeSet<(u64, Address)>,
25 revert_range: RangeInclusive<u64>,
26 revert_account: HashMap<(u64, Address), Option<Option<AccountInfo>>>,
27 revert_storage: HashMap<(u64, Address), Vec<(U256, U256)>>,
28
29 contracts: HashMap<B256, Bytecode>,
30}
31
32#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
34pub enum OriginalValuesKnown {
35 Yes,
40 No,
46}
47impl OriginalValuesKnown {
48 pub fn is_not_known(&self) -> bool {
50 matches!(self, Self::No)
51 }
52}
53
54impl Default for BundleBuilder {
55 fn default() -> Self {
56 BundleBuilder {
57 states: HashSet::default(),
58 state_original: HashMap::default(),
59 state_present: HashMap::default(),
60 state_storage: HashMap::default(),
61 reverts: BTreeSet::new(),
62 revert_range: 0..=0,
63 revert_account: HashMap::default(),
64 revert_storage: HashMap::default(),
65 contracts: HashMap::default(),
66 }
67 }
68}
69
70impl BundleBuilder {
71 pub fn new(revert_range: RangeInclusive<u64>) -> Self {
75 BundleBuilder {
76 revert_range,
77 ..Default::default()
78 }
79 }
80
81 pub fn apply<F>(self, f: F) -> Self
83 where
84 F: FnOnce(Self) -> Self,
85 {
86 f(self)
87 }
88
89 pub fn apply_mut<F>(&mut self, f: F) -> &mut Self
91 where
92 F: FnOnce(&mut Self),
93 {
94 f(self);
95 self
96 }
97
98 pub fn state_address(mut self, address: Address) -> Self {
100 self.set_state_address(address);
101 self
102 }
103
104 pub fn state_original_account_info(mut self, address: Address, original: AccountInfo) -> Self {
106 self.set_state_original_account_info(address, original);
107 self
108 }
109
110 pub fn state_present_account_info(mut self, address: Address, present: AccountInfo) -> Self {
112 self.set_state_present_account_info(address, present);
113 self
114 }
115
116 pub fn state_storage(mut self, address: Address, storage: HashMap<U256, (U256, U256)>) -> Self {
118 self.set_state_storage(address, storage);
119 self
120 }
121
122 pub fn revert_address(mut self, block_number: u64, address: Address) -> Self {
127 self.set_revert_address(block_number, address);
128 self
129 }
130
131 pub fn revert_account_info(
136 mut self,
137 block_number: u64,
138 address: Address,
139 account: Option<Option<AccountInfo>>,
140 ) -> Self {
141 self.set_revert_account_info(block_number, address, account);
142 self
143 }
144
145 pub fn revert_storage(
150 mut self,
151 block_number: u64,
152 address: Address,
153 storage: Vec<(U256, U256)>,
154 ) -> Self {
155 self.set_revert_storage(block_number, address, storage);
156 self
157 }
158
159 pub fn contract(mut self, address: B256, bytecode: Bytecode) -> Self {
161 self.set_contract(address, bytecode);
162 self
163 }
164
165 pub fn set_state_address(&mut self, address: Address) -> &mut Self {
167 self.states.insert(address);
168 self
169 }
170
171 pub fn set_state_original_account_info(
173 &mut self,
174 address: Address,
175 original: AccountInfo,
176 ) -> &mut Self {
177 self.states.insert(address);
178 self.state_original.insert(address, original);
179 self
180 }
181
182 pub fn set_state_present_account_info(
184 &mut self,
185 address: Address,
186 present: AccountInfo,
187 ) -> &mut Self {
188 self.states.insert(address);
189 self.state_present.insert(address, present);
190 self
191 }
192
193 pub fn set_state_storage(
195 &mut self,
196 address: Address,
197 storage: HashMap<U256, (U256, U256)>,
198 ) -> &mut Self {
199 self.states.insert(address);
200 self.state_storage.insert(address, storage);
201 self
202 }
203
204 pub fn set_revert_address(&mut self, block_number: u64, address: Address) -> &mut Self {
206 self.reverts.insert((block_number, address));
207 self
208 }
209
210 pub fn set_revert_account_info(
212 &mut self,
213 block_number: u64,
214 address: Address,
215 account: Option<Option<AccountInfo>>,
216 ) -> &mut Self {
217 self.reverts.insert((block_number, address));
218 self.revert_account.insert((block_number, address), account);
219 self
220 }
221
222 pub fn set_revert_storage(
224 &mut self,
225 block_number: u64,
226 address: Address,
227 storage: Vec<(U256, U256)>,
228 ) -> &mut Self {
229 self.reverts.insert((block_number, address));
230 self.revert_storage.insert((block_number, address), storage);
231 self
232 }
233
234 pub fn set_contract(&mut self, address: B256, bytecode: Bytecode) -> &mut Self {
236 self.contracts.insert(address, bytecode);
237 self
238 }
239
240 pub fn build(mut self) -> BundleState {
242 let mut state_size = 0;
243 let state = self
244 .states
245 .into_iter()
246 .map(|address| {
247 let storage = self
248 .state_storage
249 .remove(&address)
250 .map(|s| {
251 s.into_iter()
252 .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
253 .collect()
254 })
255 .unwrap_or_default();
256 let bundle_account = BundleAccount::new(
257 self.state_original.remove(&address),
258 self.state_present.remove(&address),
259 storage,
260 AccountStatus::Changed,
261 );
262 state_size += bundle_account.size_hint();
263 (address, bundle_account)
264 })
265 .collect();
266
267 let mut reverts_size = 0;
268 let mut reverts_map = BTreeMap::new();
269 for block_number in self.revert_range {
270 reverts_map.insert(block_number, Vec::new());
271 }
272 self.reverts
273 .into_iter()
274 .for_each(|(block_number, address)| {
275 let account = match self
276 .revert_account
277 .remove(&(block_number, address))
278 .unwrap_or_default()
279 {
280 Some(Some(account)) => AccountInfoRevert::RevertTo(account),
281 Some(None) => AccountInfoRevert::DeleteIt,
282 None => AccountInfoRevert::DoNothing,
283 };
284 let storage = self
285 .revert_storage
286 .remove(&(block_number, address))
287 .map(|s| {
288 s.into_iter()
289 .map(|(k, v)| (k, RevertToSlot::Some(v)))
290 .collect()
291 })
292 .unwrap_or_default();
293 let account_revert = AccountRevert {
294 account,
295 storage,
296 previous_status: AccountStatus::Changed,
297 wipe_storage: false,
298 };
299
300 if reverts_map.contains_key(&block_number) {
301 reverts_size += account_revert.size_hint();
302 reverts_map
303 .entry(block_number)
304 .or_insert(Vec::new())
305 .push((address, account_revert));
306 }
307 });
308
309 BundleState {
310 state,
311 contracts: self.contracts,
312 reverts: Reverts::new(reverts_map.into_values().collect()),
313 state_size,
314 reverts_size,
315 }
316 }
317
318 pub fn get_states(&self) -> &HashSet<Address> {
320 &self.states
321 }
322
323 pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
325 &mut self.states
326 }
327
328 pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
330 &mut self.state_original
331 }
332
333 pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
335 &mut self.state_present
336 }
337
338 pub fn get_state_storage_mut(&mut self) -> &mut HashMap<Address, HashMap<U256, (U256, U256)>> {
340 &mut self.state_storage
341 }
342
343 pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> {
345 &mut self.reverts
346 }
347
348 pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive<u64> {
350 &mut self.revert_range
351 }
352
353 pub fn get_revert_account_mut(
355 &mut self,
356 ) -> &mut HashMap<(u64, Address), Option<Option<AccountInfo>>> {
357 &mut self.revert_account
358 }
359
360 pub fn get_revert_storage_mut(&mut self) -> &mut HashMap<(u64, Address), Vec<(U256, U256)>> {
362 &mut self.revert_storage
363 }
364
365 pub fn get_contracts_mut(&mut self) -> &mut HashMap<B256, Bytecode> {
367 &mut self.contracts
368 }
369}
370
371#[derive(Debug)]
373pub enum BundleRetention {
374 PlainState,
376 Reverts,
378}
379
380impl BundleRetention {
381 pub fn includes_reverts(&self) -> bool {
383 matches!(self, Self::Reverts)
384 }
385}
386
387#[derive(Default, Clone, Debug, PartialEq, Eq)]
397#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
398pub struct BundleState {
399 pub state: HashMap<Address, BundleAccount>,
401 pub contracts: HashMap<B256, Bytecode>,
403 pub reverts: Reverts,
409 pub state_size: usize,
411 pub reverts_size: usize,
413}
414
415impl BundleState {
416 pub fn builder(revert_range: RangeInclusive<u64>) -> BundleBuilder {
418 BundleBuilder::new(revert_range)
419 }
420
421 pub fn new(
423 state: impl IntoIterator<
424 Item = (
425 Address,
426 Option<AccountInfo>,
427 Option<AccountInfo>,
428 HashMap<U256, (U256, U256)>,
429 ),
430 >,
431 reverts: impl IntoIterator<
432 Item = impl IntoIterator<
433 Item = (
434 Address,
435 Option<Option<AccountInfo>>,
436 impl IntoIterator<Item = (U256, U256)>,
437 ),
438 >,
439 >,
440 contracts: impl IntoIterator<Item = (B256, Bytecode)>,
441 ) -> Self {
442 let mut state_size = 0;
444 let state = state
445 .into_iter()
446 .map(|(address, original, present, storage)| {
447 let account = BundleAccount::new(
448 original,
449 present,
450 storage
451 .into_iter()
452 .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
453 .collect(),
454 AccountStatus::Changed,
455 );
456 state_size += account.size_hint();
457 (address, account)
458 })
459 .collect();
460
461 let mut reverts_size = 0;
463 let reverts = reverts
464 .into_iter()
465 .map(|block_reverts| {
466 block_reverts
467 .into_iter()
468 .map(|(address, account, storage)| {
469 let account = match account {
470 Some(Some(account)) => AccountInfoRevert::RevertTo(account),
471 Some(None) => AccountInfoRevert::DeleteIt,
472 None => AccountInfoRevert::DoNothing,
473 };
474 let revert = AccountRevert {
475 account,
476 storage: storage
477 .into_iter()
478 .map(|(k, v)| (k, RevertToSlot::Some(v)))
479 .collect(),
480 previous_status: AccountStatus::Changed,
481 wipe_storage: false,
482 };
483 reverts_size += revert.size_hint();
484 (address, revert)
485 })
486 .collect::<Vec<_>>()
487 })
488 .collect::<Vec<_>>();
489
490 Self {
491 state,
492 contracts: contracts.into_iter().collect(),
493 reverts: Reverts::new(reverts),
494 state_size,
495 reverts_size,
496 }
497 }
498
499 pub fn size_hint(&self) -> usize {
504 self.state_size + self.reverts_size + self.contracts.len()
505 }
506
507 pub fn state(&self) -> &HashMap<Address, BundleAccount> {
509 &self.state
510 }
511
512 pub fn is_empty(&self) -> bool {
514 self.len() == 0
515 }
516
517 pub fn len(&self) -> usize {
519 self.state.len()
520 }
521
522 pub fn account(&self, address: &Address) -> Option<&BundleAccount> {
524 self.state.get(address)
525 }
526
527 pub fn bytecode(&self, hash: &B256) -> Option<Bytecode> {
529 self.contracts.get(hash).cloned()
530 }
531
532 pub fn apply_transitions_and_create_reverts(
538 &mut self,
539 transitions: TransitionState,
540 retention: BundleRetention,
541 ) {
542 let include_reverts = retention.includes_reverts();
543 let reverts_capacity = if include_reverts {
545 transitions.transitions.len()
546 } else {
547 0
548 };
549 let mut reverts = Vec::with_capacity(reverts_capacity);
550
551 for (address, transition) in transitions.transitions.into_iter() {
552 if let Some((hash, new_bytecode)) = transition.has_new_contract() {
554 self.contracts.insert(hash, new_bytecode.clone());
555 }
556 let revert = match self.state.entry(address) {
558 Entry::Occupied(mut entry) => {
559 let entry = entry.get_mut();
560 self.state_size -= entry.size_hint();
561 let revert = entry.update_and_create_revert(transition);
563 self.state_size += entry.size_hint();
565 revert
566 }
567 Entry::Vacant(entry) => {
568 let present_bundle = transition.present_bundle_account();
570 let revert = transition.create_revert();
571 if revert.is_some() {
572 self.state_size += present_bundle.size_hint();
573 entry.insert(present_bundle);
574 }
575 revert
576 }
577 };
578
579 if let Some(revert) = revert.filter(|_| include_reverts) {
581 self.reverts_size += revert.size_hint();
582 reverts.push((address, revert));
583 }
584 }
585
586 self.reverts.push(reverts);
587 }
588
589 pub fn to_plain_state(&self, is_value_known: OriginalValuesKnown) -> StateChangeset {
592 let state_len = self.state.len();
594 let mut accounts = Vec::with_capacity(state_len);
595 let mut storage = Vec::with_capacity(state_len);
596
597 for (address, account) in &self.state {
598 let was_destroyed = account.was_destroyed();
600 if is_value_known.is_not_known() || account.is_info_changed() {
601 let info = account.info.as_ref().map(AccountInfo::copy_without_code);
602 accounts.push((*address, info));
603 }
604
605 let mut account_storage_changed = Vec::with_capacity(account.storage.len());
610
611 for (key, slot) in account.storage.iter().map(|(k, v)| (*k, *v)) {
612 let destroyed_and_not_zero = was_destroyed && !slot.present_value.is_zero();
615
616 let not_destroyed_and_changed = !was_destroyed && slot.is_changed();
619
620 if is_value_known.is_not_known()
621 || destroyed_and_not_zero
622 || not_destroyed_and_changed
623 {
624 account_storage_changed.push((key, slot.present_value));
625 }
626 }
627
628 if !account_storage_changed.is_empty() || was_destroyed {
629 storage.push(PlainStorageChangeset {
631 address: *address,
632 wipe_storage: was_destroyed,
633 storage: account_storage_changed,
634 });
635 }
636 }
637
638 let contracts = self
639 .contracts
640 .iter()
641 .filter(|(b, _)| **b != KECCAK_EMPTY)
643 .map(|(b, code)| (*b, code.clone()))
644 .collect::<Vec<_>>();
645 StateChangeset {
646 accounts,
647 storage,
648 contracts,
649 }
650 }
651
652 #[deprecated = "Use `to_plain_state` instead"]
654 pub fn into_plain_state(self, is_value_known: OriginalValuesKnown) -> StateChangeset {
655 self.to_plain_state(is_value_known)
656 }
657
658 pub fn to_plain_state_and_reverts(
661 &self,
662 is_value_known: OriginalValuesKnown,
663 ) -> (StateChangeset, PlainStateReverts) {
664 (
665 self.to_plain_state(is_value_known),
666 self.reverts.to_plain_state_reverts(),
667 )
668 }
669
670 #[deprecated = "Use `to_plain_state_and_reverts` instead"]
673 pub fn into_plain_state_and_reverts(
674 self,
675 is_value_known: OriginalValuesKnown,
676 ) -> (StateChangeset, PlainStateReverts) {
677 self.to_plain_state_and_reverts(is_value_known)
678 }
679
680 pub fn extend_state(&mut self, other_state: HashMap<Address, BundleAccount>) {
684 for (address, other_account) in other_state {
685 match self.state.entry(address) {
686 Entry::Occupied(mut entry) => {
687 let this = entry.get_mut();
688 self.state_size -= this.size_hint();
689
690 if other_account.was_destroyed() {
693 this.storage = other_account.storage;
694 } else {
695 for (key, storage_slot) in other_account.storage {
697 this.storage
699 .entry(key)
700 .or_insert(storage_slot)
701 .present_value = storage_slot.present_value;
702 }
703 }
704 this.info = other_account.info;
705 this.status.transition(other_account.status);
706
707 self.state_size += this.size_hint();
709 }
710 Entry::Vacant(entry) => {
711 self.state_size += other_account.size_hint();
713 entry.insert(other_account);
714 }
715 }
716 }
717 }
718 pub fn extend(&mut self, mut other: Self) {
729 for (address, revert) in other.reverts.iter_mut().flatten() {
732 if revert.wipe_storage {
733 if let Some(this_account) = self.state.get_mut(address) {
736 for (key, value) in this_account.storage.drain() {
739 revert
740 .storage
741 .entry(key)
742 .or_insert(RevertToSlot::Some(value.present_value));
743 }
744
745 if this_account.was_destroyed() {
747 revert.wipe_storage = false;
748 }
749 }
750 }
751
752 self.reverts_size += revert.size_hint();
754 }
755 self.extend_state(other.state);
757 self.contracts.extend(other.contracts);
759 self.reverts.extend(other.reverts);
761 }
762
763 pub fn take_n_reverts(&mut self, reverts_to_take: usize) -> Reverts {
765 if reverts_to_take > self.reverts.len() {
767 return self.take_all_reverts();
768 }
769 let (detach, this) = self.reverts.split_at(reverts_to_take);
770 let detached_reverts = Reverts::new(detach.to_vec());
771 self.reverts_size = this
772 .iter()
773 .flatten()
774 .fold(0, |acc, (_, revert)| acc + revert.size_hint());
775 self.reverts = Reverts::new(this.to_vec());
776 detached_reverts
777 }
778
779 pub fn take_all_reverts(&mut self) -> Reverts {
781 self.reverts_size = 0;
782 mem::take(&mut self.reverts)
783 }
784
785 pub fn revert_latest(&mut self) -> bool {
791 if let Some(reverts) = self.reverts.pop() {
793 for (address, revert_account) in reverts.into_iter() {
794 self.reverts_size -= revert_account.size_hint();
795 match self.state.entry(address) {
796 Entry::Occupied(mut entry) => {
797 let account = entry.get_mut();
798 self.state_size -= account.size_hint();
799 if account.revert(revert_account) {
800 entry.remove();
801 } else {
802 self.state_size += account.size_hint();
803 }
804 }
805 Entry::Vacant(entry) => {
806 let mut account = BundleAccount::new(
809 None,
810 None,
811 HashMap::default(),
812 AccountStatus::LoadedNotExisting,
813 );
814 if !account.revert(revert_account) {
815 self.state_size += account.size_hint();
816 entry.insert(account);
817 }
818 }
819 }
820 }
821 return true;
822 }
823
824 false
825 }
826
827 pub fn revert(&mut self, mut num_transitions: usize) {
831 if num_transitions == 0 {
832 return;
833 }
834
835 while self.revert_latest() {
836 num_transitions -= 1;
837 if num_transitions == 0 {
838 break;
840 }
841 }
842 }
843
844 pub fn prepend_state(&mut self, mut other: BundleState) {
850 let this_bundle = mem::take(self);
852 other.extend_state(this_bundle.state);
854 other.contracts.extend(this_bundle.contracts);
856 mem::swap(self, &mut other)
858 }
859}
860
861#[cfg(test)]
862mod tests {
863 use super::*;
864 use crate::{StorageWithOriginalValues, TransitionAccount};
865
866 #[test]
867 fn transition_states() {
868 let address = Address::new([0x01; 20]);
870 let acc1 = AccountInfo {
871 balance: U256::from(10),
872 nonce: 1,
873 code_hash: KECCAK_EMPTY,
874 code: None,
875 };
876
877 let mut bundle_state = BundleState::default();
878
879 let transition = TransitionAccount {
882 info: Some(acc1),
883 status: AccountStatus::InMemoryChange,
884 previous_info: None,
885 previous_status: AccountStatus::LoadedNotExisting,
886 storage: StorageWithOriginalValues::default(),
887 storage_was_destroyed: false,
888 };
889
890 bundle_state.apply_transitions_and_create_reverts(
892 TransitionState::single(address, transition.clone()),
893 BundleRetention::Reverts,
894 );
895 }
896
897 const fn account1() -> Address {
898 Address::new([0x60; 20])
899 }
900
901 const fn account2() -> Address {
902 Address::new([0x61; 20])
903 }
904
905 fn slot1() -> U256 {
906 U256::from(5)
907 }
908
909 fn slot2() -> U256 {
910 U256::from(7)
911 }
912
913 fn test_bundle1() -> BundleState {
915 BundleState::new(
917 vec![
918 (
919 account1(),
920 None,
921 Some(AccountInfo {
922 nonce: 1,
923 balance: U256::from(10),
924 code_hash: KECCAK_EMPTY,
925 code: None,
926 }),
927 HashMap::from_iter([
928 (slot1(), (U256::from(0), U256::from(10))),
929 (slot2(), (U256::from(0), U256::from(15))),
930 ]),
931 ),
932 (
933 account2(),
934 None,
935 Some(AccountInfo {
936 nonce: 1,
937 balance: U256::from(10),
938 code_hash: KECCAK_EMPTY,
939 code: None,
940 }),
941 HashMap::default(),
942 ),
943 ],
944 vec![vec![
945 (
946 account1(),
947 Some(None),
948 vec![(slot1(), U256::from(0)), (slot2(), U256::from(0))],
949 ),
950 (account2(), Some(None), vec![]),
951 ]],
952 vec![],
953 )
954 }
955
956 fn test_bundle2() -> BundleState {
958 BundleState::new(
960 vec![(
961 account1(),
962 None,
963 Some(AccountInfo {
964 nonce: 3,
965 balance: U256::from(20),
966 code_hash: KECCAK_EMPTY,
967 code: None,
968 }),
969 HashMap::from_iter([(slot1(), (U256::from(0), U256::from(15)))]),
970 )],
971 vec![vec![(
972 account1(),
973 Some(Some(AccountInfo {
974 nonce: 1,
975 balance: U256::from(10),
976 code_hash: KECCAK_EMPTY,
977 code: None,
978 })),
979 vec![(slot1(), U256::from(10))],
980 )]],
981 vec![],
982 )
983 }
984
985 fn test_bundle3() -> BundleState {
987 BundleState::builder(0..=0)
988 .state_present_account_info(
989 account1(),
990 AccountInfo {
991 nonce: 1,
992 balance: U256::from(10),
993 code_hash: KECCAK_EMPTY,
994 code: None,
995 },
996 )
997 .state_storage(
998 account1(),
999 HashMap::from_iter([(slot1(), (U256::from(0), U256::from(10)))]),
1000 )
1001 .state_address(account2())
1002 .state_present_account_info(
1003 account2(),
1004 AccountInfo {
1005 nonce: 1,
1006 balance: U256::from(10),
1007 code_hash: KECCAK_EMPTY,
1008 code: None,
1009 },
1010 )
1011 .revert_address(0, account1())
1012 .revert_account_info(0, account1(), Some(None))
1013 .revert_storage(0, account1(), vec![(slot1(), U256::from(0))])
1014 .revert_account_info(0, account2(), Some(None))
1015 .build()
1016 }
1017
1018 fn test_bundle4() -> BundleState {
1020 BundleState::builder(0..=0)
1021 .state_present_account_info(
1022 account1(),
1023 AccountInfo {
1024 nonce: 3,
1025 balance: U256::from(20),
1026 code_hash: KECCAK_EMPTY,
1027 code: None,
1028 },
1029 )
1030 .state_storage(
1031 account1(),
1032 HashMap::from_iter([(slot1(), (U256::from(0), U256::from(15)))]),
1033 )
1034 .revert_address(0, account1())
1035 .revert_account_info(
1036 0,
1037 account1(),
1038 Some(Some(AccountInfo {
1039 nonce: 1,
1040 balance: U256::from(10),
1041 code_hash: KECCAK_EMPTY,
1042 code: None,
1043 })),
1044 )
1045 .revert_storage(0, account1(), vec![(slot1(), U256::from(10))])
1046 .build()
1047 }
1048
1049 fn sanity_path(bundle1: BundleState, bundle2: BundleState) {
1050 let mut extended = bundle1.clone();
1051 extended.extend(bundle2.clone());
1052
1053 let mut reverted = extended.clone();
1054 reverted.revert(0);
1056 assert_eq!(reverted, extended);
1057
1058 reverted.revert(1);
1060 assert_eq!(reverted, bundle1);
1061
1062 reverted.revert(1);
1064 assert_eq!(reverted, BundleState::default());
1065
1066 let mut reverted = extended.clone();
1067
1068 reverted.revert(10);
1070 assert_eq!(reverted, BundleState::default());
1071 }
1072
1073 #[test]
1074 fn extend_on_destroyed_values() {
1075 let base_bundle1 = test_bundle1();
1076 let base_bundle2 = test_bundle2();
1077
1078 let mut b1 = base_bundle1.clone();
1083 let mut b2 = base_bundle2.clone();
1084 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1085 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1086 b1.extend(b2);
1087 assert_eq!(
1088 b1.state.get_mut(&account1()).unwrap().status,
1089 AccountStatus::DestroyedChanged
1090 );
1091
1092 let mut b1 = base_bundle1.clone();
1097 let mut b2 = base_bundle2.clone();
1098 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1099 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1100 b2.reverts[0][0].1.wipe_storage = true;
1101 b1.extend(b2);
1102 assert_eq!(
1103 b1.state.get_mut(&account1()).unwrap().status,
1104 AccountStatus::Destroyed
1105 );
1106
1107 let mut revert1 = base_bundle2.reverts[0][0].clone();
1110 revert1.1.wipe_storage = true;
1111 revert1
1112 .1
1113 .storage
1114 .insert(slot2(), RevertToSlot::Some(U256::from(15)));
1115
1116 assert_eq!(
1117 b1.reverts.as_ref(),
1118 vec![base_bundle1.reverts[0].clone(), vec![revert1]],
1119 );
1120
1121 let mut b1 = base_bundle1.clone();
1127 let mut b2 = base_bundle2.clone();
1128 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::InMemoryChange;
1129 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1130 b1.extend(b2);
1131 assert_eq!(
1132 b1.state.get_mut(&account1()).unwrap().status,
1133 AccountStatus::InMemoryChange
1134 );
1135 }
1136
1137 #[test]
1138 fn test_sanity_path() {
1139 sanity_path(test_bundle1(), test_bundle2());
1140 sanity_path(test_bundle3(), test_bundle4());
1141 }
1142
1143 #[test]
1144 fn test_multi_reverts_with_delete() {
1145 let mut state = BundleBuilder::new(0..=3)
1146 .revert_address(0, account1())
1147 .revert_account_info(2, account1(), Some(Some(AccountInfo::default())))
1148 .revert_account_info(3, account1(), Some(None))
1149 .build();
1150
1151 state.revert_latest();
1152 assert_eq!(state.state.get(&account1()), None);
1154
1155 state.revert_latest();
1156 assert_eq!(
1158 state.state.get(&account1()),
1159 Some(&BundleAccount::new(
1160 None,
1161 Some(AccountInfo::default()),
1162 HashMap::default(),
1163 AccountStatus::Changed
1164 ))
1165 );
1166 }
1167
1168 #[test]
1169 fn test_revert_capacity() {
1170 let state = BundleState::builder(0..=3)
1171 .revert_address(0, account1())
1172 .revert_address(2, account2())
1173 .revert_account_info(0, account1(), Some(None))
1174 .revert_account_info(2, account2(), None)
1175 .revert_storage(0, account1(), vec![(slot1(), U256::from(10))])
1176 .build();
1177
1178 assert_eq!(state.reverts.len(), 4);
1179 assert_eq!(state.reverts[1], vec![]);
1180 assert_eq!(state.reverts[3], vec![]);
1181 assert_eq!(state.reverts[0].len(), 1);
1182 assert_eq!(state.reverts[2].len(), 1);
1183
1184 let (addr1, revert1) = &state.reverts[0][0];
1185 assert_eq!(addr1, &account1());
1186 assert_eq!(revert1.account, AccountInfoRevert::DeleteIt);
1187
1188 let (addr2, revert2) = &state.reverts[2][0];
1189 assert_eq!(addr2, &account2());
1190 assert_eq!(revert2.account, AccountInfoRevert::DoNothing);
1191 }
1192
1193 #[test]
1194 fn take_reverts() {
1195 let bundle1 = test_bundle1();
1196 let bundle2 = test_bundle2();
1197
1198 let mut extended = bundle1.clone();
1199 extended.extend(bundle2.clone());
1200
1201 assert_eq!(extended.reverts.len(), 2);
1203
1204 let mut extended2 = extended.clone();
1206 assert_eq!(extended2.take_n_reverts(100), extended.reverts);
1207
1208 let mut extended2 = extended.clone();
1210 assert_eq!(extended2.take_all_reverts(), extended.reverts);
1211
1212 let taken_reverts = extended.take_n_reverts(0);
1214 assert_eq!(taken_reverts, Reverts::default());
1215 assert_eq!(extended.reverts.len(), 2);
1216
1217 let taken_reverts = extended.take_n_reverts(1);
1219 assert_eq!(taken_reverts, bundle1.reverts);
1220
1221 let taken_reverts = extended.take_n_reverts(1);
1223 assert_eq!(taken_reverts, bundle2.reverts);
1224 }
1225
1226 #[test]
1227 fn prepend_state() {
1228 let address1 = account1();
1229 let address2 = account2();
1230
1231 let account1 = AccountInfo {
1232 nonce: 1,
1233 ..Default::default()
1234 };
1235 let account1_changed = AccountInfo {
1236 nonce: 1,
1237 ..Default::default()
1238 };
1239 let account2 = AccountInfo {
1240 nonce: 1,
1241 ..Default::default()
1242 };
1243
1244 let present_state = BundleState::builder(2..=2)
1245 .state_present_account_info(address1, account1_changed.clone())
1246 .build();
1247 assert_eq!(present_state.reverts.len(), 1);
1248 let previous_state = BundleState::builder(1..=1)
1249 .state_present_account_info(address1, account1)
1250 .state_present_account_info(address2, account2.clone())
1251 .build();
1252 assert_eq!(previous_state.reverts.len(), 1);
1253
1254 let mut test = present_state;
1255
1256 test.prepend_state(previous_state);
1257
1258 assert_eq!(test.state.len(), 2);
1259 assert_eq!(test.reverts.len(), 1);
1261 assert_eq!(
1263 test.state.get(&address1).unwrap().info,
1264 Some(account1_changed)
1265 );
1266 assert_eq!(test.state.get(&address2).unwrap().info, Some(account2));
1268 }
1269
1270 #[test]
1271 fn test_getters() {
1272 let mut builder = BundleBuilder::new(0..=3);
1273
1274 assert!(builder.get_states().is_empty());
1276 builder.get_states_mut().insert(account1());
1277 assert!(builder.get_states().contains(&account1()));
1278
1279 assert!(builder.get_state_original_mut().is_empty());
1281 builder
1282 .get_state_original_mut()
1283 .insert(account1(), AccountInfo::default());
1284 assert!(builder.get_state_original_mut().contains_key(&account1()));
1285
1286 assert!(builder.get_state_present_mut().is_empty());
1288 builder
1289 .get_state_present_mut()
1290 .insert(account1(), AccountInfo::default());
1291 assert!(builder.get_state_present_mut().contains_key(&account1()));
1292
1293 assert!(builder.get_state_storage_mut().is_empty());
1295 builder
1296 .get_state_storage_mut()
1297 .insert(account1(), HashMap::default());
1298 assert!(builder.get_state_storage_mut().contains_key(&account1()));
1299
1300 assert!(builder.get_reverts_mut().is_empty());
1302 builder.get_reverts_mut().insert((0, account1()));
1303 assert!(builder.get_reverts_mut().contains(&(0, account1())));
1304
1305 assert_eq!(builder.get_revert_range_mut().clone(), 0..=3);
1307
1308 assert!(builder.get_revert_account_mut().is_empty());
1310 builder
1311 .get_revert_account_mut()
1312 .insert((0, account1()), Some(None));
1313 assert!(builder
1314 .get_revert_account_mut()
1315 .contains_key(&(0, account1())));
1316
1317 assert!(builder.get_revert_storage_mut().is_empty());
1319 builder
1320 .get_revert_storage_mut()
1321 .insert((0, account1()), vec![(slot1(), U256::from(0))]);
1322 assert!(builder
1323 .get_revert_storage_mut()
1324 .contains_key(&(0, account1())));
1325
1326 assert!(builder.get_contracts_mut().is_empty());
1328 builder
1329 .get_contracts_mut()
1330 .insert(B256::default(), Bytecode::default());
1331 assert!(builder.get_contracts_mut().contains_key(&B256::default()));
1332 }
1333}