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