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::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}
18
19impl Serialize for StateSnapshot {
20    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
21    where
22        S: Serializer,
23    {
24        let accounts = self.accounts.iter().collect::<BTreeMap<_, _>>();
25        let storage = self
26            .storage
27            .iter()
28            .map(|(address, storage)| (address, storage.iter().collect::<BTreeMap<_, _>>()))
29            .collect::<BTreeMap<_, _>>();
30        let block_hashes = self.block_hashes.iter().collect::<BTreeMap<_, _>>();
31
32        let mut state = serializer.serialize_struct("StateSnapshot", 3)?;
33        state.serialize_field("accounts", &accounts)?;
34        state.serialize_field("storage", &storage)?;
35        state.serialize_field("block_hashes", &block_hashes)?;
36        state.end()
37    }
38}
39
40/// Represents a state snapshot taken during evm execution
41#[derive(Clone, Debug)]
42pub struct BackendStateSnapshot<T, SPEC, BLOCK> {
43    pub db: T,
44    /// Complete context state at a specific point.
45    pub journaled_state: JournaledState,
46    /// Contains the evm env at the time of the snapshot
47    pub snap_evm_env: EvmEnv<SPEC, BLOCK>,
48}
49
50impl<T, SPEC, BLOCK> BackendStateSnapshot<T, SPEC, BLOCK> {
51    /// Takes a new state snapshot.
52    pub const fn new(db: T, journaled_state: JournaledState, evm_env: EvmEnv<SPEC, BLOCK>) -> Self {
53        Self { db, journaled_state, snap_evm_env: evm_env }
54    }
55
56    /// Called when this state snapshot is reverted.
57    ///
58    /// Since we want to keep all additional logs that were emitted since the snapshot was taken
59    /// we'll merge additional logs into the snapshot's `revm::JournaledState`. Additional logs are
60    /// those logs that are missing in the snapshot's journaled_state, since the current
61    /// journaled_state includes the same logs, we can simply replace use that See also
62    /// `DatabaseExt::revert`.
63    pub fn merge(&mut self, current: &JournaledState) {
64        self.journaled_state.logs.clone_from(&current.logs);
65    }
66}
67
68/// What to do when reverting a state snapshot.
69///
70/// Whether to remove the state snapshot or keep it.
71#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
72pub enum RevertStateSnapshotAction {
73    /// Remove the state snapshot after reverting.
74    #[default]
75    RevertRemove,
76    /// Keep the state snapshot after reverting.
77    RevertKeep,
78}
79
80impl RevertStateSnapshotAction {
81    /// Returns `true` if the action is to keep the state snapshot.
82    pub const fn is_keep(&self) -> bool {
83        matches!(self, Self::RevertKeep)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use alloy_primitives::Address;
91
92    #[test]
93    fn state_snapshot_serializes_maps_in_key_order() {
94        let low_address = Address::from_word(B256::from(U256::from(1)));
95        let high_address = Address::from_word(B256::from(U256::from(2)));
96        let low_slot = U256::from(1);
97        let high_slot = U256::from(2);
98        let mut snapshot = StateSnapshot::default();
99
100        snapshot.accounts.insert(high_address, AccountInfo::from_balance(U256::from(2)));
101        snapshot.accounts.insert(low_address, AccountInfo::from_balance(U256::from(1)));
102        snapshot.storage.insert(
103            high_address,
104            [(high_slot, U256::from(200)), (low_slot, U256::from(100))].into_iter().collect(),
105        );
106        snapshot.storage.insert(
107            low_address,
108            [(high_slot, U256::from(200)), (low_slot, U256::from(100))].into_iter().collect(),
109        );
110        snapshot.block_hashes.insert(high_slot, B256::from(U256::from(200)));
111        snapshot.block_hashes.insert(low_slot, B256::from(U256::from(100)));
112
113        let json = serde_json::to_string(&snapshot).unwrap();
114        let storage_start = json.find("\"storage\"").unwrap();
115        let block_hashes_start = json.find("\"block_hashes\"").unwrap();
116        let accounts = &json[..storage_start];
117        let storage = &json[storage_start..block_hashes_start];
118        let block_hashes = &json[block_hashes_start..];
119        let low_address = serde_json::to_string(&low_address).unwrap();
120        let high_address = serde_json::to_string(&high_address).unwrap();
121        let low_slot = serde_json::to_string(&low_slot).unwrap();
122        let high_slot = serde_json::to_string(&high_slot).unwrap();
123
124        assert!(accounts.find(&low_address).unwrap() < accounts.find(&high_address).unwrap());
125        let low_storage_start = storage.find(&low_address).unwrap();
126        let high_storage_start = storage.find(&high_address).unwrap();
127        assert!(low_storage_start < high_storage_start);
128        let low_storage = &storage[low_storage_start..high_storage_start];
129        let high_storage = &storage[high_storage_start..];
130        assert!(low_storage.find(&low_slot).unwrap() < low_storage.find(&high_slot).unwrap());
131        assert!(high_storage.find(&low_slot).unwrap() < high_storage.find(&high_slot).unwrap());
132        assert!(block_hashes.find(&low_slot).unwrap() < block_hashes.find(&high_slot).unwrap());
133    }
134}
foundry_evm_core/backend/snapshot.rs

foundry_evm_core/backend/
snapshot.rs
