kabu_evm_db/
fast_cache_db.rs

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

kabu_evm_db/
fast_cache_db.rs
