kabu_evm_db/
kabu_db.rs

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

kabu_evm_db/
kabu_db.rs
