revm_database/
in_memory_db.rs

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

revm_database/
in_memory_db.rs
