revm_database/states/
bundle_state.rs

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/// This builder is used to help to facilitate the initialization of `BundleState` struct
17#[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/// Option for [`BundleState`] when converting it to the plain state.
33#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
34pub enum OriginalValuesKnown {
35    /// Checks changed with original values that [BundleState] has.
36    ///
37    /// If we don't expect parent blocks to be committed or unwinded from database, this option
38    /// should be used.
39    Yes,
40    /// Doesn't check original values, see the implementation of [BundleState::to_plain_state] for
41    /// more info.
42    ///
43    /// If the Bundle can be split or extended, we would not be sure about original values, in that
44    /// case this option should be used.
45    No,
46}
47impl OriginalValuesKnown {
48    /// Original value is not known for sure.
49    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    /// Creates builder instance.
72    ///
73    /// `revert_range` indicates the size of BundleState `reverts` field.
74    pub fn new(revert_range: RangeInclusive<u64>) -> Self {
75        BundleBuilder {
76            revert_range,
77            ..Default::default()
78        }
79    }
80
81    /// Applies a transformation to the builder.
82    pub fn apply<F>(self, f: F) -> Self
83    where
84        F: FnOnce(Self) -> Self,
85    {
86        f(self)
87    }
88
89    /// Applies a mutable transformation to the builder.
90    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    /// Collects address info of BundleState state.
99    pub fn state_address(mut self, address: Address) -> Self {
100        self.set_state_address(address);
101        self
102    }
103
104    /// Collects account info of BundleState state.
105    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    /// Collects account info of BundleState state.
111    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    /// Collects storage info of BundleState state.
117    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    /// Collects address info of BundleState reverts.
123    ///
124    /// `block_number` must respect `revert_range`, or the input
125    /// will be ignored during the final build process.
126    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    /// Collects account info of BundleState reverts.
132    ///
133    /// `block_number` must respect `revert_range`, or the input
134    /// will be ignored during the final build process.
135    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    /// Collects storage info of BundleState reverts.
146    ///
147    /// `block_number` must respect `revert_range`, or the input
148    /// will be ignored during the final build process.
149    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    /// Collects contracts info.
160    pub fn contract(mut self, address: B256, bytecode: Bytecode) -> Self {
161        self.set_contract(address, bytecode);
162        self
163    }
164
165    /// Sets address info of BundleState state.
166    pub fn set_state_address(&mut self, address: Address) -> &mut Self {
167        self.states.insert(address);
168        self
169    }
170
171    /// Sets original account info of BundleState state.
172    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    /// Sets present account info of BundleState state.
183    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    /// Sets storage info of BundleState state.
194    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    /// Sets address info of BundleState reverts.
205    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    /// Sets account info of BundleState reverts.
211    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    /// Sets storage info of BundleState reverts.
223    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    /// Sets contracts info.
235    pub fn set_contract(&mut self, address: B256, bytecode: Bytecode) -> &mut Self {
236        self.contracts.insert(address, bytecode);
237        self
238    }
239
240    /// Creates `BundleState` instance based on collected information.
241    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    /// Getter for `states` field
319    pub fn get_states(&self) -> &HashSet<Address> {
320        &self.states
321    }
322
323    /// Mutable getter for `states` field
324    pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
325        &mut self.states
326    }
327
328    /// Mutable getter for `state_original` field
329    pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
330        &mut self.state_original
331    }
332
333    /// Mutable getter for `state_present` field
334    pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
335        &mut self.state_present
336    }
337
338    /// Mutable getter for `state_storage` field
339    pub fn get_state_storage_mut(&mut self) -> &mut HashMap<Address, HashMap<U256, (U256, U256)>> {
340        &mut self.state_storage
341    }
342
343    /// Mutable getter for `reverts` field
344    pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> {
345        &mut self.reverts
346    }
347
348    /// Mutable getter for `revert_range` field
349    pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive<u64> {
350        &mut self.revert_range
351    }
352
353    /// Mutable getter for `revert_account` field
354    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    /// Mutable getter for `revert_storage` field
361    pub fn get_revert_storage_mut(&mut self) -> &mut HashMap<(u64, Address), Vec<(U256, U256)>> {
362        &mut self.revert_storage
363    }
364
365    /// Mutable getter for `contracts` field
366    pub fn get_contracts_mut(&mut self) -> &mut HashMap<B256, Bytecode> {
367        &mut self.contracts
368    }
369}
370
371/// Bundle retention policy for applying substate to the bundle.
372#[derive(Debug)]
373pub enum BundleRetention {
374    /// Only plain state is updated.
375    PlainState,
376    /// Both, plain state and reverts, are retained
377    Reverts,
378}
379
380impl BundleRetention {
381    /// Returns `true` if reverts should be retained.
382    pub fn includes_reverts(&self) -> bool {
383        matches!(self, Self::Reverts)
384    }
385}
386
387/// Bundle state contain only values that got changed
388///
389/// For every account it contains both original and present state.
390///
391/// This is needed to decide if there were any changes to the account.
392///
393/// Reverts and created when TransitionState is applied to BundleState.
394///
395/// And can be used to revert BundleState to the state before transition.
396#[derive(Default, Clone, Debug, PartialEq, Eq)]
397#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
398pub struct BundleState {
399    /// Account state
400    pub state: HashMap<Address, BundleAccount>,
401    /// All created contracts in this block.
402    pub contracts: HashMap<B256, Bytecode>,
403    /// Changes to revert
404    ///
405    /// **Note**: Inside vector is *not* sorted by address.
406    ///
407    /// But it is unique by address.
408    pub reverts: Reverts,
409    /// The size of the plain state in the bundle state
410    pub state_size: usize,
411    /// The size of reverts in the bundle state
412    pub reverts_size: usize,
413}
414
415impl BundleState {
416    /// Returns builder instance for further manipulation.
417    pub fn builder(revert_range: RangeInclusive<u64>) -> BundleBuilder {
418        BundleBuilder::new(revert_range)
419    }
420
421    /// Creates it with new and old values of both Storage and AccountInfo.
422    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        // Create state from iterator.
443        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        // Create reverts from iterator.
462        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    /// Returns the approximate size of changes in the bundle state.
500    ///
501    /// The estimation is not precise, because the information about the number of
502    /// destroyed entries that need to be removed is not accessible to the bundle state.
503    pub fn size_hint(&self) -> usize {
504        self.state_size + self.reverts_size + self.contracts.len()
505    }
506
507    /// Returns reference to the state.
508    pub fn state(&self) -> &HashMap<Address, BundleAccount> {
509        &self.state
510    }
511
512    /// Checks whether bundle state is empty.
513    pub fn is_empty(&self) -> bool {
514        self.len() == 0
515    }
516
517    /// Return number of changed accounts.
518    pub fn len(&self) -> usize {
519        self.state.len()
520    }
521
522    /// Gets account from state.
523    pub fn account(&self, address: &Address) -> Option<&BundleAccount> {
524        self.state.get(address)
525    }
526
527    /// Gets bytecode from state.
528    pub fn bytecode(&self, hash: &B256) -> Option<Bytecode> {
529        self.contracts.get(hash).cloned()
530    }
531
532    /// Consumes [`TransitionState`] by applying the changes and creating the
533    /// reverts.
534    ///
535    /// If [BundleRetention::includes_reverts] is `true`, then the reverts will
536    /// be retained.
537    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        // Pessimistically pre-allocate assuming _all_ accounts changed.
544        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            // Add new contract if it was created/changed.
553            if let Some((hash, new_bytecode)) = transition.has_new_contract() {
554                self.contracts.insert(hash, new_bytecode.clone());
555            }
556            // Update state and create revert.
557            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                    // Update and create revert if it is present
562                    let revert = entry.update_and_create_revert(transition);
563                    // Update the state size
564                    self.state_size += entry.size_hint();
565                    revert
566                }
567                Entry::Vacant(entry) => {
568                    // Make revert from transition account
569                    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            // Append revert if present.
580            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    /// Generate a [`StateChangeset`] from the bundle state without consuming
590    /// it.
591    pub fn to_plain_state(&self, is_value_known: OriginalValuesKnown) -> StateChangeset {
592        // Pessimistically pre-allocate assuming _all_ accounts changed.
593        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            // Append account info if it is changed.
599            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            // Append storage changes
606
607            // Note: Assumption is that revert is going to remove whole plain storage from
608            // database so we can check if plain state was wiped or not.
609            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                // If storage was destroyed that means that storage was wiped.
613                // In that case we need to check if present storage value is different then ZERO.
614                let destroyed_and_not_zero = was_destroyed && !slot.present_value.is_zero();
615
616                // If account is not destroyed check if original values was changed,
617                // so we can update it.
618                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                // Append storage changes to account.
630                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            // Remove empty bytecodes
642            .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    /// Converts the bundle state into a [`StateChangeset`].
653    #[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    /// Generates a [`StateChangeset`] and [`PlainStateReverts`] from the bundle
659    /// state.
660    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    /// Consumes the bundle state and split it into a [`StateChangeset`] and a
671    /// [`PlainStateReverts`].
672    #[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    /// Extends the bundle with other state.
681    ///
682    /// Updates the `other` state only if `other` is not flagged as destroyed.
683    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 was destroyed. replace `this` storage with
691                    // the `other one.
692                    if other_account.was_destroyed() {
693                        this.storage = other_account.storage;
694                    } else {
695                        // Otherwise extend this storage with other
696                        for (key, storage_slot) in other_account.storage {
697                            // Update present value or insert storage slot.
698                            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                    // Update the state size
708                    self.state_size += this.size_hint();
709                }
710                Entry::Vacant(entry) => {
711                    // Just insert if empty
712                    self.state_size += other_account.size_hint();
713                    entry.insert(other_account);
714                }
715            }
716        }
717    }
718    /// Extends the state with state that is build on top of it.
719    ///
720    /// If storage was wiped in `other` state, copy `this` plain state
721    /// and put it inside `other` revert (if there is no duplicates of course).
722    ///
723    /// If `this` and `other` accounts were both destroyed invalidate second
724    /// wipe flag (from `other`).
725    ///
726    /// As wiping from database should be done only once
727    /// and we already transferred all potentially missing storages to the `other` revert.
728    pub fn extend(&mut self, mut other: Self) {
729        // Iterate over reverts and if its storage is wiped try to add previous bundle
730        // state as there is potential missing slots.
731        for (address, revert) in other.reverts.iter_mut().flatten() {
732            if revert.wipe_storage {
733                // If there is wipe storage in `other` revert
734                // we need to move storage from present state.
735                if let Some(this_account) = self.state.get_mut(address) {
736                    // As this account was destroyed inside `other` bundle.
737                    // We are fine to wipe/drain this storage and put it inside revert.
738                    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                    // Nullify `other` wipe as primary database wipe is done in `this`.
746                    if this_account.was_destroyed() {
747                        revert.wipe_storage = false;
748                    }
749                }
750            }
751
752            // Increment reverts size for each of the updated reverts.
753            self.reverts_size += revert.size_hint();
754        }
755        // Extension of state
756        self.extend_state(other.state);
757        // Contract can be just extended, when counter is introduced we will take into account that.
758        self.contracts.extend(other.contracts);
759        // Reverts can be just extended
760        self.reverts.extend(other.reverts);
761    }
762
763    /// Takes first N raw reverts from the [BundleState].
764    pub fn take_n_reverts(&mut self, reverts_to_take: usize) -> Reverts {
765        // Split is done as [0, num) and [num, len].
766        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    /// Returns and clears all reverts from [BundleState].
780    pub fn take_all_reverts(&mut self) -> Reverts {
781        self.reverts_size = 0;
782        mem::take(&mut self.reverts)
783    }
784
785    /// Reverts the state changes of the latest transition.
786    ///
787    /// **Note**: This is the same as `BundleState::revert(1)`
788    ///
789    /// Returns true if the state was reverted.
790    pub fn revert_latest(&mut self) -> bool {
791        // Revert the latest recorded state
792        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                        // Create empty account that we will revert on.
807                        // Only place where this account is not existing is if revert is DeleteIt.
808                        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    /// Reverts the state changes by N transitions back.
828    ///
829    /// See also [Self::revert_latest]
830    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 the loop.
839                break;
840            }
841        }
842    }
843
844    /// Prepends present the state with the given [BundleState].
845    ///
846    /// It adds changes from the given state but does not override any existing changes.
847    ///
848    /// Reverts are not updated.
849    pub fn prepend_state(&mut self, mut other: BundleState) {
850        // Take this bundle
851        let this_bundle = mem::take(self);
852        // Extend other bundle state with this
853        other.extend_state(this_bundle.state);
854        // Extend other contracts
855        other.contracts.extend(this_bundle.contracts);
856        // Swap bundles
857        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        // Dummy data
869        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        // Have transition from loaded to all other states
880
881        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        // Apply first transition
891        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    /// Tests bundle one.
914    fn test_bundle1() -> BundleState {
915        // Block changes
916        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    /// Tests bundle two.
957    fn test_bundle2() -> BundleState {
958        // Block changes
959        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    /// Tests bundle three.
986    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    /// Tests bundle four.
1019    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        // Revert zero does nothing.
1055        reverted.revert(0);
1056        assert_eq!(reverted, extended);
1057
1058        // Revert by one gives us bundle one.
1059        reverted.revert(1);
1060        assert_eq!(reverted, bundle1);
1061
1062        // Reverted by additional one gives us empty bundle.
1063        reverted.revert(1);
1064        assert_eq!(reverted, BundleState::default());
1065
1066        let mut reverted = extended.clone();
1067
1068        // Reverted by bigger number gives us empty bundle
1069        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        // test1
1079        // bundle1 has Destroyed
1080        // bundle2 has Changed
1081        // end should be DestroyedChanged.
1082        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        // test2
1093        // bundle1 has Changed
1094        // bundle2 has Destroyed
1095        // end should be Destroyed
1096        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        // test2 extension
1108        // revert of b2 should contains plain state of b1.
1109        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        // test3
1122        // bundle1 has InMemoryChange
1123        // bundle2 has Change
1124        // end should be InMemoryChange.
1125
1126        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        // State for account one was deleted
1153        assert_eq!(state.state.get(&account1()), None);
1154
1155        state.revert_latest();
1156        // State is set to
1157        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        // Check that we have two reverts
1202        assert_eq!(extended.reverts.len(), 2);
1203
1204        // Take all by big N
1205        let mut extended2 = extended.clone();
1206        assert_eq!(extended2.take_n_reverts(100), extended.reverts);
1207
1208        // Take all reverts
1209        let mut extended2 = extended.clone();
1210        assert_eq!(extended2.take_all_reverts(), extended.reverts);
1211
1212        // Take zero revert
1213        let taken_reverts = extended.take_n_reverts(0);
1214        assert_eq!(taken_reverts, Reverts::default());
1215        assert_eq!(extended.reverts.len(), 2);
1216
1217        // Take one revert
1218        let taken_reverts = extended.take_n_reverts(1);
1219        assert_eq!(taken_reverts, bundle1.reverts);
1220
1221        // Take last revert
1222        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        // Reverts num should stay the same.
1260        assert_eq!(test.reverts.len(), 1);
1261        // Account1 is not overwritten.
1262        assert_eq!(
1263            test.state.get(&address1).unwrap().info,
1264            Some(account1_changed)
1265        );
1266        // Account2 got inserted
1267        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        // Test get_states and get_states_mut
1275        assert!(builder.get_states().is_empty());
1276        builder.get_states_mut().insert(account1());
1277        assert!(builder.get_states().contains(&account1()));
1278
1279        // Test get_state_original_mut
1280        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        // Test get_state_present_mut
1287        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        // Test get_state_storage_mut
1294        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        // Test get_reverts_mut
1301        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        // Test get_revert_range_mut
1306        assert_eq!(builder.get_revert_range_mut().clone(), 0..=3);
1307
1308        // Test get_revert_account_mut
1309        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        // Test get_revert_storage_mut
1318        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        // Test get_contracts_mut
1327        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}