anvil/eth/backend/mem/
fork_db.rs

Help
1use crate::{
2    eth::backend::db::{
3        Db, MaybeForkedDatabase, MaybeFullDatabase, SerializableAccountRecord, SerializableBlock,
4        SerializableHistoricalStates, SerializableState, SerializableTransaction, StateDb,
5    },
6    revm::primitives::AccountInfo,
7};
8use alloy_primitives::{map::HashMap, Address, B256, U256, U64};
9use alloy_rpc_types::BlockId;
10use foundry_evm::{
11    backend::{
12        BlockchainDb, DatabaseError, DatabaseResult, RevertStateSnapshotAction, StateSnapshot,
13    },
14    fork::database::ForkDbStateSnapshot,
15    revm::{primitives::BlockEnv, Database},
16};
17use revm::{db::DbAccount, DatabaseRef};
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            .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 MaybeFullDatabase for ForkedDatabase {
91    fn as_dyn(&self) -> &dyn DatabaseRef<Error = DatabaseError> {
92        self
93    }
94
95    fn maybe_as_full_db(&self) -> Option<&HashMap<Address, DbAccount>> {
96        Some(&self.database().accounts)
97    }
98
99    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
100        let db = self.inner().db();
101        let accounts = std::mem::take(&mut *db.accounts.write());
102        let storage = std::mem::take(&mut *db.storage.write());
103        let block_hashes = std::mem::take(&mut *db.block_hashes.write());
104        StateSnapshot { accounts, storage, block_hashes }
105    }
106
107    fn read_as_state_snapshot(&self) -> StateSnapshot {
108        let db = self.inner().db();
109        let accounts = db.accounts.read().clone();
110        let storage = db.storage.read().clone();
111        let block_hashes = db.block_hashes.read().clone();
112        StateSnapshot { accounts, storage, block_hashes }
113    }
114
115    fn clear(&mut self) {
116        self.flush_cache();
117        self.clear_into_state_snapshot();
118    }
119
120    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
121        let db = self.inner().db();
122        let StateSnapshot { accounts, storage, block_hashes } = state_snapshot;
123        *db.accounts.write() = accounts;
124        *db.storage.write() = storage;
125        *db.block_hashes.write() = block_hashes;
126    }
127}
128
129impl MaybeFullDatabase for ForkDbStateSnapshot {
130    fn as_dyn(&self) -> &dyn DatabaseRef<Error = DatabaseError> {
131        self
132    }
133
134    fn maybe_as_full_db(&self) -> Option<&HashMap<Address, DbAccount>> {
135        Some(&self.local.accounts)
136    }
137
138    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
139        std::mem::take(&mut self.state_snapshot)
140    }
141
142    fn read_as_state_snapshot(&self) -> StateSnapshot {
143        self.state_snapshot.clone()
144    }
145
146    fn clear(&mut self) {
147        std::mem::take(&mut self.state_snapshot);
148        self.local.clear()
149    }
150
151    fn init_from_state_snapshot(&mut self, state_snapshot: StateSnapshot) {
152        self.state_snapshot = state_snapshot;
153    }
154}
155
156impl MaybeForkedDatabase for ForkedDatabase {
157    fn maybe_reset(&mut self, url: Option<String>, block_number: BlockId) -> Result<(), String> {
158        self.reset(url, block_number)
159    }
160
161    fn maybe_flush_cache(&self) -> Result<(), String> {
162        self.flush_cache();
163        Ok(())
164    }
165
166    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
167        Ok(self.inner())
168    }
169}
anvil/eth/backend/mem/fork_db.rs

anvil/eth/backend/mem/
fork_db.rs
