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 reverts_map.contains_key(&block_number) {
307 reverts_size += account_revert.size_hint();
308 reverts_map
309 .entry(block_number)
310 .or_insert_with(Vec::new)
311 .push((address, account_revert));
312 }
313 });
314
315 BundleState {
316 state,
317 contracts: self.contracts,
318 reverts: Reverts::new(reverts_map.into_values().collect()),
319 state_size,
320 reverts_size,
321 }
322 }
323
324 pub fn get_states(&self) -> &HashSet<Address> {
326 &self.states
327 }
328
329 pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
331 &mut self.states
332 }
333
334 pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
336 &mut self.state_original
337 }
338
339 pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
341 &mut self.state_present
342 }
343
344 pub fn get_state_storage_mut(
346 &mut self,
347 ) -> &mut HashMap<Address, HashMap<StorageKey, (StorageValue, StorageValue)>> {
348 &mut self.state_storage
349 }
350
351 pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> {
353 &mut self.reverts
354 }
355
356 pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive<u64> {
358 &mut self.revert_range
359 }
360
361 pub fn get_revert_account_mut(
363 &mut self,
364 ) -> &mut HashMap<(u64, Address), Option<Option<AccountInfo>>> {
365 &mut self.revert_account
366 }
367
368 pub fn get_revert_storage_mut(
370 &mut self,
371 ) -> &mut HashMap<(u64, Address), Vec<(StorageKey, StorageValue)>> {
372 &mut self.revert_storage
373 }
374
375 pub fn get_contracts_mut(&mut self) -> &mut HashMap<B256, Bytecode> {
377 &mut self.contracts
378 }
379}
380
381#[derive(Debug)]
383pub enum BundleRetention {
384 PlainState,
386 Reverts,
388}
389
390impl BundleRetention {
391 pub fn includes_reverts(&self) -> bool {
393 matches!(self, Self::Reverts)
394 }
395}
396
397#[derive(Default, Clone, Debug, PartialEq, Eq)]
407#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
408pub struct BundleState {
409 pub state: HashMap<Address, BundleAccount>,
411 pub contracts: HashMap<B256, Bytecode>,
413 pub reverts: Reverts,
419 pub state_size: usize,
421 pub reverts_size: usize,
423}
424
425impl BundleState {
426 pub fn builder(revert_range: RangeInclusive<u64>) -> BundleBuilder {
428 BundleBuilder::new(revert_range)
429 }
430
431 pub fn new(
433 state: impl IntoIterator<
434 Item = (
435 Address,
436 Option<AccountInfo>,
437 Option<AccountInfo>,
438 HashMap<StorageKey, (StorageValue, StorageValue)>,
439 ),
440 >,
441 reverts: impl IntoIterator<
442 Item = impl IntoIterator<
443 Item = (
444 Address,
445 Option<Option<AccountInfo>>,
446 impl IntoIterator<Item = (StorageKey, StorageValue)>,
447 ),
448 >,
449 >,
450 contracts: impl IntoIterator<Item = (B256, Bytecode)>,
451 ) -> Self {
452 let mut state_size = 0;
454 let state = state
455 .into_iter()
456 .map(|(address, original, present, storage)| {
457 let account = BundleAccount::new(
458 original,
459 present,
460 storage
461 .into_iter()
462 .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
463 .collect(),
464 AccountStatus::Changed,
465 );
466 state_size += account.size_hint();
467 (address, account)
468 })
469 .collect();
470
471 let mut reverts_size = 0;
473 let reverts = reverts
474 .into_iter()
475 .map(|block_reverts| {
476 block_reverts
477 .into_iter()
478 .map(|(address, account, storage)| {
479 let account = match account {
480 Some(Some(account)) => AccountInfoRevert::RevertTo(account),
481 Some(None) => AccountInfoRevert::DeleteIt,
482 None => AccountInfoRevert::DoNothing,
483 };
484 let revert = AccountRevert {
485 account,
486 storage: storage
487 .into_iter()
488 .map(|(k, v)| (k, RevertToSlot::Some(v)))
489 .collect(),
490 previous_status: AccountStatus::Changed,
491 wipe_storage: false,
492 };
493 reverts_size += revert.size_hint();
494 (address, revert)
495 })
496 .collect::<Vec<_>>()
497 })
498 .collect::<Vec<_>>();
499
500 Self {
501 state,
502 contracts: contracts.into_iter().collect(),
503 reverts: Reverts::new(reverts),
504 state_size,
505 reverts_size,
506 }
507 }
508
509 pub fn size_hint(&self) -> usize {
514 self.state_size + self.reverts_size + self.contracts.len()
515 }
516
517 pub fn state(&self) -> &HashMap<Address, BundleAccount> {
519 &self.state
520 }
521
522 pub fn is_empty(&self) -> bool {
524 self.len() == 0
525 }
526
527 pub fn len(&self) -> usize {
529 self.state.len()
530 }
531
532 pub fn account(&self, address: &Address) -> Option<&BundleAccount> {
534 self.state.get(address)
535 }
536
537 pub fn bytecode(&self, hash: &B256) -> Option<Bytecode> {
539 self.contracts.get(hash).cloned()
540 }
541
542 pub fn apply_transitions_and_create_reverts(
548 &mut self,
549 transitions: TransitionState,
550 retention: BundleRetention,
551 ) {
552 let include_reverts = retention.includes_reverts();
553 let reverts_capacity = if include_reverts {
555 transitions.transitions.len()
556 } else {
557 0
558 };
559 let mut reverts = Vec::with_capacity(reverts_capacity);
560
561 for (address, transition) in transitions.transitions.into_iter() {
562 if let Some((hash, new_bytecode)) = transition.has_new_contract() {
564 self.contracts.insert(hash, new_bytecode.clone());
565 }
566 let revert = match self.state.entry(address) {
568 Entry::Occupied(mut entry) => {
569 let entry = entry.get_mut();
570 self.state_size -= entry.size_hint();
571 let revert = entry.update_and_create_revert(transition);
573 self.state_size += entry.size_hint();
575 revert
576 }
577 Entry::Vacant(entry) => {
578 let present_bundle = transition.present_bundle_account();
580 let revert = transition.create_revert();
581 if revert.is_some() {
582 self.state_size += present_bundle.size_hint();
583 entry.insert(present_bundle);
584 }
585 revert
586 }
587 };
588
589 if let Some(revert) = revert.filter(|_| include_reverts) {
591 self.reverts_size += revert.size_hint();
592 reverts.push((address, revert));
593 }
594 }
595
596 self.reverts.push(reverts);
597 }
598
599 pub fn to_plain_state(&self, is_value_known: OriginalValuesKnown) -> StateChangeset {
602 let state_len = self.state.len();
604 let mut accounts = Vec::with_capacity(state_len);
605 let mut storage = Vec::with_capacity(state_len);
606
607 for (address, account) in &self.state {
608 let was_destroyed = account.was_destroyed();
610 if is_value_known.is_not_known() || account.is_info_changed() {
611 let info = account.info.as_ref().map(AccountInfo::copy_without_code);
612 accounts.push((*address, info));
613 }
614
615 let mut account_storage_changed = Vec::with_capacity(account.storage.len());
620
621 for (&key, &slot) in account.storage.iter() {
622 let destroyed_and_not_zero = was_destroyed && !slot.present_value.is_zero();
625
626 let not_destroyed_and_changed = !was_destroyed && slot.is_changed();
629
630 if is_value_known.is_not_known()
631 || destroyed_and_not_zero
632 || not_destroyed_and_changed
633 {
634 account_storage_changed.push((key, slot.present_value));
635 }
636 }
637
638 if !account_storage_changed.is_empty() || was_destroyed {
639 storage.push(PlainStorageChangeset {
641 address: *address,
642 wipe_storage: was_destroyed,
643 storage: account_storage_changed,
644 });
645 }
646 }
647
648 let contracts = self
649 .contracts
650 .iter()
651 .filter(|(b, _)| **b != KECCAK_EMPTY)
653 .map(|(b, code)| (*b, code.clone()))
654 .collect::<Vec<_>>();
655 StateChangeset {
656 accounts,
657 storage,
658 contracts,
659 }
660 }
661
662 pub fn to_plain_state_and_reverts(
665 &self,
666 is_value_known: OriginalValuesKnown,
667 ) -> (StateChangeset, PlainStateReverts) {
668 (
669 self.to_plain_state(is_value_known),
670 self.reverts.to_plain_state_reverts(),
671 )
672 }
673
674 #[deprecated = "Use `to_plain_state_and_reverts` instead"]
677 pub fn into_plain_state_and_reverts(
678 self,
679 is_value_known: OriginalValuesKnown,
680 ) -> (StateChangeset, PlainStateReverts) {
681 self.to_plain_state_and_reverts(is_value_known)
682 }
683
684 pub fn extend_state(&mut self, other_state: HashMap<Address, BundleAccount>) {
688 for (address, other_account) in other_state {
689 match self.state.entry(address) {
690 Entry::Occupied(mut entry) => {
691 let this = entry.get_mut();
692 self.state_size -= this.size_hint();
693
694 if other_account.was_destroyed() {
697 this.storage = other_account.storage;
698 } else {
699 for (key, storage_slot) in other_account.storage {
701 this.storage
703 .entry(key)
704 .or_insert(storage_slot)
705 .present_value = storage_slot.present_value;
706 }
707 }
708 this.info = other_account.info;
709 this.status.transition(other_account.status);
710
711 self.state_size += this.size_hint();
713 }
714 Entry::Vacant(entry) => {
715 self.state_size += other_account.size_hint();
717 entry.insert(other_account);
718 }
719 }
720 }
721 }
722 pub fn extend(&mut self, mut other: Self) {
733 for (address, revert) in other.reverts.iter_mut().flatten() {
736 if revert.wipe_storage {
737 if let Some(this_account) = self.state.get_mut(address) {
740 for (key, value) in this_account.storage.drain() {
743 revert
744 .storage
745 .entry(key)
746 .or_insert(RevertToSlot::Some(value.present_value));
747 }
748
749 if this_account.was_destroyed() {
751 revert.wipe_storage = false;
752 }
753 }
754 }
755
756 self.reverts_size += revert.size_hint();
758 }
759 self.extend_state(other.state);
761 self.contracts.extend(other.contracts);
763 self.reverts.extend(other.reverts);
765 }
766
767 pub fn take_n_reverts(&mut self, reverts_to_take: usize) -> Reverts {
769 if reverts_to_take > self.reverts.len() {
771 return self.take_all_reverts();
772 }
773 let (detach, this) = self.reverts.split_at(reverts_to_take);
774 let detached_reverts = Reverts::new(detach.to_vec());
775 self.reverts_size = this
776 .iter()
777 .flatten()
778 .fold(0, |acc, (_, revert)| acc + revert.size_hint());
779 self.reverts = Reverts::new(this.to_vec());
780 detached_reverts
781 }
782
783 pub fn take_all_reverts(&mut self) -> Reverts {
785 self.reverts_size = 0;
786 mem::take(&mut self.reverts)
787 }
788
789 pub fn revert_latest(&mut self) -> bool {
795 if let Some(reverts) = self.reverts.pop() {
797 for (address, revert_account) in reverts.into_iter() {
798 self.reverts_size -= revert_account.size_hint();
799 match self.state.entry(address) {
800 Entry::Occupied(mut entry) => {
801 let account = entry.get_mut();
802 self.state_size -= account.size_hint();
803 if account.revert(revert_account) {
804 entry.remove();
805 } else {
806 self.state_size += account.size_hint();
807 }
808 }
809 Entry::Vacant(entry) => {
810 let mut account = BundleAccount::new(
813 None,
814 None,
815 HashMap::default(),
816 AccountStatus::LoadedNotExisting,
817 );
818 if !account.revert(revert_account) {
819 self.state_size += account.size_hint();
820 entry.insert(account);
821 }
822 }
823 }
824 }
825 return true;
826 }
827
828 false
829 }
830
831 pub fn revert(&mut self, mut num_transitions: usize) {
835 if num_transitions == 0 {
836 return;
837 }
838
839 while self.revert_latest() {
840 num_transitions -= 1;
841 if num_transitions == 0 {
842 break;
844 }
845 }
846 }
847
848 pub fn prepend_state(&mut self, mut other: BundleState) {
854 let this_bundle = mem::take(self);
856 other.extend_state(this_bundle.state);
858 other.contracts.extend(this_bundle.contracts);
860 mem::swap(self, &mut other)
862 }
863}
864
865#[cfg(test)]
866mod tests {
867 use super::*;
868 use crate::{StorageWithOriginalValues, TransitionAccount};
869 use primitives::U256;
870
871 #[test]
872 fn transition_states() {
873 let address = Address::new([0x01; 20]);
875 let acc1 = AccountInfo {
876 balance: U256::from(10),
877 nonce: 1,
878 code_hash: KECCAK_EMPTY,
879 code: None,
880 };
881
882 let mut bundle_state = BundleState::default();
883
884 let transition = TransitionAccount {
887 info: Some(acc1),
888 status: AccountStatus::InMemoryChange,
889 previous_info: None,
890 previous_status: AccountStatus::LoadedNotExisting,
891 storage: StorageWithOriginalValues::default(),
892 storage_was_destroyed: false,
893 };
894
895 bundle_state.apply_transitions_and_create_reverts(
897 TransitionState::single(address, transition.clone()),
898 BundleRetention::Reverts,
899 );
900 }
901
902 const fn account1() -> Address {
903 Address::new([0x60; 20])
904 }
905
906 const fn account2() -> Address {
907 Address::new([0x61; 20])
908 }
909
910 fn slot1() -> StorageKey {
911 StorageKey::from(5)
912 }
913
914 fn slot2() -> StorageKey {
915 StorageKey::from(7)
916 }
917
918 fn test_bundle1() -> BundleState {
920 BundleState::new(
922 vec![
923 (
924 account1(),
925 None,
926 Some(AccountInfo {
927 nonce: 1,
928 balance: U256::from(10),
929 code_hash: KECCAK_EMPTY,
930 code: None,
931 }),
932 HashMap::from_iter([
933 (slot1(), (StorageValue::from(0), StorageValue::from(10))),
934 (slot2(), (StorageValue::from(0), StorageValue::from(15))),
935 ]),
936 ),
937 (
938 account2(),
939 None,
940 Some(AccountInfo {
941 nonce: 1,
942 balance: U256::from(10),
943 code_hash: KECCAK_EMPTY,
944 code: None,
945 }),
946 HashMap::default(),
947 ),
948 ],
949 vec![vec![
950 (
951 account1(),
952 Some(None),
953 vec![
954 (slot1(), StorageValue::from(0)),
955 (slot2(), StorageValue::from(0)),
956 ],
957 ),
958 (account2(), Some(None), vec![]),
959 ]],
960 vec![],
961 )
962 }
963
964 fn test_bundle2() -> BundleState {
966 BundleState::new(
968 vec![(
969 account1(),
970 None,
971 Some(AccountInfo {
972 nonce: 3,
973 balance: U256::from(20),
974 code_hash: KECCAK_EMPTY,
975 code: None,
976 }),
977 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
978 )],
979 vec![vec![(
980 account1(),
981 Some(Some(AccountInfo {
982 nonce: 1,
983 balance: U256::from(10),
984 code_hash: KECCAK_EMPTY,
985 code: None,
986 })),
987 vec![(slot1(), StorageValue::from(10))],
988 )]],
989 vec![],
990 )
991 }
992
993 fn test_bundle3() -> BundleState {
995 BundleState::builder(0..=0)
996 .state_present_account_info(
997 account1(),
998 AccountInfo {
999 nonce: 1,
1000 balance: U256::from(10),
1001 code_hash: KECCAK_EMPTY,
1002 code: None,
1003 },
1004 )
1005 .state_storage(
1006 account1(),
1007 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(10)))]),
1008 )
1009 .state_address(account2())
1010 .state_present_account_info(
1011 account2(),
1012 AccountInfo {
1013 nonce: 1,
1014 balance: U256::from(10),
1015 code_hash: KECCAK_EMPTY,
1016 code: None,
1017 },
1018 )
1019 .revert_address(0, account1())
1020 .revert_account_info(0, account1(), Some(None))
1021 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(0))])
1022 .revert_account_info(0, account2(), Some(None))
1023 .build()
1024 }
1025
1026 fn test_bundle4() -> BundleState {
1028 BundleState::builder(0..=0)
1029 .state_present_account_info(
1030 account1(),
1031 AccountInfo {
1032 nonce: 3,
1033 balance: U256::from(20),
1034 code_hash: KECCAK_EMPTY,
1035 code: None,
1036 },
1037 )
1038 .state_storage(
1039 account1(),
1040 HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
1041 )
1042 .revert_address(0, account1())
1043 .revert_account_info(
1044 0,
1045 account1(),
1046 Some(Some(AccountInfo {
1047 nonce: 1,
1048 balance: U256::from(10),
1049 code_hash: KECCAK_EMPTY,
1050 code: None,
1051 })),
1052 )
1053 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1054 .build()
1055 }
1056
1057 fn sanity_path(bundle1: BundleState, bundle2: BundleState) {
1058 let mut extended = bundle1.clone();
1059 extended.extend(bundle2.clone());
1060
1061 let mut reverted = extended.clone();
1062 reverted.revert(0);
1064 assert_eq!(reverted, extended);
1065
1066 reverted.revert(1);
1068 assert_eq!(reverted, bundle1);
1069
1070 reverted.revert(1);
1072 assert_eq!(reverted, BundleState::default());
1073
1074 let mut reverted = extended.clone();
1075
1076 reverted.revert(10);
1078 assert_eq!(reverted, BundleState::default());
1079 }
1080
1081 #[test]
1082 fn extend_on_destroyed_values() {
1083 let base_bundle1 = test_bundle1();
1084 let base_bundle2 = test_bundle2();
1085
1086 let mut b1 = base_bundle1.clone();
1091 let mut b2 = base_bundle2.clone();
1092 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1093 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1094 b1.extend(b2);
1095 assert_eq!(
1096 b1.state.get_mut(&account1()).unwrap().status,
1097 AccountStatus::DestroyedChanged
1098 );
1099
1100 let mut b1 = base_bundle1.clone();
1105 let mut b2 = base_bundle2.clone();
1106 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1107 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1108 b2.reverts[0][0].1.wipe_storage = true;
1109 b1.extend(b2);
1110 assert_eq!(
1111 b1.state.get_mut(&account1()).unwrap().status,
1112 AccountStatus::Destroyed
1113 );
1114
1115 let mut revert1 = base_bundle2.reverts[0][0].clone();
1118 revert1.1.wipe_storage = true;
1119 revert1
1120 .1
1121 .storage
1122 .insert(slot2(), RevertToSlot::Some(StorageValue::from(15)));
1123
1124 assert_eq!(
1125 b1.reverts.as_ref(),
1126 vec![base_bundle1.reverts[0].clone(), vec![revert1]],
1127 );
1128
1129 let mut b1 = base_bundle1.clone();
1135 let mut b2 = base_bundle2.clone();
1136 b1.state.get_mut(&account1()).unwrap().status = AccountStatus::InMemoryChange;
1137 b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1138 b1.extend(b2);
1139 assert_eq!(
1140 b1.state.get_mut(&account1()).unwrap().status,
1141 AccountStatus::InMemoryChange
1142 );
1143 }
1144
1145 #[test]
1146 fn test_sanity_path() {
1147 sanity_path(test_bundle1(), test_bundle2());
1148 sanity_path(test_bundle3(), test_bundle4());
1149 }
1150
1151 #[test]
1152 fn test_multi_reverts_with_delete() {
1153 let mut state = BundleBuilder::new(0..=3)
1154 .revert_address(0, account1())
1155 .revert_account_info(2, account1(), Some(Some(AccountInfo::default())))
1156 .revert_account_info(3, account1(), Some(None))
1157 .build();
1158
1159 state.revert_latest();
1160 assert_eq!(state.state.get(&account1()), None);
1162
1163 state.revert_latest();
1164 assert_eq!(
1166 state.state.get(&account1()),
1167 Some(&BundleAccount::new(
1168 None,
1169 Some(AccountInfo::default()),
1170 HashMap::default(),
1171 AccountStatus::Changed
1172 ))
1173 );
1174 }
1175
1176 #[test]
1177 fn test_revert_capacity() {
1178 let state = BundleState::builder(0..=3)
1179 .revert_address(0, account1())
1180 .revert_address(2, account2())
1181 .revert_account_info(0, account1(), Some(None))
1182 .revert_account_info(2, account2(), None)
1183 .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1184 .build();
1185
1186 assert_eq!(state.reverts.len(), 4);
1187 assert_eq!(state.reverts[1], vec![]);
1188 assert_eq!(state.reverts[3], vec![]);
1189 assert_eq!(state.reverts[0].len(), 1);
1190 assert_eq!(state.reverts[2].len(), 1);
1191
1192 let (addr1, revert1) = &state.reverts[0][0];
1193 assert_eq!(addr1, &account1());
1194 assert_eq!(revert1.account, AccountInfoRevert::DeleteIt);
1195
1196 let (addr2, revert2) = &state.reverts[2][0];
1197 assert_eq!(addr2, &account2());
1198 assert_eq!(revert2.account, AccountInfoRevert::DoNothing);
1199 }
1200
1201 #[test]
1202 fn take_reverts() {
1203 let bundle1 = test_bundle1();
1204 let bundle2 = test_bundle2();
1205
1206 let mut extended = bundle1.clone();
1207 extended.extend(bundle2.clone());
1208
1209 assert_eq!(extended.reverts.len(), 2);
1211
1212 let mut extended2 = extended.clone();
1214 assert_eq!(extended2.take_n_reverts(100), extended.reverts);
1215
1216 let mut extended2 = extended.clone();
1218 assert_eq!(extended2.take_all_reverts(), extended.reverts);
1219
1220 let taken_reverts = extended.take_n_reverts(0);
1222 assert_eq!(taken_reverts, Reverts::default());
1223 assert_eq!(extended.reverts.len(), 2);
1224
1225 let taken_reverts = extended.take_n_reverts(1);
1227 assert_eq!(taken_reverts, bundle1.reverts);
1228
1229 let taken_reverts = extended.take_n_reverts(1);
1231 assert_eq!(taken_reverts, bundle2.reverts);
1232 }
1233
1234 #[test]
1235 fn prepend_state() {
1236 let address1 = account1();
1237 let address2 = account2();
1238
1239 let account1 = AccountInfo {
1240 nonce: 1,
1241 ..Default::default()
1242 };
1243 let account1_changed = AccountInfo {
1244 nonce: 1,
1245 ..Default::default()
1246 };
1247 let account2 = AccountInfo {
1248 nonce: 1,
1249 ..Default::default()
1250 };
1251
1252 let present_state = BundleState::builder(2..=2)
1253 .state_present_account_info(address1, account1_changed.clone())
1254 .build();
1255 assert_eq!(present_state.reverts.len(), 1);
1256 let previous_state = BundleState::builder(1..=1)
1257 .state_present_account_info(address1, account1)
1258 .state_present_account_info(address2, account2.clone())
1259 .build();
1260 assert_eq!(previous_state.reverts.len(), 1);
1261
1262 let mut test = present_state;
1263
1264 test.prepend_state(previous_state);
1265
1266 assert_eq!(test.state.len(), 2);
1267 assert_eq!(test.reverts.len(), 1);
1269 assert_eq!(
1271 test.state.get(&address1).unwrap().info,
1272 Some(account1_changed)
1273 );
1274 assert_eq!(test.state.get(&address2).unwrap().info, Some(account2));
1276 }
1277
1278 #[test]
1279 fn test_getters() {
1280 let mut builder = BundleBuilder::new(0..=3);
1281
1282 assert!(builder.get_states().is_empty());
1284 builder.get_states_mut().insert(account1());
1285 assert!(builder.get_states().contains(&account1()));
1286
1287 assert!(builder.get_state_original_mut().is_empty());
1289 builder
1290 .get_state_original_mut()
1291 .insert(account1(), AccountInfo::default());
1292 assert!(builder.get_state_original_mut().contains_key(&account1()));
1293
1294 assert!(builder.get_state_present_mut().is_empty());
1296 builder
1297 .get_state_present_mut()
1298 .insert(account1(), AccountInfo::default());
1299 assert!(builder.get_state_present_mut().contains_key(&account1()));
1300
1301 assert!(builder.get_state_storage_mut().is_empty());
1303 builder
1304 .get_state_storage_mut()
1305 .insert(account1(), HashMap::default());
1306 assert!(builder.get_state_storage_mut().contains_key(&account1()));
1307
1308 assert!(builder.get_reverts_mut().is_empty());
1310 builder.get_reverts_mut().insert((0, account1()));
1311 assert!(builder.get_reverts_mut().contains(&(0, account1())));
1312
1313 assert_eq!(builder.get_revert_range_mut().clone(), 0..=3);
1315
1316 assert!(builder.get_revert_account_mut().is_empty());
1318 builder
1319 .get_revert_account_mut()
1320 .insert((0, account1()), Some(None));
1321 assert!(builder
1322 .get_revert_account_mut()
1323 .contains_key(&(0, account1())));
1324
1325 assert!(builder.get_revert_storage_mut().is_empty());
1327 builder
1328 .get_revert_storage_mut()
1329 .insert((0, account1()), vec![(slot1(), StorageValue::from(0))]);
1330 assert!(builder
1331 .get_revert_storage_mut()
1332 .contains_key(&(0, account1())));
1333
1334 assert!(builder.get_contracts_mut().is_empty());
1336 builder
1337 .get_contracts_mut()
1338 .insert(B256::default(), Bytecode::default());
1339 assert!(builder.get_contracts_mut().contains_key(&B256::default()));
1340 }
1341}