foundry_evm_core/backend/
snapshot.rs

1use crate::backend::JournaledState;
2use alloy_evm::EvmEnv;
3use alloy_primitives::{
4    B256, U256,
5    map::{AddressHashMap, U256Map},
6};
7use revm::{database::AccountState, state::AccountInfo};
8use serde::{Deserialize, Serialize, Serializer, ser::SerializeStruct};
9use std::collections::BTreeMap;
10
11/// A minimal abstraction of a state at a certain point in time
12#[derive(Clone, Debug, Default, Deserialize)]
13pub struct StateSnapshot {
14    pub accounts: AddressHashMap<AccountInfo>,
15    pub storage: AddressHashMap<U256Map<U256>>,
16    pub block_hashes: U256Map<B256>,
17    /// Accounts that are [`AccountState::NotExisting`] or [`AccountState::StorageCleared`], whose
18    /// storage missing from `storage` is zero rather than read from the underlying database.
19    ///
20    /// `None` if the snapshot does not track account states, such as one saved by an older
21    /// version.
22    #[serde(default)]
23    pub account_states: Option<AddressHashMap<AccountState>>,
24}
25
26impl Serialize for StateSnapshot {
27    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
28    where
29        S: Serializer,
30    {
31        let accounts = self.accounts.iter().collect::<BTreeMap<_, _>>();
32        let storage = self
33            .storage
34            .iter()
35            .map(|(address, storage)| (address, storage.iter().collect::<BTreeMap<_, _>>()))
36            .collect::<BTreeMap<_, _>>();
37        let block_hashes = self.block_hashes.iter().collect::<BTreeMap<_, _>>();
38        let account_states = self
39            .account_states
40            .as_ref()
41            .map(|account_states| account_states.iter().collect::<BTreeMap<_, _>>());
42
43        let len = if account_states.is_some() { 4 } else { 3 };
44        let mut state = serializer.serialize_struct("StateSnapshot", len)?;
45        state.serialize_field("accounts", &accounts)?;
46        state.serialize_field("storage", &storage)?;
47        state.serialize_field("block_hashes", &block_hashes)?;
48        if let Some(account_states) = account_states {
49            state.serialize_field("account_states", &account_states)?;
50        }
51        state.end()
52    }
53}
54
55/// Represents a state snapshot taken during evm execution
56#[derive(Clone, Debug)]
57pub struct BackendStateSnapshot<T, SPEC, BLOCK> {
58    pub db: T,
59    /// Complete context state at a specific point.
60    pub journaled_state: JournaledState,
61    /// Contains the evm env at the time of the snapshot
62    pub snap_evm_env: EvmEnv<SPEC, BLOCK>,
63}
64
65impl<T, SPEC, BLOCK> BackendStateSnapshot<T, SPEC, BLOCK> {
66    /// Takes a new state snapshot.
67    pub const fn new(db: T, journaled_state: JournaledState, evm_env: EvmEnv<SPEC, BLOCK>) -> Self {
68        Self { db, journaled_state, snap_evm_env: evm_env }
69    }
70
71    /// Called when this state snapshot is reverted.
72    ///
73    /// Since we want to keep all additional logs that were emitted since the snapshot was taken
74    /// we'll merge additional logs into the snapshot's `revm::JournaledState`. Additional logs are
75    /// those logs that are missing in the snapshot's journaled_state, since the current
76    /// journaled_state includes the same logs, we can simply replace use that See also
77    /// `DatabaseExt::revert`.
78    pub fn merge(&mut self, current: &JournaledState) {
79        self.journaled_state.logs.clone_from(&current.logs);
80    }
81}
82
83/// What to do when reverting a state snapshot.
84///
85/// Whether to remove the state snapshot or keep it.
86#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
87pub enum RevertStateSnapshotAction {
88    /// Remove the state snapshot after reverting.
89    #[default]
90    RevertRemove,
91    /// Keep the state snapshot after reverting.
92    RevertKeep,
93}
94
95impl RevertStateSnapshotAction {
96    /// Returns `true` if the action is to keep the state snapshot.
97    pub const fn is_keep(&self) -> bool {
98        matches!(self, Self::RevertKeep)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use alloy_primitives::Address;
106
107    #[test]
108    fn state_snapshot_serializes_maps_in_key_order() {
109        let low_address = Address::with_last_byte(1);
110        let high_address = Address::with_last_byte(2);
111        let low_slot = U256::ONE;
112        let high_slot = U256::from(2);
113        let mut snapshot = StateSnapshot::default();
114
115        snapshot.accounts.insert(high_address, AccountInfo::from_balance(U256::from(2)));
116        snapshot.accounts.insert(low_address, AccountInfo::from_balance(U256::ONE));
117        snapshot.storage.insert(
118            high_address,
119            [(high_slot, U256::from(200)), (low_slot, U256::from(100))].into_iter().collect(),
120        );
121        snapshot.storage.insert(
122            low_address,
123            [(high_slot, U256::from(200)), (low_slot, U256::from(100))].into_iter().collect(),
124        );
125        snapshot.block_hashes.insert(high_slot, B256::with_last_byte(200));
126        snapshot.block_hashes.insert(low_slot, B256::with_last_byte(100));
127
128        let json = serde_json::to_string(&snapshot).unwrap();
129        let storage_start = json.find("\"storage\"").unwrap();
130        let block_hashes_start = json.find("\"block_hashes\"").unwrap();
131        let accounts = &json[..storage_start];
132        let storage = &json[storage_start..block_hashes_start];
133        let block_hashes = &json[block_hashes_start..];
134        let low_address = serde_json::to_string(&low_address).unwrap();
135        let high_address = serde_json::to_string(&high_address).unwrap();
136        let low_slot = serde_json::to_string(&low_slot).unwrap();
137        let high_slot = serde_json::to_string(&high_slot).unwrap();
138
139        assert!(accounts.find(&low_address).unwrap() < accounts.find(&high_address).unwrap());
140        let low_storage_start = storage.find(&low_address).unwrap();
141        let high_storage_start = storage.find(&high_address).unwrap();
142        assert!(low_storage_start < high_storage_start);
143        let low_storage = &storage[low_storage_start..high_storage_start];
144        let high_storage = &storage[high_storage_start..];
145        assert!(low_storage.find(&low_slot).unwrap() < low_storage.find(&high_slot).unwrap());
146        assert!(high_storage.find(&low_slot).unwrap() < high_storage.find(&high_slot).unwrap());
147        assert!(block_hashes.find(&low_slot).unwrap() < block_hashes.find(&high_slot).unwrap());
148    }
149}
foundry_evm_core/backend/snapshot.rs

foundry_evm_core/backend/
snapshot.rs
