perf(levm): refactor CacheDB to use more efficient APIs

crates/vm/levm/src/call_frame.rs

@@ -186,6 +186,22 @@ pub struct CallFrameBackup {
 }
 
 impl CallFrameBackup {
+    pub fn backup_account_info(
+        &mut self,
+        address: Address,
+        account: &Account,
+    ) -> Result<(), InternalError> {
+        self.original_accounts_info
+            .entry(address)
+            .or_insert_with(|| Account {
+                info: account.info.clone(),
+                code: account.code.clone(),
+                storage: HashMap::new(),
+            });
+
+        Ok(())
+    }
+
     pub fn clear(&mut self) {
         self.original_accounts_info.clear();
         self.original_account_storage_slots.clear();

crates/vm/levm/src/db/gen_db.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::collections::HashSet;
 use std::sync::Arc;
 
 use bytes::Bytes;
@@ -17,7 +16,8 @@ use crate::vm::VM;
 
 use super::CacheDB;
 use super::Database;
-use super::cache;
+use std::collections::HashSet;
+use std::collections::hash_map::Entry;
 
 #[derive(Clone)]
 pub struct GeneralizedDatabase {
@@ -46,11 +46,13 @@ impl GeneralizedDatabase {
     /// Gets account, first checking the cache and then the database
     /// (caching in the second case)
     pub fn get_account(&mut self, address: Address) -> Result<&Account, InternalError> {
-        if !cache::account_is_cached(&self.current_accounts_state, &address) {
+        if !self.current_accounts_state.contains_key(&address) {
             let account = self.get_account_from_database(address)?;
-            cache::insert_account(&mut self.current_accounts_state, address, account);
+            self.current_accounts_state.insert(address, account);
         }
-        cache::get_account(&self.current_accounts_state, &address)
+
+        self.current_accounts_state
+            .get(&address)
             .ok_or(InternalError::AccountNotFound)
     }
 
@@ -140,17 +142,25 @@ impl<'a> VM<'a> {
 
     */
     pub fn get_account_mut(&mut self, address: Address) -> Result<&mut Account, InternalError> {
-        if cache::is_account_cached(&self.db.current_accounts_state, &address) {
-            self.backup_account_info(address)?;
-            cache::get_account_mut(&mut self.db.current_accounts_state, &address)
-                .ok_or(InternalError::AccountNotFound)
-        } else {
-            let acc = self.db.get_account_from_database(address)?;
-            cache::insert_account(&mut self.db.current_accounts_state, address, acc);
-            self.backup_account_info(address)?;
-            cache::get_account_mut(&mut self.db.current_accounts_state, &address)
-                .ok_or(InternalError::AccountNotFound)
-        }
+        let account = match self.db.current_accounts_state.entry(address) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => {
+                let account = self.db.store.get_account(address)?;
+                self.db
+                    .initial_accounts_state
+                    .insert(address, account.clone());
+
+                entry.insert(account)
+            }
+        };
+
+        self.call_frames
+            .last_mut()
+            .ok_or(InternalError::CallFrame)?
+            .call_frame_backup
+            .backup_account_info(address, account)?;
+
+        Ok(account)
     }
 
     pub fn increase_account_balance(
@@ -220,9 +230,13 @@ impl<'a> VM<'a> {
         address: Address,
         account: Account,
     ) -> Result<(), InternalError> {
-        self.backup_account_info(address)?;
-        let _ = cache::insert_account(&mut self.db.current_accounts_state, address, account);
+        self.call_frames
+            .last_mut()
+            .ok_or(InternalError::CallFrame)?
+            .call_frame_backup
+            .backup_account_info(address, &account)?;
 
+        self.db.current_accounts_state.insert(address, account);
         Ok(())
     }
 
@@ -270,7 +284,7 @@ impl<'a> VM<'a> {
         address: Address,
         key: H256,
     ) -> Result<U256, InternalError> {
-        if let Some(account) = cache::get_account(&self.db.current_accounts_state, &address) {
+        if let Some(account) = self.db.current_accounts_state.get(&address) {
             if let Some(value) = account.storage.get(&key) {
                 return Ok(*value);
             }
@@ -320,37 +334,4 @@ impl<'a> VM<'a> {
 
         Ok(())
     }
-
-    pub fn backup_account_info(&mut self, address: Address) -> Result<(), InternalError> {
-        if self.call_frames.is_empty() {
-            return Ok(());
-        }
-
-        let is_not_backed_up = !self
-            .current_call_frame_mut()?
-            .call_frame_backup
-            .original_accounts_info
-            .contains_key(&address);
-
-        if is_not_backed_up {
-            let account = cache::get_account(&self.db.current_accounts_state, &address)
-                .ok_or(InternalError::AccountNotFound)?;
-            let info = account.info.clone();
-            let code = account.code.clone();
-
-            self.current_call_frame_mut()?
-                .call_frame_backup
-                .original_accounts_info
-                .insert(
-                    address,
-                    Account {
-                        info,
-                        code,
-                        storage: HashMap::new(),
-                    },
-                );
-        }
-
-        Ok(())
-    }
 }

crates/vm/levm/src/db/mod.rs

@@ -1,14 +1,15 @@
 use crate::errors::DatabaseError;
 use bytes::Bytes;
-pub use cache::CacheDB;
 use ethrex_common::{
     Address, H256, U256,
     types::{Account, ChainConfig},
 };
+use std::collections::HashMap;
 
-pub mod cache;
 pub mod gen_db;
 
+pub type CacheDB = HashMap<Address, Account>;
+
 pub trait Database: Send + Sync {
     fn get_account(&self, address: Address) -> Result<Account, DatabaseError>;
     fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError>;

crates/vm/levm/src/utils.rs

@@ -2,7 +2,7 @@ use crate::{
     EVMConfig,
     call_frame::CallFrameBackup,
     constants::*,
-    db::{cache, gen_db::GeneralizedDatabase},
+    db::gen_db::GeneralizedDatabase,
     errors::{ExceptionalHalt, InternalError, TxValidationError, VMError},
     gas_cost::{
         self, ACCESS_LIST_ADDRESS_COST, ACCESS_LIST_STORAGE_KEY_COST, BLOB_GAS_PER_BLOB,
@@ -124,9 +124,7 @@ pub fn restore_cache_state(
     callframe_backup: CallFrameBackup,
 ) -> Result<(), VMError> {
     for (address, account) in callframe_backup.original_accounts_info {
-        if let Some(current_account) =
-            cache::get_account_mut(&mut db.current_accounts_state, &address)
-        {
+        if let Some(current_account) = db.current_accounts_state.get_mut(&address) {
             current_account.info = account.info;
             current_account.code = account.code;
         }
@@ -136,7 +134,9 @@ pub fn restore_cache_state(
         // This call to `get_account_mut` should never return None, because we are looking up accounts
         // that had their storage modified, which means they should be in the cache. That's why
         // we return an internal error in case we haven't found it.
-        let account = cache::get_account_mut(&mut db.current_accounts_state, &address)
+        let account = db
+            .current_accounts_state
+            .get_mut(&address)
             .ok_or(InternalError::AccountNotFound)?;
 
         for (key, value) in storage {