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