kabu_evm_db/
in_memory_db.rs

1use crate::fast_cache_db::FastCacheDB;
2use alloy::primitives::Address;
3use alloy::rpc::types::trace::geth::AccountState as GethAccountState;
4use revm::database::{AccountState, EmptyDB};
5use revm::state::Bytecode;
6use std::collections::BTreeMap;
7use std::sync::Arc;
8use tracing::{error, trace};
9
10pub type LoomInMemoryDB = FastCacheDB<Arc<FastCacheDB<EmptyDB>>>;
11
12impl LoomInMemoryDB {
13    pub fn with_db(self, db: Arc<FastCacheDB<EmptyDB>>) -> Self {
14        Self { db, ..self }
15    }
16
17    pub fn merge(&self) -> FastCacheDB<EmptyDB> {
18        let mut db: FastCacheDB<EmptyDB> = FastCacheDB {
19            accounts: self.db.accounts.clone(),
20            logs: self.db.logs.clone(),
21            contracts: self.db.contracts.clone(),
22            block_hashes: self.db.block_hashes.clone(),
23            db: self.db.db,
24        };
25        for (k, v) in self.block_hashes.iter() {
26            db.block_hashes.insert(*k, *v);
27        }
28        for (k, v) in self.contracts.iter() {
29            db.contracts.insert(*k, v.clone());
30        }
31        db.logs.clone_from(&self.logs);
32        for (address, account) in self.accounts.iter() {
33            let mut info = account.info.clone();
34            db.insert_contract(&mut info);
35
36            let entry = db.accounts.entry(*address).or_default();
37            entry.info = info;
38            for (k, v) in account.storage.iter() {
39                entry.storage.insert(*k, *v);
40            }
41        }
42        db
43    }
44
45    pub fn update_accounts(&self) -> FastCacheDB<EmptyDB> {
46        let mut db = (*self.db.as_ref()).clone();
47
48        for (k, v) in self.block_hashes.iter() {
49            db.block_hashes.insert(*k, *v);
50        }
51        for (k, v) in self.contracts.iter() {
52            db.contracts.entry(*k).and_modify(|k| k.clone_from(v));
53        }
54        db.logs.clone_from(&self.logs);
55
56        for (address, account) in self.accounts.iter() {
57            db.accounts.entry(*address).and_modify(|db_account| {
58                let info = account.info.clone();
59                db_account.info = info;
60                for (k, v) in account.storage.iter() {
61                    db_account.storage.insert(*k, *v);
62                }
63                db_account.account_state = AccountState::Touched
64            });
65        }
66        db
67    }
68
69    pub fn update_cells(&self) -> FastCacheDB<EmptyDB> {
70        let mut db = self.db.as_ref().clone();
71
72        for (k, v) in self.block_hashes.iter() {
73            db.block_hashes.insert(*k, *v);
74        }
75        for (k, v) in self.contracts.iter() {
76            db.contracts.entry(*k).and_modify(|k| k.clone_from(v));
77        }
78        db.logs.clone_from(&self.logs);
79
80        for (address, account) in self.accounts.iter() {
81            db.accounts.entry(*address).and_modify(|db_account| {
82                let info = account.info.clone();
83                db_account.info = info;
84                for (k, v) in account.storage.iter() {
85                    db_account.storage.entry(*k).and_modify(|cv| cv.clone_from(v));
86                }
87                db_account.account_state = AccountState::Touched
88            });
89        }
90        db
91    }
92
93    #[allow(irrefutable_let_patterns)]
94    pub fn apply_geth_update(&mut self, update: BTreeMap<Address, GethAccountState>) {
95        for (addr, acc_state) in update {
96            trace!("apply_geth_update {} is code {} storage_len {} ", addr, acc_state.code.is_some(), acc_state.storage.len());
97
98            for (k, v) in acc_state.storage.iter() {
99                if let Err(e) = self.insert_account_storage(addr, (*k).into(), (*v).into()) {
100                    error!("apply_geth_update :{}", e);
101                }
102            }
103
104            if let Ok(account) = self.load_account(addr) {
105                if let Some(code) = acc_state.code.clone() {
106                    let bytecode = Bytecode::new_raw(code);
107                    account.info.code_hash = bytecode.hash_slow();
108                    account.info.code = Some(bytecode);
109                }
110                if let Some(nonce) = acc_state.nonce {
111                    //trace!("nonce : {} -> {}", account.info.nonce, nonce);
112                    account.info.nonce = nonce;
113                }
114                if let Some(balance) = acc_state.balance {
115                    account.info.balance = balance;
116                }
117                account.account_state = AccountState::Touched;
118            }
119        }
120    }
121
122    pub fn apply_geth_update_vec(&mut self, update: Vec<BTreeMap<Address, GethAccountState>>) {
123        for entry in update.into_iter() {
124            self.apply_geth_update(entry);
125        }
126    }
127}