revm_database/
in_memory_db.rs

1use core::convert::Infallible;
2use database_interface::{Database, DatabaseCommit, DatabaseRef, EmptyDB};
3use primitives::{address, hash_map::Entry, Address, HashMap, Log, B256, KECCAK_EMPTY, U256};
4use state::{Account, AccountInfo, Bytecode};
5use std::vec::Vec;
6
7/// A [Database] implementation that stores all state changes in memory.
8pub type InMemoryDB = CacheDB<EmptyDB>;
9
10/// A [Database] implementation that stores all state changes in memory.
11///
12/// This implementation wraps a [DatabaseRef] that is used to load data ([AccountInfo]).
13///
14/// Accounts and code are stored in two separate maps, the `accounts` map maps addresses to [DbAccount],
15/// whereas contracts are identified by their code hash, and are stored in the `contracts` map.
16/// The [DbAccount] holds the code hash of the contract, which is used to look up the contract in the `contracts` map.
17#[derive(Debug, Clone)]
18#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
19pub struct CacheDB<ExtDB> {
20    /// Account info where None means it is not existing. Not existing state is needed for Pre TANGERINE forks.
21    /// `code` is always `None`, and bytecode can be found in `contracts`.
22    pub accounts: HashMap<Address, DbAccount>,
23    /// Tracks all contracts by their code hash.
24    pub contracts: HashMap<B256, Bytecode>,
25    /// All logs that were committed via [DatabaseCommit::commit].
26    pub logs: Vec<Log>,
27    /// All cached block hashes from the [DatabaseRef].
28    pub block_hashes: HashMap<U256, B256>,
29    /// The underlying database ([DatabaseRef]) that is used to load data.
30    ///
31    /// Note: This is read-only, data is never written to this database.
32    pub db: ExtDB,
33}
34
35impl<ExtDB: Default> Default for CacheDB<ExtDB> {
36    fn default() -> Self {
37        Self::new(ExtDB::default())
38    }
39}
40
41impl<ExtDb> CacheDB<CacheDB<ExtDb>> {
42    /// Flattens a nested cache by applying the outer cache to the inner cache.
43    ///
44    /// The behavior is as follows:
45    /// - Accounts are overridden with outer accounts
46    /// - Contracts are overridden with outer contracts
47    /// - Logs are appended
48    /// - Block hashes are overridden with outer block hashes
49    pub fn flatten(self) -> CacheDB<ExtDb> {
50        let CacheDB {
51            accounts,
52            contracts,
53            logs,
54            block_hashes,
55            db: mut inner,
56        } = self;
57
58        inner.accounts.extend(accounts);
59        inner.contracts.extend(contracts);
60        inner.logs.extend(logs);
61        inner.block_hashes.extend(block_hashes);
62        inner
63    }
64
65    /// Discards the outer cache and return the inner cache.
66    pub fn discard_outer(self) -> CacheDB<ExtDb> {
67        self.db
68    }
69}
70
71impl<ExtDB> CacheDB<ExtDB> {
72    /// Creates a new cache with the given external database.
73    pub fn new(db: ExtDB) -> Self {
74        let mut contracts = HashMap::default();
75        contracts.insert(KECCAK_EMPTY, Bytecode::default());
76        contracts.insert(B256::ZERO, Bytecode::default());
77        Self {
78            accounts: HashMap::default(),
79            contracts,
80            logs: Vec::default(),
81            block_hashes: HashMap::default(),
82            db,
83        }
84    }
85
86    /// Inserts the account's code into the cache.
87    ///
88    /// Accounts objects and code are stored separately in the cache, this will take the code from the account and instead map it to the code hash.
89    ///
90    /// Note: This will not insert into the underlying external database.
91    pub fn insert_contract(&mut self, account: &mut AccountInfo) {
92        if let Some(code) = &account.code {
93            if !code.is_empty() {
94                if account.code_hash == KECCAK_EMPTY {
95                    account.code_hash = code.hash_slow();
96                }
97                self.contracts
98                    .entry(account.code_hash)
99                    .or_insert_with(|| code.clone());
100            }
101        }
102        if account.code_hash.is_zero() {
103            account.code_hash = KECCAK_EMPTY;
104        }
105    }
106
107    /// Inserts account info but not override storage
108    pub fn insert_account_info(&mut self, address: Address, mut info: AccountInfo) {
109        self.insert_contract(&mut info);
110        self.accounts.entry(address).or_default().info = info;
111    }
112
113    /// Wraps the cache in a [CacheDB], creating a nested cache.
114    pub fn nest(self) -> CacheDB<Self> {
115        CacheDB::new(self)
116    }
117}
118
119impl<ExtDB: DatabaseRef> CacheDB<ExtDB> {
120    /// Returns the account for the given address.
121    ///
122    /// If the account was not found in the cache, it will be loaded from the underlying database.
123    pub fn load_account(&mut self, address: Address) -> Result<&mut DbAccount, ExtDB::Error> {
124        let db = &self.db;
125        match self.accounts.entry(address) {
126            Entry::Occupied(entry) => Ok(entry.into_mut()),
127            Entry::Vacant(entry) => Ok(entry.insert(
128                db.basic_ref(address)?
129                    .map(|info| DbAccount {
130                        info,
131                        ..Default::default()
132                    })
133                    .unwrap_or_else(DbAccount::new_not_existing),
134            )),
135        }
136    }
137
138    /// Inserts account storage without overriding account info
139    pub fn insert_account_storage(
140        &mut self,
141        address: Address,
142        slot: U256,
143        value: U256,
144    ) -> Result<(), ExtDB::Error> {
145        let account = self.load_account(address)?;
146        account.storage.insert(slot, value);
147        Ok(())
148    }
149
150    /// Replaces account storage without overriding account info
151    pub fn replace_account_storage(
152        &mut self,
153        address: Address,
154        storage: HashMap<U256, U256>,
155    ) -> Result<(), ExtDB::Error> {
156        let account = self.load_account(address)?;
157        account.account_state = AccountState::StorageCleared;
158        account.storage = storage.into_iter().collect();
159        Ok(())
160    }
161}
162
163impl<ExtDB> DatabaseCommit for CacheDB<ExtDB> {
164    fn commit(&mut self, changes: HashMap<Address, Account>) {
165        for (address, mut account) in changes {
166            if !account.is_touched() {
167                continue;
168            }
169            if account.is_selfdestructed() {
170                let db_account = self.accounts.entry(address).or_default();
171                db_account.storage.clear();
172                db_account.account_state = AccountState::NotExisting;
173                db_account.info = AccountInfo::default();
174                continue;
175            }
176            let is_newly_created = account.is_created();
177            self.insert_contract(&mut account.info);
178
179            let db_account = self.accounts.entry(address).or_default();
180            db_account.info = account.info;
181
182            db_account.account_state = if is_newly_created {
183                db_account.storage.clear();
184                AccountState::StorageCleared
185            } else if db_account.account_state.is_storage_cleared() {
186                // Preserve old account state if it already exists
187                AccountState::StorageCleared
188            } else {
189                AccountState::Touched
190            };
191            db_account.storage.extend(
192                account
193                    .storage
194                    .into_iter()
195                    .map(|(key, value)| (key, value.present_value())),
196            );
197        }
198    }
199}
200
201impl<ExtDB: DatabaseRef> Database for CacheDB<ExtDB> {
202    type Error = ExtDB::Error;
203
204    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
205        let basic = match self.accounts.entry(address) {
206            Entry::Occupied(entry) => entry.into_mut(),
207            Entry::Vacant(entry) => entry.insert(
208                self.db
209                    .basic_ref(address)?
210                    .map(|info| DbAccount {
211                        info,
212                        ..Default::default()
213                    })
214                    .unwrap_or_else(DbAccount::new_not_existing),
215            ),
216        };
217        Ok(basic.info())
218    }
219
220    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
221        match self.contracts.entry(code_hash) {
222            Entry::Occupied(entry) => Ok(entry.get().clone()),
223            Entry::Vacant(entry) => {
224                // If you return code bytes when basic fn is called this function is not needed.
225                Ok(entry.insert(self.db.code_by_hash_ref(code_hash)?).clone())
226            }
227        }
228    }
229
230    /// Get the value in an account's storage slot.
231    ///
232    /// It is assumed that account is already loaded.
233    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
234        match self.accounts.entry(address) {
235            Entry::Occupied(mut acc_entry) => {
236                let acc_entry = acc_entry.get_mut();
237                match acc_entry.storage.entry(index) {
238                    Entry::Occupied(entry) => Ok(*entry.get()),
239                    Entry::Vacant(entry) => {
240                        if matches!(
241                            acc_entry.account_state,
242                            AccountState::StorageCleared | AccountState::NotExisting
243                        ) {
244                            Ok(U256::ZERO)
245                        } else {
246                            let slot = self.db.storage_ref(address, index)?;
247                            entry.insert(slot);
248                            Ok(slot)
249                        }
250                    }
251                }
252            }
253            Entry::Vacant(acc_entry) => {
254                // Acc needs to be loaded for us to access slots.
255                let info = self.db.basic_ref(address)?;
256                let (account, value) = if info.is_some() {
257                    let value = self.db.storage_ref(address, index)?;
258                    let mut account: DbAccount = info.into();
259                    account.storage.insert(index, value);
260                    (account, value)
261                } else {
262                    (info.into(), U256::ZERO)
263                };
264                acc_entry.insert(account);
265                Ok(value)
266            }
267        }
268    }
269
270    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
271        match self.block_hashes.entry(U256::from(number)) {
272            Entry::Occupied(entry) => Ok(*entry.get()),
273            Entry::Vacant(entry) => {
274                let hash = self.db.block_hash_ref(number)?;
275                entry.insert(hash);
276                Ok(hash)
277            }
278        }
279    }
280}
281
282impl<ExtDB: DatabaseRef> DatabaseRef for CacheDB<ExtDB> {
283    type Error = ExtDB::Error;
284
285    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
286        match self.accounts.get(&address) {
287            Some(acc) => Ok(acc.info()),
288            None => self.db.basic_ref(address),
289        }
290    }
291
292    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
293        match self.contracts.get(&code_hash) {
294            Some(entry) => Ok(entry.clone()),
295            None => self.db.code_by_hash_ref(code_hash),
296        }
297    }
298
299    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
300        match self.accounts.get(&address) {
301            Some(acc_entry) => match acc_entry.storage.get(&index) {
302                Some(entry) => Ok(*entry),
303                None => {
304                    if matches!(
305                        acc_entry.account_state,
306                        AccountState::StorageCleared | AccountState::NotExisting
307                    ) {
308                        Ok(U256::ZERO)
309                    } else {
310                        self.db.storage_ref(address, index)
311                    }
312                }
313            },
314            None => self.db.storage_ref(address, index),
315        }
316    }
317
318    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
319        match self.block_hashes.get(&U256::from(number)) {
320            Some(entry) => Ok(*entry),
321            None => self.db.block_hash_ref(number),
322        }
323    }
324}
325
326#[derive(Debug, Clone, Default)]
327#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
328pub struct DbAccount {
329    pub info: AccountInfo,
330    /// If account is selfdestructed or newly created, storage will be cleared.
331    pub account_state: AccountState,
332    /// Storage slots
333    pub storage: HashMap<U256, U256>,
334}
335
336impl DbAccount {
337    pub fn new_not_existing() -> Self {
338        Self {
339            account_state: AccountState::NotExisting,
340            ..Default::default()
341        }
342    }
343
344    pub fn info(&self) -> Option<AccountInfo> {
345        if matches!(self.account_state, AccountState::NotExisting) {
346            None
347        } else {
348            Some(self.info.clone())
349        }
350    }
351}
352
353impl From<Option<AccountInfo>> for DbAccount {
354    fn from(from: Option<AccountInfo>) -> Self {
355        from.map(Self::from).unwrap_or_else(Self::new_not_existing)
356    }
357}
358
359impl From<AccountInfo> for DbAccount {
360    fn from(info: AccountInfo) -> Self {
361        Self {
362            info,
363            account_state: AccountState::None,
364            ..Default::default()
365        }
366    }
367}
368
369#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
370#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
371pub enum AccountState {
372    /// Before Spurious Dragon hardfork there was a difference between empty and not existing.
373    /// And we are flagging it here.
374    NotExisting,
375    /// EVM touched this account. For newer hardfork this means it can be cleared/removed from state.
376    Touched,
377    /// EVM cleared storage of this account, mostly by selfdestruct, we don't ask database for storage slots
378    /// and assume they are U256::ZERO
379    StorageCleared,
380    /// EVM didn't interacted with this account
381    #[default]
382    None,
383}
384
385impl AccountState {
386    /// Returns `true` if EVM cleared storage of this account
387    pub fn is_storage_cleared(&self) -> bool {
388        matches!(self, AccountState::StorageCleared)
389    }
390}
391
392/// Custom benchmarking DB that only has account info for the zero address.
393///
394/// Any other address will return an empty account.
395#[derive(Debug, Default, Clone)]
396pub struct BenchmarkDB(pub Bytecode, B256);
397
398impl BenchmarkDB {
399    pub fn new_bytecode(bytecode: Bytecode) -> Self {
400        let hash = bytecode.hash_slow();
401        Self(bytecode, hash)
402    }
403}
404
405/// BYTECODE address
406pub const FFADDRESS: Address = address!("ffffffffffffffffffffffffffffffffffffffff");
407pub const BENCH_TARGET: Address = FFADDRESS;
408/// CALLER address
409pub const EEADDRESS: Address = address!("eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee");
410pub const BENCH_CALLER: Address = EEADDRESS;
411
412impl Database for BenchmarkDB {
413    type Error = Infallible;
414    /// Get basic account information.
415    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
416        if address == FFADDRESS {
417            return Ok(Some(AccountInfo {
418                nonce: 1,
419                balance: U256::from(10000000),
420                code: Some(self.0.clone()),
421                code_hash: self.1,
422            }));
423        }
424        if address == EEADDRESS {
425            return Ok(Some(AccountInfo {
426                nonce: 0,
427                balance: U256::from(10000000),
428                code: None,
429                code_hash: KECCAK_EMPTY,
430            }));
431        }
432        Ok(None)
433    }
434
435    /// Get account code by its hash
436    fn code_by_hash(&mut self, _code_hash: B256) -> Result<Bytecode, Self::Error> {
437        Ok(Bytecode::default())
438    }
439
440    /// Get storage value of address at index.
441    fn storage(&mut self, _address: Address, _index: U256) -> Result<U256, Self::Error> {
442        Ok(U256::default())
443    }
444
445    // History related
446    fn block_hash(&mut self, _number: u64) -> Result<B256, Self::Error> {
447        Ok(B256::default())
448    }
449}
450
451#[cfg(test)]
452mod tests {
453    use super::{CacheDB, EmptyDB};
454    use database_interface::Database;
455    use primitives::{Address, HashMap, U256};
456    use state::AccountInfo;
457
458    #[test]
459    fn test_insert_account_storage() {
460        let account = Address::with_last_byte(42);
461        let nonce = 42;
462        let mut init_state = CacheDB::new(EmptyDB::default());
463        init_state.insert_account_info(
464            account,
465            AccountInfo {
466                nonce,
467                ..Default::default()
468            },
469        );
470
471        let (key, value) = (U256::from(123), U256::from(456));
472        let mut new_state = CacheDB::new(init_state);
473        new_state
474            .insert_account_storage(account, key, value)
475            .unwrap();
476
477        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce);
478        assert_eq!(new_state.storage(account, key), Ok(value));
479    }
480
481    #[test]
482    fn test_replace_account_storage() {
483        let account = Address::with_last_byte(42);
484        let nonce = 42;
485        let mut init_state = CacheDB::new(EmptyDB::default());
486        init_state.insert_account_info(
487            account,
488            AccountInfo {
489                nonce,
490                ..Default::default()
491            },
492        );
493
494        let (key0, value0) = (U256::from(123), U256::from(456));
495        let (key1, value1) = (U256::from(789), U256::from(999));
496        init_state
497            .insert_account_storage(account, key0, value0)
498            .unwrap();
499
500        let mut new_state = CacheDB::new(init_state);
501        new_state
502            .replace_account_storage(account, HashMap::from_iter([(key1, value1)]))
503            .unwrap();
504
505        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce);
506        assert_eq!(new_state.storage(account, key0), Ok(U256::ZERO));
507        assert_eq!(new_state.storage(account, key1), Ok(value1));
508    }
509
510    #[cfg(feature = "serde")]
511    #[test]
512    fn test_serialize_deserialize_cachedb() {
513        let account = Address::with_last_byte(69);
514        let nonce = 420;
515        let mut init_state = CacheDB::new(EmptyDB::default());
516        init_state.insert_account_info(
517            account,
518            AccountInfo {
519                nonce,
520                ..Default::default()
521            },
522        );
523
524        let serialized = serde_json::to_string(&init_state).unwrap();
525        let deserialized: CacheDB<EmptyDB> = serde_json::from_str(&serialized).unwrap();
526
527        assert!(deserialized.accounts.contains_key(&account));
528        assert_eq!(
529            deserialized.accounts.get(&account).unwrap().info.nonce,
530            nonce
531        );
532    }
533}
revm_database/in_memory_db.rs

revm_database/
in_memory_db.rs
