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::{BlockchainDb, DatabaseResult, RevertStateSnapshotAction, StateSnapshot},
11    fork::database::ForkDbStateSnapshot,
12};
13use revm::{
14    context::BlockEnv,
15    database::{Database, DbAccount},
16    state::AccountInfo,
17};
18
19pub use foundry_evm::fork::database::ForkedDatabase;
20
21impl<N: Network> Db for ForkedDatabase<N> {
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        cache_block_hash(&mut self.inner().block_hashes().write(), number, hash);
34    }
35
36    fn set_block_hashes(&mut self, block_hashes: Vec<(U256, B256)>) {
37        *self.inner().block_hashes().write() = block_hashes.into_iter().collect();
38    }
39
40    fn dump_state(
41        &self,
42        at: BlockEnv,
43        best_number: u64,
44        blocks: Vec<SerializableBlock>,
45        transactions: Vec<SerializableTransaction>,
46        historical_states: Option<SerializableHistoricalStates>,
47    ) -> DatabaseResult<Option<SerializableState>> {
48        let mut db = self.database().clone();
49        let accounts = self
50            .database()
51            .cache
52            .accounts
53            .clone()
54            .into_iter()
55            .map(|(k, v)| -> DatabaseResult<_> {
56                let code = if let Some(code) = v.info.code {
57                    code
58                } else {
59                    db.code_by_hash(v.info.code_hash)?
60                };
61                Ok((
62                    k,
63                    SerializableAccountRecord {
64                        nonce: v.info.nonce,
65                        balance: v.info.balance,
66                        code: code.original_bytes(),
67                        storage: v.storage.into_iter().map(|(k, v)| (k.into(), v.into())).collect(),
68                    },
69                ))
70            })
71            .collect::<Result<_, _>>()?;
72        Ok(Some(SerializableState {
73            block: Some(at),
74            accounts,
75            best_block_number: Some(best_number),
76            blocks,
77            transactions,
78            historical_states,
79        }))
80    }
81
82    fn snapshot_state(&mut self) -> U256 {
83        self.insert_state_snapshot()
84    }
85
86    fn revert_state(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
87        self.revert_state_snapshot(id, action)
88    }
89
90    fn current_state(&self) -> StateDb {
91        StateDb::new(self.create_state_snapshot())
92    }
93}
94
95impl<N: Network> MaybeFullDatabase for ForkedDatabase<N> {
96    fn maybe_as_full_db(&self) -> Option<&AddressMap<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<N: Network> MaybeFullDatabase for ForkDbStateSnapshot<N> {
131    fn maybe_as_full_db(&self) -> Option<&AddressMap<DbAccount>> {
132        Some(&self.local.cache.accounts)
133    }
134
135    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
136        let mut state_snapshot = std::mem::take(&mut self.state_snapshot);
137        let local_state_snapshot = self.local.clear_into_state_snapshot();
138        state_snapshot.accounts.extend(local_state_snapshot.accounts);
139        state_snapshot.storage.extend(local_state_snapshot.storage);
140        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
141        state_snapshot
142    }
143
144    fn read_as_state_snapshot(&self) -> StateSnapshot {
145        let mut state_snapshot = self.state_snapshot.clone();
146        let local_state_snapshot = self.local.read_as_state_snapshot();
147        state_snapshot.accounts.extend(local_state_snapshot.accounts);
148        state_snapshot.storage.extend(local_state_snapshot.storage);
149        state_snapshot.block_hashes.extend(local_state_snapshot.block_hashes);
150        state_snapshot
151    }
152
153    fn clear(&mut self) {
154        std::mem::take(&mut self.state_snapshot);
155        self.local.clear()
156    }
157
158    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
159        self.state_snapshot = state_snapshot;
160    }
161}
162
163impl<N: Network> MaybeForkedDatabase for ForkedDatabase<N> {
164    fn maybe_reset(&mut self, urls: Vec<String>, block_number: BlockId) -> Result<(), String> {
165        self.reset(urls, block_number)
166    }
167
168    fn maybe_flush_cache(&self) -> Result<(), String> {
169        self.flush_cache();
170        Ok(())
171    }
172
173    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
174        Ok(self.inner())
175    }
176}
anvil/eth/backend/mem/fork_db.rs

anvil/eth/backend/mem/
fork_db.rs
