Skip to main content

anvil/eth/backend/mem/
fork_db.rs

1use crate::eth::backend::db::{
2    Db, MaybeForkedDatabase, MaybeFullDatabase, SerializableAccountRecord, SerializableBlock,
3    SerializableHistoricalStates, SerializableState, SerializableTransaction, StateDb,
4    cache_block_hash,
5};
6use alloy_network::Network;
7use alloy_primitives::{Address, B256, U256, map::AddressMap};
8use alloy_rpc_types::BlockId;
9use foundry_evm::{
10    backend::{
11        BlockchainDb, DatabaseResult, MemDb, RevertStateSnapshotAction, SharedBackend,
12        StateSnapshot,
13    },
14    fork::database::ForkDbStateSnapshot,
15};
16use revm::{
17    context::BlockEnv,
18    database::{CacheDB, Database, DbAccount},
19    state::AccountInfo,
20};
21
22pub use foundry_evm::fork::database::ForkedDatabase;
23
24impl<N: Network> MaybeFullDatabase for SharedBackend<N> {
25    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
26        StateSnapshot::default()
27    }
28
29    fn read_as_state_snapshot(&self) -> StateSnapshot {
30        StateSnapshot::default()
31    }
32
33    fn clear(&mut self) {}
34
35    fn init_from_state_snapshot(&mut self, _state_snapshot: StateSnapshot) {}
36}
37
38impl<N: Network> Db for ForkedDatabase<N> {
39    fn insert_account(&mut self, address: Address, account: AccountInfo) {
40        self.database_mut().insert_account(address, account)
41    }
42
43    fn set_storage_at(&mut self, address: Address, slot: B256, val: B256) -> DatabaseResult<()> {
44        // this ensures the account is loaded first
45        let _ = Database::basic(self, address)?;
46        self.database_mut().set_storage_at(address, slot, val)
47    }
48
49    fn insert_block_hash(&mut self, number: U256, hash: B256) {
50        cache_block_hash(&mut self.database_mut().cache.block_hashes, number, hash);
51    }
52
53    fn set_block_hashes(&mut self, block_hashes: Vec<(U256, B256)>) {
54        self.database_mut().cache.block_hashes = block_hashes.into_iter().collect();
55    }
56
57    fn dump_state(
58        &self,
59        at: BlockEnv,
60        best_number: u64,
61        blocks: Vec<SerializableBlock>,
62        transactions: Vec<SerializableTransaction>,
63        historical_states: Option<SerializableHistoricalStates>,
64    ) -> DatabaseResult<Option<SerializableState>> {
65        let mut db = self.database().clone();
66        let accounts = self
67            .database()
68            .cache
69            .accounts
70            .clone()
71            .into_iter()
72            .map(|(k, v)| -> DatabaseResult<_> {
73                let code = if let Some(code) = v.info.code {
74                    code
75                } else {
76                    db.code_by_hash(v.info.code_hash())?
77                };
78                Ok((
79                    k,
80                    SerializableAccountRecord {
81                        nonce: v.info.nonce,
82                        balance: v.info.balance,
83                        code: code.original_bytes(),
84                        storage: v.storage.into_iter().map(|(k, v)| (k.into(), v.into())).collect(),
85                    },
86                ))
87            })
88            .collect::<Result<_, _>>()?;
89        Ok(Some(SerializableState {
90            block: Some(at),
91            accounts,
92            best_block_number: Some(best_number),
93            blocks,
94            transactions,
95            #[cfg(feature = "monad")]
96            monad_block_participants: Default::default(),
97            #[cfg(feature = "monad")]
98            monad_block_replay_profiles: Default::default(),
99            historical_states,
100        }))
101    }
102
103    fn snapshot_state(&mut self) -> U256 {
104        self.insert_state_snapshot()
105    }
106
107    fn revert_state(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
108        self.revert_state_snapshot(id, action)
109    }
110
111    fn delete_state_snapshot(&mut self, id: U256) -> bool {
112        self.state_snapshots().lock().remove_at(id).is_some()
113    }
114
115    fn current_state(&self) -> StateDb {
116        StateDb::new(self.create_state_snapshot())
117    }
118
119    fn historical_state(&self, state_snapshot: StateSnapshot) -> StateDb {
120        // Fall back to the fork for anything the snapshot did not capture. Snapshots without
121        // account states cannot tell cleared storage from uncached storage, so they keep the
122        // previous standalone restore.
123        let mut state = if state_snapshot.account_states.is_some() {
124            StateDb::new(CacheDB::new(self.database().db.clone()))
125        } else {
126            StateDb::new(MemDb::default())
127        };
128        state.init_from_state_snapshot(state_snapshot);
129        state
130    }
131}
132
133impl<N: Network> MaybeFullDatabase for ForkedDatabase<N> {
134    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
135        Some(&self.database().cache.accounts)
136    }
137
138    fn maybe_as_full_db_mut(&mut self) -> Option<&mut AddressMap<DbAccount>> {
139        Some(&mut self.database_mut().cache.accounts)
140    }
141
142    fn maybe_full_db(&self) -> Option<AddressMap<DbAccount>> {
143        None
144    }
145
146    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
147        let db = self.inner().db();
148        let accounts = std::mem::take(&mut *db.accounts.write());
149        let storage = std::mem::take(&mut *db.storage.write());
150        let block_hashes = std::mem::take(&mut *db.block_hashes.write());
151        StateSnapshot { accounts, storage, block_hashes, ..Default::default() }
152    }
153
154    fn read_as_state_snapshot(&self) -> StateSnapshot {
155        let db = self.inner().db();
156        let accounts = db.accounts.read().clone();
157        let storage = db.storage.read().clone();
158        let block_hashes = db.block_hashes.read().clone();
159        StateSnapshot { accounts, storage, block_hashes, ..Default::default() }
160    }
161
162    fn clear(&mut self) {
163        self.flush_cache();
164        self.clear_into_state_snapshot();
165    }
166
167    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
168        let db = self.inner().db();
169        let StateSnapshot { accounts, storage, block_hashes, .. } = state_snapshot;
170        *db.accounts.write() = accounts;
171        *db.storage.write() = storage;
172        *db.block_hashes.write() = block_hashes;
173    }
174}
175
176impl<N: Network> MaybeFullDatabase for ForkDbStateSnapshot<N> {
177    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
178        Some(&self.local.cache.accounts)
179    }
180
181    fn maybe_full_db(&self) -> Option<AddressMap<DbAccount>> {
182        None
183    }
184
185    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
186        let mut state_snapshot = std::mem::take(&mut self.state_snapshot);
187        let local_state_snapshot = self.local.clear_into_state_snapshot();
188        state_snapshot.accounts.extend(local_state_snapshot.accounts);
189        state_snapshot.storage.extend(local_state_snapshot.storage);
190        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
191        if let Some(account_states) = local_state_snapshot.account_states {
192            state_snapshot.account_states.get_or_insert_default().extend(account_states);
193        }
194        state_snapshot
195    }
196
197    fn read_as_state_snapshot(&self) -> StateSnapshot {
198        let mut state_snapshot = self.state_snapshot.clone();
199        let local_state_snapshot = self.local.read_as_state_snapshot();
200        state_snapshot.accounts.extend(local_state_snapshot.accounts);
201        state_snapshot.storage.extend(local_state_snapshot.storage);
202        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
203        if let Some(account_states) = local_state_snapshot.account_states {
204            state_snapshot.account_states.get_or_insert_default().extend(account_states);
205        }
206        state_snapshot
207    }
208
209    fn clear(&mut self) {
210        std::mem::take(&mut self.state_snapshot);
211        self.local.clear()
212    }
213
214    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
215        self.state_snapshot = state_snapshot;
216    }
217}
218
219impl<N: Network> MaybeForkedDatabase for ForkedDatabase<N> {
220    fn maybe_reset(&mut self, urls: Vec<String>, block_number: BlockId) -> Result<(), String> {
221        self.reset(urls, block_number)
222    }
223
224    fn maybe_flush_cache(&self) -> Result<(), String> {
225        self.flush_cache();
226        Ok(())
227    }
228
229    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
230        Ok(self.inner())
231    }
232}