kabu_evm_db/
kabu_db.rs

1use crate::alloydb::AlloyDB;
2use crate::error::KabuDBError;
3use crate::fast_cache_db::FastDbAccount;
4use crate::fast_hasher::SimpleBuildHasher;
5use crate::kabu_db_helper::KabuDBHelper;
6use crate::DatabaseKabuExt;
7use alloy::consensus::constants::KECCAK_EMPTY;
8use alloy::eips::BlockNumberOrTag;
9use alloy::primitives::map::HashMap;
10use alloy::primitives::{Address, BlockNumber, Log, B256, U256};
11use alloy::providers::{Network, Provider, ProviderBuilder};
12use alloy::rpc::client::ClientBuilder;
13use alloy::rpc::types::trace::geth::AccountState as GethAccountState;
14use alloy::transports::Transport;
15use eyre::{OptionExt, Result};
16use revm::database::AccountState as DBAccountState;
17use revm::database::EmptyDBTyped;
18use revm::state::{Account, AccountInfo, Bytecode};
19use revm::{Database, DatabaseCommit, DatabaseRef};
20use std::collections::hash_map::Entry;
21use std::collections::BTreeMap;
22use std::fmt::{Debug, Formatter};
23use std::sync::Arc;
24use tracing::{error, trace};
25
26#[derive(Clone)]
27pub struct KabuDB
28where
29    Self: Sized + Send + Sync,
30{
31    pub accounts: HashMap<Address, FastDbAccount>,
32    pub contracts: HashMap<B256, Bytecode, SimpleBuildHasher>,
33    pub logs: Vec<Log>,
34    pub block_hashes: HashMap<BlockNumber, B256>,
35    pub read_only_db: Option<Arc<KabuDB>>,
36    pub ext_db: Option<Arc<dyn DatabaseRef<Error = KabuDBError> + Send + Sync>>,
37}
38
39impl Debug for KabuDB {
40    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("KabuDB").field("accounts", &self.accounts).finish()
42    }
43}
44
45impl Default for KabuDB {
46    fn default() -> Self {
47        KabuDB::new().with_ext_db(EmptyDBTyped::<KabuDBError>::new())
48    }
49}
50
51#[allow(dead_code)]
52impl KabuDB {
53    pub fn empty() -> Self {
54        Self::default()
55    }
56    pub fn new() -> Self {
57        let mut contracts = HashMap::with_hasher(SimpleBuildHasher::default());
58        contracts.insert(KECCAK_EMPTY, Bytecode::default());
59        contracts.insert(B256::ZERO, Bytecode::default());
60
61        Self {
62            read_only_db: None,
63            ext_db: None,
64            accounts: Default::default(),
65            contracts,
66            logs: Default::default(),
67            block_hashes: Default::default(),
68        }
69    }
70
71    pub fn is_rw_ro_account(&self, address: &Address) -> bool {
72        self.accounts.contains_key(address) || if let Some(db) = &self.read_only_db { db.accounts.contains_key(address) } else { false }
73    }
74
75    pub fn is_rw_ro_slot(&self, address: &Address, slot: &U256) -> bool {
76        let is_rw_slot = if let Some(account) = self.accounts.get(address) { account.storage.contains_key(slot) } else { false };
77
78        if is_rw_slot {
79            true
80        } else if let Some(read_only_db) = &self.read_only_db {
81            if let Some(account) = read_only_db.accounts.get(address) {
82                account.storage.contains_key(slot)
83            } else {
84                false
85            }
86        } else {
87            false
88        }
89    }
90
91    pub fn read_only_db(&self) -> Self {
92        self.read_only_db.clone().map_or(KabuDB::empty(), |a| a.as_ref().clone())
93    }
94    pub fn rw_contracts_len(&self) -> usize {
95        self.contracts.len()
96    }
97    pub fn rw_accounts_len(&self) -> usize {
98        self.accounts.len()
99    }
100
101    pub fn ro_contracts_len(&self) -> usize {
102        self.read_only_db.as_ref().map_or(0, |db| db.contracts_len())
103    }
104
105    pub fn ro_accounts_len(&self) -> usize {
106        self.read_only_db.as_ref().map_or(0, |db| db.accounts_len())
107    }
108
109    pub fn rw_storage_len(&self) -> usize {
110        self.accounts.values().map(|a| a.storage.len()).sum()
111    }
112    pub fn ro_storage_len(&self) -> usize {
113        self.read_only_db.as_ref().map_or(0, |db| db.accounts.values().map(|a| a.storage.len()).sum())
114    }
115
116    pub fn with_ext_db<ExtDB>(self, ext_db: ExtDB) -> Self
117    where
118        ExtDB: DatabaseRef<Error = KabuDBError> + Send + Sync + 'static,
119        Self: Sized,
120    {
121        let ext_db = Arc::new(ext_db) as Arc<dyn DatabaseRef<Error = KabuDBError> + Send + Sync>;
122        Self { ext_db: Some(ext_db), ..self }
123    }
124
125    pub fn without_ext_db(self) -> Self
126    where
127        Self: Sized,
128    {
129        Self { ext_db: None, ..self }
130    }
131
132    pub fn with_ro_db(self, db: Option<KabuDB>) -> Self {
133        Self { read_only_db: db.map(Arc::new), ..self }
134    }
135
136    pub fn new_with_ext_db<ExtDB>(db: KabuDB, ext_db: ExtDB) -> Self
137    where
138        ExtDB: DatabaseRef<Error = KabuDBError> + Send + Sync + 'static,
139        Self: Sized,
140    {
141        Self::new().with_ro_db(Some(db)).with_ext_db(ext_db)
142    }
143
144    // Returns the account for the given address.
145    ///
146    /// If the account was not found in the cache, it will be loaded from the underlying database.
147    pub fn load_ro_rw_ext_account(&mut self, address: Address) -> Result<&mut FastDbAccount> {
148        match self.accounts.entry(address) {
149            Entry::Occupied(entry) => Ok(entry.into_mut()),
150            Entry::Vacant(entry) => Ok(entry.insert(
151                KabuDBHelper::get_or_fetch_basic(&self.read_only_db, &self.ext_db, address)
152                    .unwrap_or_default()
153                    .map(|info| FastDbAccount { info, ..Default::default() })
154                    .unwrap_or_else(FastDbAccount::new_not_existing),
155            )),
156        }
157    }
158
159    // Returns the account for the given address.
160    ///
161    /// If the account was not found in the cache, it will be loaded from the underlying database.
162    pub fn load_ro_rw_account(&mut self, address: Address) -> Result<&mut FastDbAccount> {
163        match self.accounts.entry(address) {
164            Entry::Occupied(entry) => Ok(entry.into_mut()),
165            Entry::Vacant(entry) => Ok(entry.insert(
166                KabuDBHelper::get_basic(&self.read_only_db, address)
167                    .unwrap_or_default()
168                    .map(|info| FastDbAccount { info, ..Default::default() })
169                    .unwrap_or_else(FastDbAccount::new_not_existing),
170            )),
171        }
172    }
173
174    pub fn new_with_ro_db_and_provider<P, N>(read_only_db: Option<KabuDB>, client: P) -> Result<Self>
175    where
176        N: Network,
177        P: Provider<N> + 'static,
178        Self: Sized,
179    {
180        let box_transport = client.client().transport().clone().boxed();
181
182        let rpc_client = ClientBuilder::default().transport(box_transport, true);
183
184        let provider = ProviderBuilder::new().disable_recommended_fillers().connect_client(rpc_client);
185
186        let ext_db = AlloyDB::new(provider, BlockNumberOrTag::Latest.into());
187
188        let ext_db = ext_db.ok_or_eyre("EXT_DB_NOT_CREATED")?;
189
190        Ok(Self::new().with_ro_db(read_only_db).with_ext_db(ext_db))
191    }
192
193    pub fn insert_contract(&mut self, account: &mut AccountInfo) {
194        if let Some(code) = &account.code {
195            if !code.is_empty() {
196                if account.code_hash == KECCAK_EMPTY {
197                    account.code_hash = code.hash_slow();
198                }
199                self.contracts.entry(account.code_hash).or_insert_with(|| code.clone());
200            }
201        }
202        if account.code_hash == B256::ZERO {
203            account.code_hash = KECCAK_EMPTY;
204        }
205    }
206
207    /// Insert account info but not override storage
208    pub fn insert_account_info(&mut self, address: Address, mut info: AccountInfo) {
209        self.insert_contract(&mut info);
210        self.accounts.entry(address).or_default().info = info;
211    }
212
213    /// insert account storage without overriding account info
214    pub fn insert_account_storage(&mut self, address: Address, slot: U256, value: U256) -> Result<()> {
215        let account = self.load_account(address)?;
216        account.storage.insert(slot, value);
217        Ok(())
218    }
219
220    /// replace account storage without overriding account info
221    pub fn replace_account_storage(&mut self, address: Address, storage: HashMap<U256, U256>) -> Result<()> {
222        let account = self.load_account(address)?;
223        account.account_state = DBAccountState::StorageCleared;
224        account.storage = storage.into_iter().collect();
225        Ok(())
226    }
227
228    pub fn merge_all(self) -> KabuDB {
229        let mut read_only_db = self.read_only_db.unwrap_or_default().as_ref().clone();
230
231        for (k, v) in self.block_hashes.iter() {
232            read_only_db.block_hashes.insert(*k, *v);
233        }
234
235        for (k, v) in self.contracts.iter() {
236            read_only_db.contracts.insert(*k, v.clone());
237        }
238        read_only_db.logs.clone_from(&self.logs);
239
240        for (address, account) in self.accounts.iter() {
241            let mut info = account.info.clone();
242            read_only_db.insert_contract(&mut info);
243
244            let entry = read_only_db.accounts.entry(*address).or_default();
245            entry.info = info;
246            for (k, v) in account.storage.iter() {
247                entry.storage.insert(*k, *v);
248            }
249        }
250
251        let read_only_db = Some(Arc::new(read_only_db));
252
253        KabuDB { read_only_db, ext_db: self.ext_db, ..Default::default() }
254    }
255
256    pub fn merge_accounts(self) -> KabuDB {
257        let read_only_db = if let Some(read_only_db) = self.read_only_db {
258            let mut read_only_db_clone = (*read_only_db).clone();
259
260            for (k, v) in self.block_hashes.iter() {
261                read_only_db_clone.block_hashes.insert(*k, *v);
262            }
263            for (k, v) in self.contracts.iter() {
264                read_only_db_clone.contracts.entry(*k).and_modify(|k| k.clone_from(v));
265            }
266            read_only_db_clone.logs.clone_from(&self.logs);
267
268            for (address, account) in self.accounts.iter() {
269                read_only_db_clone.accounts.entry(*address).and_modify(|db_account| {
270                    let info = account.info.clone();
271                    db_account.info = info;
272                    for (k, v) in account.storage.iter() {
273                        db_account.storage.insert(*k, *v);
274                    }
275                    db_account.account_state = DBAccountState::Touched
276                });
277            }
278            Some(Arc::new(read_only_db_clone))
279        } else {
280            None
281        };
282
283        KabuDB { read_only_db, ext_db: self.ext_db, ..Default::default() }
284    }
285
286    pub fn merge_cells(self) -> KabuDB {
287        let read_only_db = if let Some(read_only_db) = self.read_only_db {
288            let mut read_only_db_clone = (*read_only_db).clone();
289
290            for (k, v) in self.block_hashes.iter() {
291                read_only_db_clone.block_hashes.insert(*k, *v);
292            }
293            for (k, v) in self.contracts.iter() {
294                read_only_db_clone.contracts.entry(*k).and_modify(|k| k.clone_from(v));
295            }
296            read_only_db_clone.logs.clone_from(&self.logs);
297
298            for (address, account) in self.accounts.iter() {
299                read_only_db_clone.accounts.entry(*address).and_modify(|db_account| {
300                    let info = account.info.clone();
301                    db_account.info = info;
302                    for (k, v) in account.storage.iter() {
303                        db_account.storage.entry(*k).and_modify(|cv| cv.clone_from(v));
304                    }
305                    db_account.account_state = DBAccountState::Touched
306                });
307            }
308            Some(Arc::new(read_only_db_clone))
309        } else {
310            None
311        };
312
313        KabuDB { read_only_db, ext_db: self.ext_db, ..Default::default() }
314    }
315
316    pub fn apply_geth_update(&mut self, update: BTreeMap<Address, GethAccountState>) {
317        for (addr, acc_state) in update {
318            trace!("apply_geth_update {} is code {} storage_len {} ", addr, acc_state.code.is_some(), acc_state.storage.len());
319
320            for (k, v) in acc_state.storage.iter() {
321                if let Err(e) = self.insert_account_storage(addr, (*k).into(), (*v).into()) {
322                    error!("apply_geth_update :{}", e);
323                }
324            }
325            if let Ok(account) = self.load_account(addr) {
326                if let Some(code) = acc_state.code.clone() {
327                    let bytecode = Bytecode::new_raw(code);
328                    account.info.code_hash = bytecode.hash_slow();
329                    account.info.code = Some(bytecode);
330                }
331                if let Some(nonce) = acc_state.nonce {
332                    //trace!("nonce : {} -> {}", account.info.nonce, nonce);
333                    account.info.nonce = nonce;
334                }
335                if let Some(balance) = acc_state.balance {
336                    account.info.balance = balance;
337                }
338                account.account_state = DBAccountState::Touched;
339            }
340        }
341    }
342
343    pub fn apply_geth_update_vec(&mut self, update: Vec<BTreeMap<Address, GethAccountState>>) {
344        for entry in update.into_iter() {
345            self.apply_geth_update(entry);
346        }
347    }
348
349    pub fn apply_account_info_btree(
350        &mut self,
351        address: &Address,
352        account_updated_state: &alloy::rpc::types::trace::geth::AccountState,
353        insert: bool,
354        only_new: bool,
355    ) {
356        let account = self.load_cached_account(*address);
357
358        if let Ok(account) = account {
359            if insert
360                || ((account.account_state == DBAccountState::NotExisting || account.account_state == DBAccountState::None) && only_new)
361                || (!only_new
362                    && (account.account_state == DBAccountState::Touched || account.account_state == DBAccountState::StorageCleared))
363            {
364                let code: Option<Bytecode> = match &account_updated_state.code {
365                    Some(c) => {
366                        if c.len() < 2 {
367                            account.info.code.clone()
368                        } else {
369                            Some(Bytecode::new_raw(c.clone()))
370                        }
371                    }
372                    None => account.info.code.clone(),
373                };
374
375                trace!(
376                    "apply_account_info {address}.  code len: {} storage len: {}",
377                    code.clone().map_or(0, |x| x.len()),
378                    account.storage.len()
379                );
380
381                let account_info = AccountInfo {
382                    balance: account_updated_state.balance.unwrap_or_default(),
383                    nonce: account_updated_state.nonce.unwrap_or_default(),
384                    code_hash: if code.is_some() { KECCAK_EMPTY } else { Default::default() },
385                    code,
386                };
387
388                self.insert_account_info(*address, account_info);
389            } else {
390                trace!("apply_account_info exists {address}. storage len: {}", account.storage.len(),);
391            }
392        }
393
394        if let Ok(account) = self.load_cached_account(*address) {
395            account.account_state = DBAccountState::Touched;
396            trace!(
397                "after apply_account_info account: {address} state: {:?} storage len: {} code len : {}",
398                account.account_state,
399                account.storage.len(),
400                account.info.code.clone().map_or(0, |c| c.len())
401            );
402        } else {
403            trace!(%address, "account not found after apply");
404        }
405    }
406
407    pub fn apply_account_storage(&mut self, address: &Address, acc_state: &GethAccountState, insert: bool, only_new: bool) {
408        if insert {
409            if let Ok(account) = self.load_cached_account(*address) {
410                for (slot, value) in acc_state.storage.iter() {
411                    trace!(%address, ?slot, ?value, "Inserting storage");
412                    account.storage.insert((*slot).into(), (*value).into());
413                }
414            }
415        } else if self.is_account(address) {
416            let slots_to_insert: Vec<_> = acc_state
417                .storage
418                .iter()
419                .filter_map(|(slot, value)| {
420                    let is_slot = self.is_slot(address, &(*slot).into());
421                    if is_slot || only_new {
422                        Some(((*slot).into(), (*value).into()))
423                    } else {
424                        None
425                    }
426                })
427                .collect();
428
429            if let Ok(account) = self.load_cached_account(*address) {
430                for (slot, value) in slots_to_insert {
431                    account.storage.insert(slot, value);
432                    trace!(%address, ?slot, ?value, "Inserting storage");
433                }
434            }
435        }
436    }
437
438    pub fn apply_geth_state_update(
439        &mut self,
440        update_vec: &Vec<BTreeMap<Address, GethAccountState>>,
441        insert: bool,
442        only_new: bool,
443    ) -> &mut Self {
444        for update_record in update_vec {
445            for (address, acc_state) in update_record {
446                trace!(
447                    "updating {address} insert: {insert} only_new: {only_new} storage len {} code: {}",
448                    acc_state.storage.len(),
449                    acc_state.code.is_some()
450                );
451                self.apply_account_info_btree(address, acc_state, insert, only_new);
452                self.apply_account_storage(address, acc_state, insert, only_new);
453            }
454        }
455        self
456    }
457}
458
459impl DatabaseKabuExt for KabuDB {
460    fn with_ext_db(&mut self, arc_db: impl DatabaseRef<Error = KabuDBError> + Send + Sync + 'static) {
461        self.ext_db = Some(Arc::new(arc_db))
462    }
463
464    fn is_account(&self, address: &Address) -> bool {
465        self.is_rw_ro_account(address)
466    }
467
468    fn is_slot(&self, address: &Address, slot: &U256) -> bool {
469        self.is_rw_ro_slot(address, slot)
470    }
471
472    fn contracts_len(&self) -> usize {
473        self.rw_contracts_len() + self.ro_contracts_len()
474    }
475
476    fn accounts_len(&self) -> usize {
477        self.rw_accounts_len() + self.ro_contracts_len()
478    }
479
480    fn storage_len(&self) -> usize {
481        self.rw_storage_len() + self.ro_storage_len()
482    }
483
484    fn load_account(&mut self, address: Address) -> Result<&mut FastDbAccount> {
485        self.load_ro_rw_ext_account(address)
486    }
487
488    fn load_cached_account(&mut self, address: Address) -> Result<&mut FastDbAccount> {
489        self.load_ro_rw_account(address)
490    }
491
492    fn insert_contract(&mut self, account: &mut AccountInfo) {
493        self.insert_contract(account)
494    }
495
496    fn insert_account_info(&mut self, address: Address, info: AccountInfo) {
497        self.insert_account_info(address, info)
498    }
499
500    fn insert_account_storage(&mut self, address: Address, slot: U256, value: U256) -> Result<()> {
501        self.insert_account_storage(address, slot, value)
502    }
503
504    fn replace_account_storage(&mut self, address: Address, storage: HashMap<U256, U256>) -> Result<()> {
505        self.replace_account_storage(address, storage)
506    }
507
508    fn maintain(self) -> Self {
509        self.merge_all()
510    }
511}
512
513impl DatabaseRef for KabuDB {
514    type Error = KabuDBError;
515    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
516        trace!(%address, "basic_ref");
517        let result = match address {
518            Address::ZERO => Ok(Some(AccountInfo::default())),
519            _ => match self.accounts.get(&address) {
520                Some(acc) => {
521                    trace!(%address, "account found");
522                    Ok(acc.info())
523                }
524                None => Ok(KabuDBHelper::get_or_fetch_basic(&self.read_only_db, &self.ext_db, address).unwrap_or_default()),
525            },
526        };
527
528        result
529    }
530
531    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
532        match self.contracts.get(&code_hash) {
533            Some(entry) => Ok(entry.clone()),
534            None => KabuDBHelper::get_code_by_hash(&self.read_only_db, code_hash),
535        }
536    }
537
538    fn storage_ref(&self, address: Address, slot: U256) -> Result<U256, Self::Error> {
539        trace!(%address, ?slot, "storage_ref");
540
541        match self.accounts.get(&address) {
542            Some(acc_entry) => match acc_entry.storage.get(&slot) {
543                Some(entry) => {
544                    trace!(%address, ?slot, %entry,  "storage_ref");
545                    Ok(*entry)
546                }
547                None => {
548                    if matches!(acc_entry.account_state, DBAccountState::StorageCleared | DBAccountState::NotExisting) {
549                        trace!(%address, ?slot, state=?acc_entry.account_state, "storage_ref ZERO");
550                        Ok(U256::ZERO)
551                    } else {
552                        KabuDBHelper::get_or_fetch_storage(&self.read_only_db, &self.ext_db, address, slot)
553                    }
554                }
555            },
556            None => KabuDBHelper::get_or_fetch_storage(&self.read_only_db, &self.ext_db, address, slot),
557        }
558    }
559
560    fn block_hash_ref(&self, number: BlockNumber) -> Result<B256, Self::Error> {
561        match self.block_hashes.get(&number) {
562            Some(entry) => Ok(*entry),
563            None => KabuDBHelper::get_or_fetch_block_hash(&self.read_only_db, &self.ext_db, number),
564        }
565    }
566}
567
568impl Database for KabuDB {
569    type Error = KabuDBError;
570
571    fn basic(&mut self, address: Address) -> std::result::Result<Option<AccountInfo>, Self::Error> {
572        trace!(%address, "basic");
573
574        let basic = match self.accounts.entry(address) {
575            Entry::Occupied(entry) => entry.into_mut(),
576            Entry::Vacant(entry) => entry.insert(
577                KabuDBHelper::get_or_fetch_basic(&self.read_only_db, &self.ext_db, address)
578                    .unwrap_or_default()
579                    .map(|info| FastDbAccount { info, ..Default::default() })
580                    .unwrap_or_else(FastDbAccount::new_not_existing),
581            ),
582        };
583        Ok(basic.info())
584    }
585
586    fn code_by_hash(&mut self, code_hash: B256) -> std::result::Result<Bytecode, Self::Error> {
587        match self.contracts.entry(code_hash) {
588            Entry::Occupied(entry) => Ok(entry.get().clone()),
589            Entry::Vacant(entry) => {
590                // if you return code bytes when basic fn is called this function is not needed.
591                Ok(entry.insert(KabuDBHelper::get_code_by_hash(&self.read_only_db, code_hash)?).clone())
592            }
593        }
594    }
595
596    /// Get the value in an account's storage slot.
597    ///
598    /// It is assumed that account is already loaded.
599    fn storage(&mut self, address: Address, slot: U256) -> std::result::Result<U256, Self::Error> {
600        trace!(%address, ?slot, "storage");
601
602        match self.accounts.entry(address) {
603            Entry::Occupied(mut acc_entry) => {
604                let acc_entry = acc_entry.get_mut();
605                match acc_entry.storage.entry(slot) {
606                    Entry::Occupied(entry) => Ok(*entry.get()),
607                    Entry::Vacant(entry) => {
608                        if matches!(acc_entry.account_state, DBAccountState::StorageCleared | DBAccountState::NotExisting) {
609                            Ok(U256::ZERO)
610                        } else {
611                            let slot = KabuDBHelper::get_or_fetch_storage(&self.read_only_db, &self.ext_db, address, slot)?;
612                            entry.insert(slot);
613                            Ok(slot)
614                        }
615                    }
616                }
617            }
618            Entry::Vacant(acc_entry) => {
619                let info = KabuDBHelper::get_or_fetch_basic(&self.read_only_db, &self.ext_db, address)?;
620                let (account, value) = if info.is_some() {
621                    let value = KabuDBHelper::get_or_fetch_storage(&self.read_only_db, &self.ext_db, address, slot)?;
622                    let mut account: FastDbAccount = info.into();
623                    account.storage.insert(slot, value);
624                    (account, value)
625                } else {
626                    (info.into(), U256::ZERO)
627                };
628                acc_entry.insert(account);
629                Ok(value)
630            }
631        }
632    }
633
634    fn block_hash(&mut self, number: BlockNumber) -> std::result::Result<B256, Self::Error> {
635        match self.block_hashes.entry(number) {
636            Entry::Occupied(entry) => Ok(*entry.get()),
637            Entry::Vacant(entry) => {
638                let hash = KabuDBHelper::get_or_fetch_block_hash(&self.read_only_db, &self.ext_db, number)?;
639                entry.insert(hash);
640                Ok(hash)
641            }
642        }
643    }
644}
645
646impl DatabaseCommit for KabuDB {
647    fn commit(&mut self, changes: HashMap<Address, Account>) {
648        for (address, mut account) in changes {
649            if !account.is_touched() {
650                continue;
651            }
652            if account.is_selfdestructed() {
653                let db_account = self.accounts.entry(address).or_default();
654                db_account.storage.clear();
655                db_account.account_state = DBAccountState::NotExisting;
656                db_account.info = AccountInfo::default();
657                continue;
658            }
659            let is_newly_created = account.is_created();
660            self.insert_contract(&mut account.info);
661
662            let db_account = self.accounts.entry(address).or_default();
663            db_account.info = account.info;
664
665            db_account.account_state = if is_newly_created {
666                db_account.storage.clear();
667                DBAccountState::StorageCleared
668            } else if db_account.account_state.is_storage_cleared() {
669                // Preserve old account state if it already exists
670                DBAccountState::StorageCleared
671            } else {
672                DBAccountState::Touched
673            };
674            db_account.storage.extend(account.storage.into_iter().map(|(key, value)| (key, value.present_value())));
675        }
676    }
677}
678
679#[cfg(test)]
680mod test {
681    use super::GethAccountState;
682    use crate::alloydb::AlloyDB;
683    use crate::error::KabuDBError;
684    use crate::kabu_db::KabuDB;
685    use alloy::eips::BlockNumberOrTag;
686    use alloy::primitives::map::HashMap;
687    use alloy::primitives::{Address, Bytes, B256, I256, U256};
688    use alloy::providers::{Provider, ProviderBuilder};
689    use eyre::ErrReport;
690    use revm::database::EmptyDBTyped;
691    use revm::primitives::KECCAK_EMPTY;
692    use revm::state::{AccountInfo, Bytecode};
693    use revm::{Database, DatabaseRef};
694    use std::collections::BTreeMap;
695
696    #[test]
697    fn test_new_with_provider() {
698        let db = KabuDB::new();
699        let provider = ProviderBuilder::new().connect_anvil_with_wallet();
700
701        let rt = tokio::runtime::Runtime::new().unwrap();
702
703        rt.block_on(async move {
704            let test_addr = Address::parse_checksummed("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266", None).unwrap();
705
706            let balance = provider.get_balance(test_addr).await?;
707
708            let db = KabuDB::new_with_ro_db_and_provider(Some(db), provider.clone()).unwrap();
709
710            let info = db.basic_ref(test_addr).unwrap().unwrap();
711
712            assert_eq!(info.balance, U256::from(10000000000000000000000u128));
713            assert_eq!(info.balance, balance);
714            Ok::<(), ErrReport>(())
715        })
716        .unwrap();
717    }
718
719    #[test]
720    fn test_new_with_ext_db() {
721        let db = KabuDB::new();
722        let provider = ProviderBuilder::new().connect_anvil_with_wallet();
723
724        let rt = tokio::runtime::Runtime::new().unwrap();
725
726        rt.block_on(async move {
727            let test_addr = Address::parse_checksummed("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266", None).unwrap();
728
729            let balance = provider.get_balance(test_addr).await?;
730
731            let ext_db = AlloyDB::new(provider.clone(), BlockNumberOrTag::Latest.into()).unwrap();
732
733            let db = KabuDB::new_with_ext_db(db, ext_db);
734
735            let info = db.basic_ref(test_addr).unwrap().unwrap();
736            assert_eq!(info.balance, U256::from(10000000000000000000000u128));
737            assert_eq!(info.balance, balance);
738
739            Ok::<(), ErrReport>(())
740        })
741        .unwrap();
742    }
743
744    #[test]
745    fn test_insert_account_storage() {
746        let account = Address::with_last_byte(42);
747        let nonce = 42;
748        let mut init_state = KabuDB::new();
749        init_state.insert_account_info(account, AccountInfo { nonce, ..Default::default() });
750
751        let (key, value) = (U256::from(123), U256::from(456));
752        let mut new_state = KabuDB::new().with_ro_db(Some(init_state));
753
754        new_state.insert_account_storage(account, key, value).unwrap();
755
756        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce);
757        assert_eq!(new_state.storage(account, key).unwrap(), value);
758    }
759
760    #[test]
761    fn test_insert_account_storage_inherited() {
762        let account = Address::with_last_byte(42);
763        let nonce = 42;
764        let mut init_state = KabuDB::new();
765        init_state.insert_account_info(account, AccountInfo { nonce, ..Default::default() });
766
767        let (key, value) = (U256::from(123), U256::from(456));
768        let mut new_state = KabuDB::new().with_ro_db(Some(init_state));
769        new_state.insert_account_storage(account, key, value).unwrap();
770
771        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce);
772        assert_eq!(new_state.storage(account, key).unwrap(), value);
773    }
774
775    #[test]
776    fn test_replace_account_storage() {
777        let account = Address::with_last_byte(42);
778        let nonce = 42;
779        let mut init_state = KabuDB::new();
780        init_state.insert_account_info(account, AccountInfo { nonce, ..Default::default() });
781
782        let (key0, value0) = (U256::from(123), U256::from(456));
783        let (key1, value1) = (U256::from(789), U256::from(999));
784        init_state.insert_account_storage(account, key0, value0).unwrap();
785
786        let mut new_state = KabuDB::new().with_ro_db(Some(init_state));
787        assert_eq!(new_state.accounts.len(), 0);
788        let mut hm: HashMap<U256, U256> = Default::default();
789        hm.insert(key1, value1);
790
791        new_state.replace_account_storage(account, hm).unwrap();
792
793        let mut new_state = new_state.merge_all();
794
795        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce);
796        assert_eq!(new_state.storage(account, key0).unwrap(), value0);
797        assert_eq!(new_state.storage(account, key1).unwrap(), value1);
798        assert_eq!(new_state.accounts.len(), 1);
799    }
800
801    #[test]
802    fn test_apply_geth_update() {
803        let account = Address::with_last_byte(42);
804        let nonce = 42;
805        let code = Bytecode::new_raw(Bytes::from(vec![1, 2, 3]));
806        let mut init_state = KabuDB::new();
807        init_state.insert_account_info(account, AccountInfo { nonce, code: Some(code.clone()), ..Default::default() });
808
809        let (key0, value0) = (U256::from(123), U256::from(456));
810        let (key1, value1) = (U256::from(789), U256::from(999));
811        init_state.insert_account_storage(account, key0, value0).unwrap();
812        init_state.insert_account_storage(account, key1, value1).unwrap();
813
814        let mut new_state = KabuDB::new().with_ro_db(Some(init_state));
815        assert_eq!(new_state.accounts.len(), 0);
816
817        let update_record = GethAccountState {
818            balance: None,
819            code: Some(Bytes::from(vec![1, 2, 3])),
820            nonce: Some(nonce + 1),
821            storage: [(B256::from(I256::try_from(123).unwrap()), B256::from(I256::try_from(333).unwrap()))].into(),
822        };
823
824        let update: BTreeMap<Address, GethAccountState> = [(account, update_record)].into();
825
826        new_state.apply_geth_update(update);
827
828        assert_eq!(new_state.basic(account).unwrap().unwrap().code, Some(code.clone()));
829        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce + 1);
830        assert_eq!(new_state.storage_ref(account, key0).unwrap(), U256::from(333));
831        assert_eq!(new_state.storage_ref(account, key1).unwrap(), value1);
832        assert_eq!(new_state.accounts.len(), 1);
833
834        let mut new_state = new_state.merge_all();
835
836        assert_eq!(new_state.basic(account).unwrap().unwrap().code, Some(code.clone()));
837        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce + 1);
838        assert_eq!(new_state.storage_ref(account, key0).unwrap(), U256::from(333));
839        assert_eq!(new_state.storage_ref(account, key1).unwrap(), value1);
840        assert_eq!(new_state.accounts.len(), 1);
841    }
842
843    #[test]
844    fn test_merge() {
845        let account = Address::with_last_byte(42);
846        let nonce = 42;
847        let code = Bytecode::new_raw(Bytes::from(vec![1, 2, 3]));
848        let mut init_state = KabuDB::new();
849        init_state.insert_account_info(account, AccountInfo { nonce, code: Some(code.clone()), ..Default::default() });
850
851        let (key0, value0) = (U256::from(123), U256::from(456));
852        let (key1, value1) = (U256::from(789), U256::from(999));
853        let (key2, value2) = (U256::from(999), U256::from(111));
854        init_state.insert_account_storage(account, key0, value0).unwrap();
855        init_state.insert_account_storage(account, key1, value1).unwrap();
856
857        let mut new_state = KabuDB::new().with_ro_db(Some(init_state));
858        assert_eq!(new_state.accounts.len(), 0);
859
860        new_state.insert_account_info(
861            account,
862            AccountInfo {
863                balance: U256::ZERO,
864                code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2, 2]))),
865                nonce: nonce + 1,
866                code_hash: KECCAK_EMPTY,
867            },
868        );
869
870        new_state.insert_account_storage(account, key0, U256::from(333)).unwrap();
871        new_state.insert_account_storage(account, key2, value2).unwrap();
872
873        let mut new_state = new_state.merge_all();
874
875        assert_eq!(new_state.basic(account).unwrap().unwrap().code, Some(Bytecode::new_raw(Bytes::from(vec![1, 2, 2]))));
876        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce + 1);
877        assert_eq!(new_state.storage_ref(account, key0).unwrap(), U256::from(333));
878        assert_eq!(new_state.storage_ref(account, key1).unwrap(), value1);
879        assert_eq!(new_state.storage_ref(account, key2).unwrap(), value2);
880        assert_eq!(new_state.accounts.len(), 1);
881    }
882
883    #[test]
884    fn test_update_cell() {
885        let account = Address::with_last_byte(42);
886        let account2 = Address::with_last_byte(43);
887        let nonce = 42;
888        let code = Bytecode::new_raw(Bytes::from(vec![1, 2, 3]));
889        let mut init_state = KabuDB::new();
890        init_state.insert_account_info(account, AccountInfo { nonce, code: Some(code.clone()), ..Default::default() });
891
892        let (key0, value0) = (U256::from(123), U256::from(456));
893        let (key1, value1) = (U256::from(789), U256::from(999));
894        let (key2, value2) = (U256::from(999), U256::from(111));
895        init_state.insert_account_storage(account, key0, value0).unwrap();
896        init_state.insert_account_storage(account, key1, value1).unwrap();
897
898        let mut new_state = KabuDB::new().with_ro_db(Some(init_state)).with_ext_db(EmptyDBTyped::<KabuDBError>::new());
899        assert_eq!(new_state.accounts.len(), 0);
900
901        new_state.insert_account_info(
902            account,
903            AccountInfo {
904                balance: U256::ZERO,
905                code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2, 2]))),
906                nonce: nonce + 1,
907                code_hash: KECCAK_EMPTY,
908            },
909        );
910
911        new_state.insert_account_info(
912            account2,
913            AccountInfo {
914                balance: U256::ZERO,
915                code: Some(Bytecode::new_raw(Bytes::from(vec![1, 2, 2]))),
916                nonce: nonce + 1,
917                code_hash: KECCAK_EMPTY,
918            },
919        );
920
921        new_state.insert_account_storage(account, key0, U256::from(333)).unwrap();
922        new_state.insert_account_storage(account, key2, value2).unwrap();
923
924        let mut new_state = new_state.merge_cells();
925
926        assert_eq!(new_state.basic(account).unwrap().unwrap().code, Some(Bytecode::new_raw(Bytes::from(vec![1, 2, 2]))));
927        assert_eq!(new_state.basic(account).unwrap().unwrap().nonce, nonce + 1);
928        assert_eq!(new_state.storage_ref(account, key0).unwrap(), U256::from(333));
929        assert_eq!(new_state.storage_ref(account, key1).unwrap(), value1);
930        assert_eq!(new_state.storage_ref(account, key2).unwrap(), U256::ZERO);
931        assert_eq!(new_state.storage_ref(account2, key0).unwrap(), U256::ZERO);
932        assert_eq!(new_state.accounts.len(), 1);
933        assert_eq!(new_state.basic(account2).unwrap(), None);
934        assert_eq!(new_state.accounts.len(), 2);
935        assert_eq!(new_state.basic(account2).unwrap(), None);
936        assert_eq!(new_state.accounts.len(), 2);
937    }
938}
kabu_evm_db/kabu_db.rs

kabu_evm_db/
kabu_db.rs
