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::{
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/// This builder is used to help to facilitate the initialization of `BundleState` struct
19#[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/// Option for [`BundleState`] when converting it to the plain state.
35#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
36pub enum OriginalValuesKnown {
37    /// Checks changed with original values that [BundleState] has.
38    ///
39    /// If we don't expect parent blocks to be committed or unwinded from database, this option
40    /// should be used.
41    Yes,
42    /// Doesn't check original values, see the implementation of [BundleState::to_plain_state] for
43    /// more info.
44    ///
45    /// If the Bundle can be split or extended, we would not be sure about original values, in that
46    /// case this option should be used.
47    No,
48}
49impl OriginalValuesKnown {
50    /// Original value is not known for sure.
51    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    /// Creates builder instance.
74    ///
75    /// `revert_range` indicates the size of BundleState `reverts` field.
76    pub fn new(revert_range: RangeInclusive<u64>) -> Self {
77        BundleBuilder {
78            revert_range,
79            ..Default::default()
80        }
81    }
82
83    /// Applies a transformation to the builder.
84    pub fn apply<F>(self, f: F) -> Self
85    where
86        F: FnOnce(Self) -> Self,
87    {
88        f(self)
89    }
90
91    /// Applies a mutable transformation to the builder.
92    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    /// Collects address info of BundleState state.
101    pub fn state_address(mut self, address: Address) -> Self {
102        self.set_state_address(address);
103        self
104    }
105
106    /// Collects account info of BundleState state.
107    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    /// Collects account info of BundleState state.
113    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    /// Collects storage info of BundleState state.
119    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    /// Collects address info of BundleState reverts.
129    ///
130    /// `block_number` must respect `revert_range`, or the input
131    /// will be ignored during the final build process.
132    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    /// Collects account info of BundleState reverts.
138    ///
139    /// `block_number` must respect `revert_range`, or the input
140    /// will be ignored during the final build process.
141    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    /// Collects storage info of BundleState reverts.
152    ///
153    /// `block_number` must respect `revert_range`, or the input
154    /// will be ignored during the final build process.
155    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    /// Collects contracts info.
166    pub fn contract(mut self, address: B256, bytecode: Bytecode) -> Self {
167        self.set_contract(address, bytecode);
168        self
169    }
170
171    /// Sets address info of BundleState state.
172    pub fn set_state_address(&mut self, address: Address) -> &mut Self {
173        self.states.insert(address);
174        self
175    }
176
177    /// Sets original account info of BundleState state.
178    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    /// Sets present account info of BundleState state.
189    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    /// Sets storage info of BundleState state.
200    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    /// Sets address info of BundleState reverts.
211    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    /// Sets account info of BundleState reverts.
217    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    /// Sets storage info of BundleState reverts.
229    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    /// Sets contracts info.
241    pub fn set_contract(&mut self, address: B256, bytecode: Bytecode) -> &mut Self {
242        self.contracts.insert(address, bytecode);
243        self
244    }
245
246    /// Creates `BundleState` instance based on collected information.
247    pub fn build(mut self) -> BundleState {
248        let mut state_size = 0;
249        let state = self
250            .states
251            .into_iter()
252            .map(|address| {
253                let storage = self
254                    .state_storage
255                    .remove(&address)
256                    .map(|s| {
257                        s.into_iter()
258                            .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
259                            .collect()
260                    })
261                    .unwrap_or_default();
262                let bundle_account = BundleAccount::new(
263                    self.state_original.remove(&address),
264                    self.state_present.remove(&address),
265                    storage,
266                    AccountStatus::Changed,
267                );
268                state_size += bundle_account.size_hint();
269                (address, bundle_account)
270            })
271            .collect();
272
273        let mut reverts_size = 0;
274        let mut reverts_map = BTreeMap::new();
275        for block_number in self.revert_range {
276            reverts_map.insert(block_number, Vec::new());
277        }
278        self.reverts
279            .into_iter()
280            .for_each(|(block_number, address)| {
281                let account = match self
282                    .revert_account
283                    .remove(&(block_number, address))
284                    .unwrap_or_default()
285                {
286                    Some(Some(account)) => AccountInfoRevert::RevertTo(account),
287                    Some(None) => AccountInfoRevert::DeleteIt,
288                    None => AccountInfoRevert::DoNothing,
289                };
290                let storage = self
291                    .revert_storage
292                    .remove(&(block_number, address))
293                    .map(|s| {
294                        s.into_iter()
295                            .map(|(k, v)| (k, RevertToSlot::Some(v)))
296                            .collect()
297                    })
298                    .unwrap_or_default();
299                let account_revert = AccountRevert {
300                    account,
301                    storage,
302                    previous_status: AccountStatus::Changed,
303                    wipe_storage: false,
304                };
305
306                if let Some(vec) = reverts_map.get_mut(&block_number) {
307                    reverts_size += account_revert.size_hint();
308                    vec.push((address, account_revert));
309                }
310            });
311
312        BundleState {
313            state,
314            contracts: self.contracts,
315            reverts: Reverts::new(reverts_map.into_values().collect()),
316            state_size,
317            reverts_size,
318        }
319    }
320
321    /// Getter for `states` field
322    pub fn get_states(&self) -> &HashSet<Address> {
323        &self.states
324    }
325
326    /// Mutable getter for `states` field
327    pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
328        &mut self.states
329    }
330
331    /// Mutable getter for `state_original` field
332    pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
333        &mut self.state_original
334    }
335
336    /// Mutable getter for `state_present` field
337    pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
338        &mut self.state_present
339    }
340
341    /// Mutable getter for `state_storage` field
342    pub fn get_state_storage_mut(
343        &mut self,
344    ) -> &mut HashMap<Address, HashMap<StorageKey, (StorageValue, StorageValue)>> {
345        &mut self.state_storage
346    }
347
348    /// Mutable getter for `reverts` field
349    pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> {
350        &mut self.reverts
351    }
352
353    /// Mutable getter for `revert_range` field
354    pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive<u64> {
355        &mut self.revert_range
356    }
357
358    /// Mutable getter for `revert_account` field
359    pub fn get_revert_account_mut(
360        &mut self,
361    ) -> &mut HashMap<(u64, Address), Option<Option<AccountInfo>>> {
362        &mut self.revert_account
363    }
364
365    /// Mutable getter for `revert_storage` field
366    pub fn get_revert_storage_mut(
367        &mut self,
368    ) -> &mut HashMap<(u64, Address), Vec<(StorageKey, StorageValue)>> {
369        &mut self.revert_storage
370    }
371
372    /// Mutable getter for `contracts` field
373    pub fn get_contracts_mut(&mut self) -> &mut HashMap<B256, Bytecode> {
374        &mut self.contracts
375    }
376}
377
378/// Bundle retention policy for applying substate to the bundle.
379#[derive(Debug)]
380pub enum BundleRetention {
381    /// Only plain state is updated.
382    PlainState,
383    /// Both, plain state and reverts, are retained
384    Reverts,
385}
386
387impl BundleRetention {
388    /// Returns `true` if reverts should be retained.
389    pub fn includes_reverts(&self) -> bool {
390        matches!(self, Self::Reverts)
391    }
392}
393
394/// Bundle state contain only values that got changed
395///
396/// For every account it contains both original and present state.
397///
398/// This is needed to decide if there were any changes to the account.
399///
400/// Reverts and created when TransitionState is applied to BundleState.
401///
402/// And can be used to revert BundleState to the state before transition.
403#[derive(Default, Clone, Debug, PartialEq, Eq)]
404#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
405pub struct BundleState {
406    /// Account state
407    pub state: HashMap<Address, BundleAccount>,
408    /// All created contracts in this block.
409    pub contracts: HashMap<B256, Bytecode>,
410    /// Changes to revert
411    ///
412    /// **Note**: Inside vector is *not* sorted by address.
413    ///
414    /// But it is unique by address.
415    pub reverts: Reverts,
416    /// The size of the plain state in the bundle state
417    pub state_size: usize,
418    /// The size of reverts in the bundle state
419    pub reverts_size: usize,
420}
421
422impl BundleState {
423    /// Returns builder instance for further manipulation.
424    pub fn builder(revert_range: RangeInclusive<u64>) -> BundleBuilder {
425        BundleBuilder::new(revert_range)
426    }
427
428    /// Creates it with new and old values of both Storage and AccountInfo.
429    pub fn new(
430        state: impl IntoIterator<
431            Item = (
432                Address,
433                Option<AccountInfo>,
434                Option<AccountInfo>,
435                HashMap<StorageKey, (StorageValue, StorageValue)>,
436            ),
437        >,
438        reverts: impl IntoIterator<
439            Item = impl IntoIterator<
440                Item = (
441                    Address,
442                    Option<Option<AccountInfo>>,
443                    impl IntoIterator<Item = (StorageKey, StorageValue)>,
444                ),
445            >,
446        >,
447        contracts: impl IntoIterator<Item = (B256, Bytecode)>,
448    ) -> Self {
449        // Create state from iterator.
450        let mut state_size = 0;
451        let state = state
452            .into_iter()
453            .map(|(address, original, present, storage)| {
454                let account = BundleAccount::new(
455                    original,
456                    present,
457                    storage
458                        .into_iter()
459                        .map(|(k, (o_val, p_val))| (k, StorageSlot::new_changed(o_val, p_val)))
460                        .collect(),
461                    AccountStatus::Changed,
462                );
463                state_size += account.size_hint();
464                (address, account)
465            })
466            .collect();
467
468        // Create reverts from iterator.
469        let mut reverts_size = 0;
470        let reverts = reverts
471            .into_iter()
472            .map(|block_reverts| {
473                block_reverts
474                    .into_iter()
475                    .map(|(address, account, storage)| {
476                        let account = match account {
477                            Some(Some(account)) => AccountInfoRevert::RevertTo(account),
478                            Some(None) => AccountInfoRevert::DeleteIt,
479                            None => AccountInfoRevert::DoNothing,
480                        };
481                        let revert = AccountRevert {
482                            account,
483                            storage: storage
484                                .into_iter()
485                                .map(|(k, v)| (k, RevertToSlot::Some(v)))
486                                .collect(),
487                            previous_status: AccountStatus::Changed,
488                            wipe_storage: false,
489                        };
490                        reverts_size += revert.size_hint();
491                        (address, revert)
492                    })
493                    .collect::<Vec<_>>()
494            })
495            .collect::<Vec<_>>();
496
497        Self {
498            state,
499            contracts: contracts.into_iter().collect(),
500            reverts: Reverts::new(reverts),
501            state_size,
502            reverts_size,
503        }
504    }
505
506    /// Returns the approximate size of changes in the bundle state.
507    ///
508    /// The estimation is not precise, because the information about the number of
509    /// destroyed entries that need to be removed is not accessible to the bundle state.
510    pub fn size_hint(&self) -> usize {
511        self.state_size + self.reverts_size + self.contracts.len()
512    }
513
514    /// Returns reference to the state.
515    pub fn state(&self) -> &HashMap<Address, BundleAccount> {
516        &self.state
517    }
518
519    /// Checks whether bundle state is empty.
520    pub fn is_empty(&self) -> bool {
521        self.len() == 0
522    }
523
524    /// Return number of changed accounts.
525    pub fn len(&self) -> usize {
526        self.state.len()
527    }
528
529    /// Gets account from state.
530    pub fn account(&self, address: &Address) -> Option<&BundleAccount> {
531        self.state.get(address)
532    }
533
534    /// Gets bytecode from state.
535    pub fn bytecode(&self, hash: &B256) -> Option<Bytecode> {
536        self.contracts.get(hash).cloned()
537    }
538
539    /// Consumes [`TransitionState`] by applying the changes and creating the
540    /// reverts.
541    ///
542    /// If [BundleRetention::includes_reverts] is `true`, then the reverts will
543    /// be retained.
544    pub fn apply_transitions_and_create_reverts(
545        &mut self,
546        transitions: TransitionState,
547        retention: BundleRetention,
548    ) {
549        let include_reverts = retention.includes_reverts();
550        // Pessimistically pre-allocate assuming _all_ accounts changed.
551        let reverts_capacity = if include_reverts {
552            transitions.transitions.len()
553        } else {
554            0
555        };
556        let mut reverts = Vec::with_capacity(reverts_capacity);
557
558        for (address, transition) in transitions.transitions.into_iter() {
559            // Add new contract if it was created/changed.
560            if let Some((hash, new_bytecode)) = transition.has_new_contract() {
561                self.contracts.insert(hash, new_bytecode.clone());
562            }
563            // Update state and create revert.
564            let revert = match self.state.entry(address) {
565                Entry::Occupied(mut entry) => {
566                    let entry = entry.get_mut();
567                    self.state_size -= entry.size_hint();
568                    // Update and create revert if it is present
569                    let revert = entry.update_and_create_revert(transition);
570                    // Update the state size
571                    self.state_size += entry.size_hint();
572                    revert
573                }
574                Entry::Vacant(entry) => {
575                    // Make revert from transition account
576                    let present_bundle = transition.present_bundle_account();
577                    let revert = transition.create_revert();
578                    if revert.is_some() {
579                        self.state_size += present_bundle.size_hint();
580                        entry.insert(present_bundle);
581                    }
582                    revert
583                }
584            };
585
586            // Append revert if present.
587            if let Some(revert) = revert.filter(|_| include_reverts) {
588                self.reverts_size += revert.size_hint();
589                reverts.push((address, revert));
590            }
591        }
592
593        self.reverts.push(reverts);
594    }
595
596    /// Generate a [`StateChangeset`] from the bundle state without consuming
597    /// it.
598    pub fn to_plain_state(&self, is_value_known: OriginalValuesKnown) -> StateChangeset {
599        // Pessimistically pre-allocate assuming _all_ accounts changed.
600        let state_len = self.state.len();
601        let mut accounts = Vec::with_capacity(state_len);
602        let mut storage = Vec::with_capacity(state_len);
603
604        for (address, account) in &self.state {
605            // Append account info if it is changed.
606            let was_destroyed = account.was_destroyed();
607            if is_value_known.is_not_known() || account.is_info_changed() {
608                let info = account.info.as_ref().map(AccountInfo::copy_without_code);
609                accounts.push((*address, info));
610            }
611
612            // Append storage changes
613
614            // Note: Assumption is that revert is going to remove whole plain storage from
615            // database so we can check if plain state was wiped or not.
616            let mut account_storage_changed = Vec::with_capacity(account.storage.len());
617
618            for (&key, &slot) in account.storage.iter() {
619                // If storage was destroyed that means that storage was wiped.
620                // In that case we need to check if present storage value is different then ZERO.
621                let destroyed_and_not_zero = was_destroyed && !slot.present_value.is_zero();
622
623                // If account is not destroyed check if original values was changed,
624                // so we can update it.
625                let not_destroyed_and_changed = !was_destroyed && slot.is_changed();
626
627                if is_value_known.is_not_known()
628                    || destroyed_and_not_zero
629                    || not_destroyed_and_changed
630                {
631                    account_storage_changed.push((key, slot.present_value));
632                }
633            }
634
635            if !account_storage_changed.is_empty() || was_destroyed {
636                // Append storage changes to account.
637                storage.push(PlainStorageChangeset {
638                    address: *address,
639                    wipe_storage: was_destroyed,
640                    storage: account_storage_changed,
641                });
642            }
643        }
644
645        let contracts = self
646            .contracts
647            .iter()
648            // Remove empty bytecodes
649            .filter(|(b, _)| **b != KECCAK_EMPTY)
650            .map(|(b, code)| (*b, code.clone()))
651            .collect::<Vec<_>>();
652        StateChangeset {
653            accounts,
654            storage,
655            contracts,
656        }
657    }
658
659    /// Generates a [`StateChangeset`] and [`PlainStateReverts`] from the bundle
660    /// state.
661    pub fn to_plain_state_and_reverts(
662        &self,
663        is_value_known: OriginalValuesKnown,
664    ) -> (StateChangeset, PlainStateReverts) {
665        (
666            self.to_plain_state(is_value_known),
667            self.reverts.to_plain_state_reverts(),
668        )
669    }
670
671    /// Consumes the bundle state and split it into a [`StateChangeset`] and a
672    /// [`PlainStateReverts`].
673    #[deprecated = "Use `to_plain_state_and_reverts` instead"]
674    pub fn into_plain_state_and_reverts(
675        self,
676        is_value_known: OriginalValuesKnown,
677    ) -> (StateChangeset, PlainStateReverts) {
678        self.to_plain_state_and_reverts(is_value_known)
679    }
680
681    /// Extends the bundle with other state.
682    ///
683    /// Updates the `other` state only if `other` is not flagged as destroyed.
684    pub fn extend_state(&mut self, other_state: HashMap<Address, BundleAccount>) {
685        for (address, other_account) in other_state {
686            match self.state.entry(address) {
687                Entry::Occupied(mut entry) => {
688                    let this = entry.get_mut();
689                    self.state_size -= this.size_hint();
690
691                    // If other was destroyed. replace `this` storage with
692                    // the `other one.
693                    if other_account.was_destroyed() {
694                        this.storage = other_account.storage;
695                    } else {
696                        // Otherwise extend this storage with other
697                        for (key, storage_slot) in other_account.storage {
698                            // Update present value or insert storage slot.
699                            this.storage
700                                .entry(key)
701                                .or_insert(storage_slot)
702                                .present_value = storage_slot.present_value;
703                        }
704                    }
705                    this.info = other_account.info;
706                    this.status.transition(other_account.status);
707
708                    // Update the state size
709                    self.state_size += this.size_hint();
710                }
711                Entry::Vacant(entry) => {
712                    // Just insert if empty
713                    self.state_size += other_account.size_hint();
714                    entry.insert(other_account);
715                }
716            }
717        }
718    }
719    /// Extends the state with state that is build on top of it.
720    ///
721    /// If storage was wiped in `other` state, copy `this` plain state
722    /// and put it inside `other` revert (if there is no duplicates of course).
723    ///
724    /// If `this` and `other` accounts were both destroyed invalidate second
725    /// wipe flag (from `other`).
726    ///
727    /// As wiping from database should be done only once
728    /// and we already transferred all potentially missing storages to the `other` revert.
729    pub fn extend(&mut self, mut other: Self) {
730        // Iterate over reverts and if its storage is wiped try to add previous bundle
731        // state as there is potential missing slots.
732        for (address, revert) in other.reverts.iter_mut().flatten() {
733            if revert.wipe_storage {
734                // If there is wipe storage in `other` revert
735                // we need to move storage from present state.
736                if let Some(this_account) = self.state.get_mut(address) {
737                    // As this account was destroyed inside `other` bundle.
738                    // We are fine to wipe/drain this storage and put it inside revert.
739                    for (key, value) in this_account.storage.drain() {
740                        revert
741                            .storage
742                            .entry(key)
743                            .or_insert(RevertToSlot::Some(value.present_value));
744                    }
745
746                    // Nullify `other` wipe as primary database wipe is done in `this`.
747                    if this_account.was_destroyed() {
748                        revert.wipe_storage = false;
749                    }
750                }
751            }
752
753            // Increment reverts size for each of the updated reverts.
754            self.reverts_size += revert.size_hint();
755        }
756        // Extension of state
757        self.extend_state(other.state);
758        // Contract can be just extended, when counter is introduced we will take into account that.
759        self.contracts.extend(other.contracts);
760        // Reverts can be just extended
761        self.reverts.extend(other.reverts);
762    }
763
764    /// Takes first N raw reverts from the [BundleState].
765    pub fn take_n_reverts(&mut self, reverts_to_take: usize) -> Reverts {
766        // Split is done as [0, num) and [num, len].
767        if reverts_to_take > self.reverts.len() {
768            return self.take_all_reverts();
769        }
770        let (detach, this) = self.reverts.split_at(reverts_to_take);
771        let detached_reverts = Reverts::new(detach.to_vec());
772        self.reverts_size = this
773            .iter()
774            .flatten()
775            .fold(0, |acc, (_, revert)| acc + revert.size_hint());
776        self.reverts = Reverts::new(this.to_vec());
777        detached_reverts
778    }
779
780    /// Returns and clears all reverts from [BundleState].
781    pub fn take_all_reverts(&mut self) -> Reverts {
782        self.reverts_size = 0;
783        mem::take(&mut self.reverts)
784    }
785
786    /// Reverts the state changes of the latest transition.
787    ///
788    /// **Note**: This is the same as `BundleState::revert(1)`
789    ///
790    /// Returns true if the state was reverted.
791    pub fn revert_latest(&mut self) -> bool {
792        // Revert the latest recorded state
793        if let Some(reverts) = self.reverts.pop() {
794            for (address, revert_account) in reverts.into_iter() {
795                self.reverts_size -= revert_account.size_hint();
796                match self.state.entry(address) {
797                    Entry::Occupied(mut entry) => {
798                        let account = entry.get_mut();
799                        self.state_size -= account.size_hint();
800                        if account.revert(revert_account) {
801                            entry.remove();
802                        } else {
803                            self.state_size += account.size_hint();
804                        }
805                    }
806                    Entry::Vacant(entry) => {
807                        // Create empty account that we will revert on.
808                        // Only place where this account is not existing is if revert is DeleteIt.
809                        let mut account = BundleAccount::new(
810                            None,
811                            None,
812                            HashMap::default(),
813                            AccountStatus::LoadedNotExisting,
814                        );
815                        if !account.revert(revert_account) {
816                            self.state_size += account.size_hint();
817                            entry.insert(account);
818                        }
819                    }
820                }
821            }
822            return true;
823        }
824
825        false
826    }
827
828    /// Reverts the state changes by N transitions back.
829    ///
830    /// See also [Self::revert_latest]
831    pub fn revert(&mut self, mut num_transitions: usize) {
832        if num_transitions == 0 {
833            return;
834        }
835
836        while self.revert_latest() {
837            num_transitions -= 1;
838            if num_transitions == 0 {
839                // Break the loop.
840                break;
841            }
842        }
843    }
844
845    /// Prepends present the state with the given [BundleState].
846    ///
847    /// It adds changes from the given state but does not override any existing changes.
848    ///
849    /// Reverts are not updated.
850    pub fn prepend_state(&mut self, mut other: BundleState) {
851        // Take this bundle
852        let this_bundle = mem::take(self);
853        // Extend other bundle state with this
854        other.extend_state(this_bundle.state);
855        // Extend other contracts
856        other.contracts.extend(this_bundle.contracts);
857        // Swap bundles
858        mem::swap(self, &mut other)
859    }
860}
861
862#[cfg(test)]
863mod tests {
864    use super::*;
865    use crate::{StorageWithOriginalValues, TransitionAccount};
866    use primitives::U256;
867
868    #[test]
869    fn transition_states() {
870        // Dummy data
871        let address = Address::new([0x01; 20]);
872        let acc1 = AccountInfo {
873            balance: U256::from(10),
874            nonce: 1,
875            code_hash: KECCAK_EMPTY,
876            code: None,
877        };
878
879        let mut bundle_state = BundleState::default();
880
881        // Have transition from loaded to all other states
882
883        let transition = TransitionAccount {
884            info: Some(acc1),
885            status: AccountStatus::InMemoryChange,
886            previous_info: None,
887            previous_status: AccountStatus::LoadedNotExisting,
888            storage: StorageWithOriginalValues::default(),
889            storage_was_destroyed: false,
890        };
891
892        // Apply first transition
893        bundle_state.apply_transitions_and_create_reverts(
894            TransitionState::single(address, transition.clone()),
895            BundleRetention::Reverts,
896        );
897    }
898
899    const fn account1() -> Address {
900        Address::new([0x60; 20])
901    }
902
903    const fn account2() -> Address {
904        Address::new([0x61; 20])
905    }
906
907    fn slot1() -> StorageKey {
908        StorageKey::from(5)
909    }
910
911    fn slot2() -> StorageKey {
912        StorageKey::from(7)
913    }
914
915    /// Tests bundle one.
916    fn test_bundle1() -> BundleState {
917        // Block changes
918        BundleState::new(
919            vec![
920                (
921                    account1(),
922                    None,
923                    Some(AccountInfo {
924                        nonce: 1,
925                        balance: U256::from(10),
926                        code_hash: KECCAK_EMPTY,
927                        code: None,
928                    }),
929                    HashMap::from_iter([
930                        (slot1(), (StorageValue::from(0), StorageValue::from(10))),
931                        (slot2(), (StorageValue::from(0), StorageValue::from(15))),
932                    ]),
933                ),
934                (
935                    account2(),
936                    None,
937                    Some(AccountInfo {
938                        nonce: 1,
939                        balance: U256::from(10),
940                        code_hash: KECCAK_EMPTY,
941                        code: None,
942                    }),
943                    HashMap::default(),
944                ),
945            ],
946            vec![vec![
947                (
948                    account1(),
949                    Some(None),
950                    vec![
951                        (slot1(), StorageValue::from(0)),
952                        (slot2(), StorageValue::from(0)),
953                    ],
954                ),
955                (account2(), Some(None), vec![]),
956            ]],
957            vec![],
958        )
959    }
960
961    /// Tests bundle two.
962    fn test_bundle2() -> BundleState {
963        // Block changes
964        BundleState::new(
965            vec![(
966                account1(),
967                None,
968                Some(AccountInfo {
969                    nonce: 3,
970                    balance: U256::from(20),
971                    code_hash: KECCAK_EMPTY,
972                    code: None,
973                }),
974                HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
975            )],
976            vec![vec![(
977                account1(),
978                Some(Some(AccountInfo {
979                    nonce: 1,
980                    balance: U256::from(10),
981                    code_hash: KECCAK_EMPTY,
982                    code: None,
983                })),
984                vec![(slot1(), StorageValue::from(10))],
985            )]],
986            vec![],
987        )
988    }
989
990    /// Tests bundle three.
991    fn test_bundle3() -> BundleState {
992        BundleState::builder(0..=0)
993            .state_present_account_info(
994                account1(),
995                AccountInfo {
996                    nonce: 1,
997                    balance: U256::from(10),
998                    code_hash: KECCAK_EMPTY,
999                    code: None,
1000                },
1001            )
1002            .state_storage(
1003                account1(),
1004                HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(10)))]),
1005            )
1006            .state_address(account2())
1007            .state_present_account_info(
1008                account2(),
1009                AccountInfo {
1010                    nonce: 1,
1011                    balance: U256::from(10),
1012                    code_hash: KECCAK_EMPTY,
1013                    code: None,
1014                },
1015            )
1016            .revert_address(0, account1())
1017            .revert_account_info(0, account1(), Some(None))
1018            .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(0))])
1019            .revert_account_info(0, account2(), Some(None))
1020            .build()
1021    }
1022
1023    /// Tests bundle four.
1024    fn test_bundle4() -> BundleState {
1025        BundleState::builder(0..=0)
1026            .state_present_account_info(
1027                account1(),
1028                AccountInfo {
1029                    nonce: 3,
1030                    balance: U256::from(20),
1031                    code_hash: KECCAK_EMPTY,
1032                    code: None,
1033                },
1034            )
1035            .state_storage(
1036                account1(),
1037                HashMap::from_iter([(slot1(), (StorageValue::from(0), StorageValue::from(15)))]),
1038            )
1039            .revert_address(0, account1())
1040            .revert_account_info(
1041                0,
1042                account1(),
1043                Some(Some(AccountInfo {
1044                    nonce: 1,
1045                    balance: U256::from(10),
1046                    code_hash: KECCAK_EMPTY,
1047                    code: None,
1048                })),
1049            )
1050            .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1051            .build()
1052    }
1053
1054    fn sanity_path(bundle1: BundleState, bundle2: BundleState) {
1055        let mut extended = bundle1.clone();
1056        extended.extend(bundle2.clone());
1057
1058        let mut reverted = extended.clone();
1059        // Revert zero does nothing.
1060        reverted.revert(0);
1061        assert_eq!(reverted, extended);
1062
1063        // Revert by one gives us bundle one.
1064        reverted.revert(1);
1065        assert_eq!(reverted, bundle1);
1066
1067        // Reverted by additional one gives us empty bundle.
1068        reverted.revert(1);
1069        assert_eq!(reverted, BundleState::default());
1070
1071        let mut reverted = extended.clone();
1072
1073        // Reverted by bigger number gives us empty bundle
1074        reverted.revert(10);
1075        assert_eq!(reverted, BundleState::default());
1076    }
1077
1078    #[test]
1079    fn extend_on_destroyed_values() {
1080        let base_bundle1 = test_bundle1();
1081        let base_bundle2 = test_bundle2();
1082
1083        // test1
1084        // bundle1 has Destroyed
1085        // bundle2 has Changed
1086        // end should be DestroyedChanged.
1087        let mut b1 = base_bundle1.clone();
1088        let mut b2 = base_bundle2.clone();
1089        b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1090        b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1091        b1.extend(b2);
1092        assert_eq!(
1093            b1.state.get_mut(&account1()).unwrap().status,
1094            AccountStatus::DestroyedChanged
1095        );
1096
1097        // test2
1098        // bundle1 has Changed
1099        // bundle2 has Destroyed
1100        // end should be Destroyed
1101        let mut b1 = base_bundle1.clone();
1102        let mut b2 = base_bundle2.clone();
1103        b1.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1104        b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Destroyed;
1105        b2.reverts[0][0].1.wipe_storage = true;
1106        b1.extend(b2);
1107        assert_eq!(
1108            b1.state.get_mut(&account1()).unwrap().status,
1109            AccountStatus::Destroyed
1110        );
1111
1112        // test2 extension
1113        // revert of b2 should contains plain state of b1.
1114        let mut revert1 = base_bundle2.reverts[0][0].clone();
1115        revert1.1.wipe_storage = true;
1116        revert1
1117            .1
1118            .storage
1119            .insert(slot2(), RevertToSlot::Some(StorageValue::from(15)));
1120
1121        assert_eq!(
1122            b1.reverts.as_ref(),
1123            vec![base_bundle1.reverts[0].clone(), vec![revert1]],
1124        );
1125
1126        // test3
1127        // bundle1 has InMemoryChange
1128        // bundle2 has Change
1129        // end should be InMemoryChange.
1130
1131        let mut b1 = base_bundle1.clone();
1132        let mut b2 = base_bundle2.clone();
1133        b1.state.get_mut(&account1()).unwrap().status = AccountStatus::InMemoryChange;
1134        b2.state.get_mut(&account1()).unwrap().status = AccountStatus::Changed;
1135        b1.extend(b2);
1136        assert_eq!(
1137            b1.state.get_mut(&account1()).unwrap().status,
1138            AccountStatus::InMemoryChange
1139        );
1140    }
1141
1142    #[test]
1143    fn test_sanity_path() {
1144        sanity_path(test_bundle1(), test_bundle2());
1145        sanity_path(test_bundle3(), test_bundle4());
1146    }
1147
1148    #[test]
1149    fn test_multi_reverts_with_delete() {
1150        let mut state = BundleBuilder::new(0..=3)
1151            .revert_address(0, account1())
1152            .revert_account_info(2, account1(), Some(Some(AccountInfo::default())))
1153            .revert_account_info(3, account1(), Some(None))
1154            .build();
1155
1156        state.revert_latest();
1157        // State for account one was deleted
1158        assert_eq!(state.state.get(&account1()), None);
1159
1160        state.revert_latest();
1161        // State is set to
1162        assert_eq!(
1163            state.state.get(&account1()),
1164            Some(&BundleAccount::new(
1165                None,
1166                Some(AccountInfo::default()),
1167                HashMap::default(),
1168                AccountStatus::Changed
1169            ))
1170        );
1171    }
1172
1173    #[test]
1174    fn test_revert_capacity() {
1175        let state = BundleState::builder(0..=3)
1176            .revert_address(0, account1())
1177            .revert_address(2, account2())
1178            .revert_account_info(0, account1(), Some(None))
1179            .revert_account_info(2, account2(), None)
1180            .revert_storage(0, account1(), vec![(slot1(), StorageValue::from(10))])
1181            .build();
1182
1183        assert_eq!(state.reverts.len(), 4);
1184        assert_eq!(state.reverts[1], vec![]);
1185        assert_eq!(state.reverts[3], vec![]);
1186        assert_eq!(state.reverts[0].len(), 1);
1187        assert_eq!(state.reverts[2].len(), 1);
1188
1189        let (addr1, revert1) = &state.reverts[0][0];
1190        assert_eq!(addr1, &account1());
1191        assert_eq!(revert1.account, AccountInfoRevert::DeleteIt);
1192
1193        let (addr2, revert2) = &state.reverts[2][0];
1194        assert_eq!(addr2, &account2());
1195        assert_eq!(revert2.account, AccountInfoRevert::DoNothing);
1196    }
1197
1198    #[test]
1199    fn take_reverts() {
1200        let bundle1 = test_bundle1();
1201        let bundle2 = test_bundle2();
1202
1203        let mut extended = bundle1.clone();
1204        extended.extend(bundle2.clone());
1205
1206        // Check that we have two reverts
1207        assert_eq!(extended.reverts.len(), 2);
1208
1209        // Take all by big N
1210        let mut extended2 = extended.clone();
1211        assert_eq!(extended2.take_n_reverts(100), extended.reverts);
1212
1213        // Take all reverts
1214        let mut extended2 = extended.clone();
1215        assert_eq!(extended2.take_all_reverts(), extended.reverts);
1216
1217        // Take zero revert
1218        let taken_reverts = extended.take_n_reverts(0);
1219        assert_eq!(taken_reverts, Reverts::default());
1220        assert_eq!(extended.reverts.len(), 2);
1221
1222        // Take one revert
1223        let taken_reverts = extended.take_n_reverts(1);
1224        assert_eq!(taken_reverts, bundle1.reverts);
1225
1226        // Take last revert
1227        let taken_reverts = extended.take_n_reverts(1);
1228        assert_eq!(taken_reverts, bundle2.reverts);
1229    }
1230
1231    #[test]
1232    fn prepend_state() {
1233        let address1 = account1();
1234        let address2 = account2();
1235
1236        let account1 = AccountInfo {
1237            nonce: 1,
1238            ..Default::default()
1239        };
1240        let account1_changed = AccountInfo {
1241            nonce: 1,
1242            ..Default::default()
1243        };
1244        let account2 = AccountInfo {
1245            nonce: 1,
1246            ..Default::default()
1247        };
1248
1249        let present_state = BundleState::builder(2..=2)
1250            .state_present_account_info(address1, account1_changed.clone())
1251            .build();
1252        assert_eq!(present_state.reverts.len(), 1);
1253        let previous_state = BundleState::builder(1..=1)
1254            .state_present_account_info(address1, account1)
1255            .state_present_account_info(address2, account2.clone())
1256            .build();
1257        assert_eq!(previous_state.reverts.len(), 1);
1258
1259        let mut test = present_state;
1260
1261        test.prepend_state(previous_state);
1262
1263        assert_eq!(test.state.len(), 2);
1264        // Reverts num should stay the same.
1265        assert_eq!(test.reverts.len(), 1);
1266        // Account1 is not overwritten.
1267        assert_eq!(
1268            test.state.get(&address1).unwrap().info,
1269            Some(account1_changed)
1270        );
1271        // Account2 got inserted
1272        assert_eq!(test.state.get(&address2).unwrap().info, Some(account2));
1273    }
1274
1275    #[test]
1276    fn test_getters() {
1277        let mut builder = BundleBuilder::new(0..=3);
1278
1279        // Test get_states and get_states_mut
1280        assert!(builder.get_states().is_empty());
1281        builder.get_states_mut().insert(account1());
1282        assert!(builder.get_states().contains(&account1()));
1283
1284        // Test get_state_original_mut
1285        assert!(builder.get_state_original_mut().is_empty());
1286        builder
1287            .get_state_original_mut()
1288            .insert(account1(), AccountInfo::default());
1289        assert!(builder.get_state_original_mut().contains_key(&account1()));
1290
1291        // Test get_state_present_mut
1292        assert!(builder.get_state_present_mut().is_empty());
1293        builder
1294            .get_state_present_mut()
1295            .insert(account1(), AccountInfo::default());
1296        assert!(builder.get_state_present_mut().contains_key(&account1()));
1297
1298        // Test get_state_storage_mut
1299        assert!(builder.get_state_storage_mut().is_empty());
1300        builder
1301            .get_state_storage_mut()
1302            .insert(account1(), HashMap::default());
1303        assert!(builder.get_state_storage_mut().contains_key(&account1()));
1304
1305        // Test get_reverts_mut
1306        assert!(builder.get_reverts_mut().is_empty());
1307        builder.get_reverts_mut().insert((0, account1()));
1308        assert!(builder.get_reverts_mut().contains(&(0, account1())));
1309
1310        // Test get_revert_range_mut
1311        assert_eq!(builder.get_revert_range_mut().clone(), 0..=3);
1312
1313        // Test get_revert_account_mut
1314        assert!(builder.get_revert_account_mut().is_empty());
1315        builder
1316            .get_revert_account_mut()
1317            .insert((0, account1()), Some(None));
1318        assert!(builder
1319            .get_revert_account_mut()
1320            .contains_key(&(0, account1())));
1321
1322        // Test get_revert_storage_mut
1323        assert!(builder.get_revert_storage_mut().is_empty());
1324        builder
1325            .get_revert_storage_mut()
1326            .insert((0, account1()), vec![(slot1(), StorageValue::from(0))]);
1327        assert!(builder
1328            .get_revert_storage_mut()
1329            .contains_key(&(0, account1())));
1330
1331        // Test get_contracts_mut
1332        assert!(builder.get_contracts_mut().is_empty());
1333        builder
1334            .get_contracts_mut()
1335            .insert(B256::default(), Bytecode::default());
1336        assert!(builder.get_contracts_mut().contains_key(&B256::default()));
1337    }
1338}