anvil/eth/backend/mem/
in_memory_db.rs

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

anvil/eth/backend/mem/
in_memory_db.rs
