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 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}