anvil/eth/backend/mem/
fork_db.rs

1use crate::eth::backend::db::{
2    Db, MaybeForkedDatabase, MaybeFullDatabase, SerializableAccountRecord, SerializableBlock,
3    SerializableHistoricalStates, SerializableState, SerializableTransaction, StateDb,
4};
5use alloy_network::Network;
6use alloy_primitives::{Address, B256, U256, map::AddressMap};
7use alloy_rpc_types::BlockId;
8use foundry_evm::{
9    backend::{BlockchainDb, DatabaseResult, RevertStateSnapshotAction, StateSnapshot},
10    fork::database::ForkDbStateSnapshot,
11};
12use revm::{
13    context::BlockEnv,
14    database::{Database, DbAccount},
15    state::AccountInfo,
16};
17
18pub use foundry_evm::fork::database::ForkedDatabase;
19
20impl<N: Network> Db for ForkedDatabase<N> {
21    fn insert_account(&mut self, address: Address, account: AccountInfo) {
22        self.database_mut().insert_account(address, account)
23    }
24
25    fn set_storage_at(&mut self, address: Address, slot: B256, val: B256) -> DatabaseResult<()> {
26        // this ensures the account is loaded first
27        let _ = Database::basic(self, address)?;
28        self.database_mut().set_storage_at(address, slot, val)
29    }
30
31    fn insert_block_hash(&mut self, number: U256, hash: B256) {
32        self.inner().block_hashes().write().insert(number, hash);
33    }
34
35    fn dump_state(
36        &self,
37        at: BlockEnv,
38        best_number: u64,
39        blocks: Vec<SerializableBlock>,
40        transactions: Vec<SerializableTransaction>,
41        historical_states: Option<SerializableHistoricalStates>,
42    ) -> DatabaseResult<Option<SerializableState>> {
43        let mut db = self.database().clone();
44        let accounts = self
45            .database()
46            .cache
47            .accounts
48            .clone()
49            .into_iter()
50            .map(|(k, v)| -> DatabaseResult<_> {
51                let code = if let Some(code) = v.info.code {
52                    code
53                } else {
54                    db.code_by_hash(v.info.code_hash)?
55                };
56                Ok((
57                    k,
58                    SerializableAccountRecord {
59                        nonce: v.info.nonce,
60                        balance: v.info.balance,
61                        code: code.original_bytes(),
62                        storage: v.storage.into_iter().map(|(k, v)| (k.into(), v.into())).collect(),
63                    },
64                ))
65            })
66            .collect::<Result<_, _>>()?;
67        Ok(Some(SerializableState {
68            block: Some(at),
69            accounts,
70            best_block_number: Some(best_number),
71            blocks,
72            transactions,
73            historical_states,
74        }))
75    }
76
77    fn snapshot_state(&mut self) -> U256 {
78        self.insert_state_snapshot()
79    }
80
81    fn revert_state(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
82        self.revert_state_snapshot(id, action)
83    }
84
85    fn current_state(&self) -> StateDb {
86        StateDb::new(self.create_state_snapshot())
87    }
88}
89
90impl<N: Network> MaybeFullDatabase for ForkedDatabase<N> {
91    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
92        Some(&self.database().cache.accounts)
93    }
94
95    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
96        let db = self.inner().db();
97        let accounts = std::mem::take(&mut *db.accounts.write());
98        let storage = std::mem::take(&mut *db.storage.write());
99        let block_hashes = std::mem::take(&mut *db.block_hashes.write());
100        StateSnapshot { accounts, storage, block_hashes }
101    }
102
103    fn read_as_state_snapshot(&self) -> StateSnapshot {
104        let db = self.inner().db();
105        let accounts = db.accounts.read().clone();
106        let storage = db.storage.read().clone();
107        let block_hashes = db.block_hashes.read().clone();
108        StateSnapshot { accounts, storage, block_hashes }
109    }
110
111    fn clear(&mut self) {
112        self.flush_cache();
113        self.clear_into_state_snapshot();
114    }
115
116    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
117        let db = self.inner().db();
118        let StateSnapshot { accounts, storage, block_hashes } = state_snapshot;
119        *db.accounts.write() = accounts;
120        *db.storage.write() = storage;
121        *db.block_hashes.write() = block_hashes;
122    }
123}
124
125impl<N: Network> MaybeFullDatabase for ForkDbStateSnapshot<N> {
126    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
127        Some(&self.local.cache.accounts)
128    }
129
130    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
131        let mut state_snapshot = std::mem::take(&mut self.state_snapshot);
132        let local_state_snapshot = self.local.clear_into_state_snapshot();
133        state_snapshot.accounts.extend(local_state_snapshot.accounts);
134        state_snapshot.storage.extend(local_state_snapshot.storage);
135        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
136        state_snapshot
137    }
138
139    fn read_as_state_snapshot(&self) -> StateSnapshot {
140        let mut state_snapshot = self.state_snapshot.clone();
141        let local_state_snapshot = self.local.read_as_state_snapshot();
142        state_snapshot.accounts.extend(local_state_snapshot.accounts);
143        state_snapshot.storage.extend(local_state_snapshot.storage);
144        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
145        state_snapshot
146    }
147
148    fn clear(&mut self) {
149        std::mem::take(&mut self.state_snapshot);
150        self.local.clear()
151    }
152
153    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
154        self.state_snapshot = state_snapshot;
155    }
156}
157
158impl<N: Network> MaybeForkedDatabase for ForkedDatabase<N> {
159    fn maybe_reset(&mut self, url: Option<String>, block_number: BlockId) -> Result<(), String> {
160        self.reset(url, block_number)
161    }
162
163    fn maybe_flush_cache(&self) -> Result<(), String> {
164        self.flush_cache();
165        Ok(())
166    }
167
168    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
169        Ok(self.inner())
170    }
171}
anvil/eth/backend/mem/fork_db.rs

anvil/eth/backend/mem/
fork_db.rs
