anvil/eth/backend/mem/
in_memory_db.rs

1//! The in memory DB
2
3use crate::{
4    eth::backend::db::{
5        Db, MaybeForkedDatabase, MaybeFullDatabase, SerializableAccountRecord, SerializableBlock,
6        SerializableHistoricalStates, SerializableState, SerializableTransaction, StateDb,
7    },
8    mem::state::state_root,
9};
10use alloy_primitives::{Address, B256, U256, map::HashMap};
11use alloy_rpc_types::BlockId;
12use foundry_evm::backend::{BlockchainDb, DatabaseResult, StateSnapshot};
13use revm::{
14    context::BlockEnv,
15    database::{DatabaseRef, DbAccount},
16    state::AccountInfo,
17};
18
19// reexport for convenience
20pub use foundry_evm::backend::MemDb;
21use foundry_evm::backend::RevertStateSnapshotAction;
22
23impl Db for MemDb {
24    fn insert_account(&mut self, address: Address, account: AccountInfo) {
25        self.inner.insert_account_info(address, account)
26    }
27
28    fn set_storage_at(&mut self, address: Address, slot: B256, val: B256) -> DatabaseResult<()> {
29        self.inner.insert_account_storage(address, slot.into(), val.into())
30    }
31
32    fn insert_block_hash(&mut self, number: U256, hash: B256) {
33        self.inner.cache.block_hashes.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 accounts = self
45            .inner
46            .cache
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                    self.inner.code_by_hash_ref(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
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    /// Creates a new snapshot
79    fn snapshot_state(&mut self) -> U256 {
80        let id = self.state_snapshots.insert(self.inner.clone());
81        trace!(target: "backend::memdb", "Created new state snapshot {}", id);
82        id
83    }
84
85    fn revert_state(&mut self, id: U256, action: RevertStateSnapshotAction) -> bool {
86        if let Some(state_snapshot) = self.state_snapshots.remove(id) {
87            if action.is_keep() {
88                self.state_snapshots.insert_at(state_snapshot.clone(), id);
89            }
90            self.inner = state_snapshot;
91            trace!(target: "backend::memdb", "Reverted state snapshot {}", id);
92            true
93        } else {
94            warn!(target: "backend::memdb", "No state snapshot to revert for {}", id);
95            false
96        }
97    }
98
99    fn maybe_state_root(&self) -> Option<B256> {
100        Some(state_root(&self.inner.cache.accounts))
101    }
102
103    fn current_state(&self) -> StateDb {
104        StateDb::new(Self { inner: self.inner.clone(), ..Default::default() })
105    }
106}
107
108impl MaybeFullDatabase for MemDb {
109    fn maybe_as_full_db(&self) -> Option<&HashMap<Address, DbAccount>> {
110        Some(&self.inner.cache.accounts)
111    }
112
113    fn clear_into_state_snapshot(&mut self) -> StateSnapshot {
114        self.inner.clear_into_state_snapshot()
115    }
116
117    fn read_as_state_snapshot(&self) -> StateSnapshot {
118        self.inner.read_as_state_snapshot()
119    }
120
121    fn clear(&mut self) {
122        self.inner.clear();
123    }
124
125    fn init_from_state_snapshot(&mut self, snapshot: StateSnapshot) {
126        self.inner.init_from_state_snapshot(snapshot)
127    }
128}
129
130impl MaybeForkedDatabase for MemDb {
131    fn maybe_reset(&mut self, _url: Option<String>, _block_number: BlockId) -> Result<(), String> {
132        Err("not supported".to_string())
133    }
134
135    fn maybe_flush_cache(&self) -> Result<(), String> {
136        Err("not supported".to_string())
137    }
138
139    fn maybe_inner(&self) -> Result<&BlockchainDb, String> {
140        Err("not supported".to_string())
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use alloy_primitives::{Bytes, address};
148    use revm::{bytecode::Bytecode, primitives::KECCAK_EMPTY};
149    use std::collections::BTreeMap;
150
151    // verifies that all substantial aspects of a loaded account remain the same after an account
152    // is dumped and reloaded
153    #[test]
154    fn test_dump_reload_cycle() {
155        let test_addr: Address = address!("0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266");
156
157        let mut dump_db = MemDb::default();
158
159        let contract_code = Bytecode::new_raw(Bytes::from("fake contract code"));
160        dump_db.insert_account(
161            test_addr,
162            AccountInfo {
163                balance: U256::from(123456),
164                code_hash: KECCAK_EMPTY,
165                code: Some(contract_code.clone()),
166                nonce: 1234,
167            },
168        );
169        dump_db
170            .set_storage_at(test_addr, U256::from(1234567).into(), U256::from(1).into())
171            .unwrap();
172
173        // blocks dumping/loading tested in storage.rs
174        let state = dump_db
175            .dump_state(Default::default(), 0, Vec::new(), Vec::new(), Default::default())
176            .unwrap()
177            .unwrap();
178
179        let mut load_db = MemDb::default();
180
181        load_db.load_state(state).unwrap();
182
183        let loaded_account = load_db.basic_ref(test_addr).unwrap().unwrap();
184
185        assert_eq!(loaded_account.balance, U256::from(123456));
186        assert_eq!(load_db.code_by_hash_ref(loaded_account.code_hash).unwrap(), contract_code);
187        assert_eq!(loaded_account.nonce, 1234);
188        assert_eq!(load_db.storage_ref(test_addr, U256::from(1234567)).unwrap(), U256::from(1));
189    }
190
191    // verifies that multiple accounts can be loaded at a time, and storage is merged within those
192    // accounts as well.
193    #[test]
194    fn test_load_state_merge() {
195        let test_addr: Address = address!("0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266");
196        let test_addr2: Address = address!("0x70997970c51812dc3a010c7d01b50e0d17dc79c8");
197
198        let contract_code = Bytecode::new_raw(Bytes::from("fake contract code"));
199
200        let mut db = MemDb::default();
201
202        db.insert_account(
203            test_addr,
204            AccountInfo {
205                balance: U256::from(123456),
206                code_hash: KECCAK_EMPTY,
207                code: Some(contract_code.clone()),
208                nonce: 1234,
209            },
210        );
211
212        db.set_storage_at(test_addr, U256::from(1234567).into(), U256::from(1).into()).unwrap();
213        db.set_storage_at(test_addr, U256::from(1234568).into(), U256::from(2).into()).unwrap();
214
215        let mut new_state = SerializableState::default();
216
217        new_state.accounts.insert(
218            test_addr2,
219            SerializableAccountRecord {
220                balance: Default::default(),
221                code: Default::default(),
222                nonce: 1,
223                storage: Default::default(),
224            },
225        );
226
227        let mut new_storage = BTreeMap::default();
228        new_storage.insert(U256::from(1234568).into(), U256::from(5).into());
229
230        new_state.accounts.insert(
231            test_addr,
232            SerializableAccountRecord {
233                balance: U256::from(100100),
234                code: contract_code.bytes()[..contract_code.len()].to_vec().into(),
235                nonce: 100,
236                storage: new_storage,
237            },
238        );
239
240        db.load_state(new_state).unwrap();
241
242        let loaded_account = db.basic_ref(test_addr).unwrap().unwrap();
243        let loaded_account2 = db.basic_ref(test_addr2).unwrap().unwrap();
244
245        assert_eq!(loaded_account2.nonce, 1);
246
247        assert_eq!(loaded_account.balance, U256::from(100100));
248        assert_eq!(db.code_by_hash_ref(loaded_account.code_hash).unwrap(), contract_code);
249        assert_eq!(loaded_account.nonce, 1234);
250        assert_eq!(db.storage_ref(test_addr, U256::from(1234567)).unwrap(), U256::from(1));
251        assert_eq!(db.storage_ref(test_addr, U256::from(1234568)).unwrap(), U256::from(5));
252    }
253}
anvil/eth/backend/mem/in_memory_db.rs

anvil/eth/backend/mem/
in_memory_db.rs
