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

anvil/eth/backend/mem/
fork_db.rs
