use db for mapping

crates/chain-state/Cargo.toml

@@ -21,6 +21,8 @@ reth-primitives.workspace = true
 reth-storage-api.workspace = true
 reth-trie.workspace = true
 reth-revm.workspace = true
+reth-db.workspace = true
+reth-db-api.workspace = true
 
 # alloy
 alloy-primitives.workspace = true

crates/chain-state/src/cache/mod.rs

@@ -1,24 +1,34 @@
 /// State provider with cached states for execution.
 pub mod cached_provider;
-mod plain_state;
+pub mod plain_state;
 
-use crate::ExecutedBlock;
-use tracing::debug;
-
-use crate::cache::plain_state::{
-    clear_plain_state, get_account, get_code, get_storage, insert_account, insert_code,
-    insert_storage, write_plain_state,
+use crate::{
+    cache::plain_state::{
+        clear_plain_state, get_account, get_code, get_storage, insert_account, insert_code,
+        insert_storage, PlainCacheWriter,
+    },
+    ExecutedBlock,
+};
+use reth_db::{
+    cursor::{DbCursorRW, DbDupCursorRO, DbDupCursorRW},
+    tables,
 };
+use tracing::debug;
 
 /// Writes the execution outcomes of the given blocks to the cache.
-pub fn write_to_cache(blocks: Vec<ExecutedBlock>) {
-    for block in blocks {
-        debug!("Start to write block {} to cache", block.block.header.number);
-        let bundle_state = block.execution_outcome().clone().bundle;
-        write_plain_state(bundle_state);
-        debug!("Finish to write block {} to cache", block.block.header.number);
-    }
-}
+// pub fn write_to_cache(
+//     blocks: Vec<ExecutedBlock>,
+//     storage_cursor: &dyn DbDupCursorRO<tables::PlainStorageState>,
+// ) {
+//     for block in blocks {
+//         debug!("Start to write block {} to cache", block.block.header.number);
+//         let bundle_state = block.execution_outcome().clone().bundle;
+//
+//         let mut cache_writer = PlainCacheWriter::new(storage_cursor);
+//         cache_writer.write_plain_state(bundle_state);
+//         debug!("Finish to write block {} to cache", block.block.header.number);
+//     }
+// }
 
 /// Clears all cached states.
 pub fn clear_cache() {

crates/chain-state/src/cache/plain_state.rs

@@ -1,13 +1,10 @@
+use crate::ExecutedBlock;
 use lazy_static::lazy_static;
 use quick_cache::sync::Cache;
+use reth_db::{cursor::DbDupCursorRO, table::TableRow, tables, DatabaseError, PlainStorageState};
+use reth_db_api::transaction::DbTx;
 use reth_primitives::{Account, Address, Bytecode, StorageKey, StorageValue, B256, U256};
 use reth_revm::db::{BundleState, OriginalValuesKnown};
-use schnellru::{ByLength, LruMap};
-use std::{
-    collections::{HashMap, HashSet},
-    sync::{Arc, Mutex, RwLock},
-};
-use tracing::info;
 
 // Cache sizes
 const ACCOUNT_CACHE_SIZE: usize = 1000000;
@@ -21,11 +18,12 @@ lazy_static! {
     /// Account cache
     pub(crate) static ref PLAIN_ACCOUNTS: Cache<Address, Account> = Cache::new(ACCOUNT_CACHE_SIZE);
 
+    /// Storage cache
+     pub(crate) static ref PLAIN_STORAGES: Cache<AddressStorageKey, StorageValue> = Cache::new(STORAGE_CACHE_SIZE);
+
     /// Contract cache
     /// The size of contract is large and the hot contracts should be limited.
      pub(crate) static ref CONTRACT_CODES: Cache<B256, Bytecode> = Cache::new(CONTRACT_CACHE_SIZE);
-
-     pub(crate) static ref STORAGES: Mutex<LruMap<Address, LruMap<B256, U256>>> = Mutex::new(LruMap::new(ByLength::new(ACCOUNT_CACHE_SIZE as u32)));
 }
 
 pub(crate) fn insert_account(k: Address, v: Account) {
@@ -34,43 +32,7 @@ pub(crate) fn insert_account(k: Address, v: Account) {
 
 /// Insert storage into the cache
 pub(crate) fn insert_storage(k: AddressStorageKey, v: U256) {
-    let mut outer = STORAGES.lock().unwrap();
-    match outer.get(&k.0) {
-        Some(inner) => {
-            inner.insert(k.1, v);
-        }
-        None => {
-            let mut inner = LruMap::new(ByLength::new(100));
-            inner.insert(k.1, v);
-            outer.insert(k.0, inner);
-        }
-    }
-}
-
-pub(crate) fn remove_storages(address: Address) {
-    let mut outer = STORAGES.lock().unwrap();
-    outer.remove(&address);
-}
-pub(crate) fn insert_storages(address: Address, storages: Vec<(U256, U256)>) {
-    if storages.is_empty() {
-        return;
-    }
-
-    let mut outer = STORAGES.lock().unwrap();
-    match outer.get(&address) {
-        Some(inner) => {
-            for (k, v) in storages {
-                inner.insert(StorageKey::from(k), v);
-            }
-        }
-        None => {
-            let mut inner = LruMap::new(ByLength::new(100));
-            for (k, v) in storages {
-                inner.insert(StorageKey::from(k), v);
-            }
-            outer.insert(address, inner);
-        }
-    }
+    PLAIN_STORAGES.insert(k, v);
 }
 
 // Get account from cache
@@ -80,14 +42,7 @@ pub(crate) fn get_account(k: &Address) -> Option<Account> {
 
 // Get storage from cache
 pub(crate) fn get_storage(k: &AddressStorageKey) -> Option<StorageValue> {
-    let mut outer = STORAGES.lock().unwrap();
-    match outer.get(&k.0) {
-        Some(inner) => match inner.get(&k.1) {
-            Some(value) => Some(*value),
-            None => None,
-        },
-        None => None,
-    }
+    PLAIN_STORAGES.get(k)
 }
 
 // Get code from cache
@@ -100,43 +55,79 @@ pub(crate) fn insert_code(k: B256, v: Bytecode) {
     CONTRACT_CODES.insert(k, v);
 }
 
-/// Write committed state to cache.
-pub(crate) fn write_plain_state(bundle: BundleState) {
-    let change_set = bundle.into_plain_state(OriginalValuesKnown::Yes);
+pub struct PlainCacheWriter<'a, TX>(&'a TX);
 
-    // Update account cache
-    for (address, account_info) in &change_set.accounts {
-        match account_info {
-            None => {
-                PLAIN_ACCOUNTS.remove(address);
-            }
-            Some(acc) => {
-                PLAIN_ACCOUNTS.insert(
-                    *address,
-                    Account {
-                        nonce: acc.nonce,
-                        balance: acc.balance,
-                        bytecode_hash: Some(acc.code_hash),
-                    },
-                );
-            }
-        }
+impl<'a, TX> PlainCacheWriter<'a, TX> {
+    pub const fn new(tx: &'a TX) -> Self {
+        Self(tx)
     }
+}
 
-    // Update storage cache
-    for storage in &change_set.storage {
-        if storage.wipe_storage {
-            info!(target: "blockchain_tree", "wipe_storage is true.");
-            remove_storages(storage.address);
+impl<'a, TX> PlainCacheWriter<'a, TX> {
+    /// Write committed state to cache.
+    pub fn write_plain_state(&mut self, blocks: Vec<ExecutedBlock>)
+    where
+        TX: DbTx,
+    {
+        let cursor = self.0.cursor_dup_read::<tables::PlainStorageState>();
+        match cursor {
+            Ok(mut cursor) => {
+                for block in blocks {
+                    let bundle_state = block.execution_outcome().clone().bundle;
+                    let change_set = bundle_state.into_plain_state(OriginalValuesKnown::Yes);
+
+                    // Update account cache
+                    for (address, account_info) in &change_set.accounts {
+                        match account_info {
+                            None => {
+                                PLAIN_ACCOUNTS.remove(address);
+                            }
+                            Some(acc) => {
+                                PLAIN_ACCOUNTS.insert(
+                                    *address,
+                                    Account {
+                                        nonce: acc.nonce,
+                                        balance: acc.balance,
+                                        bytecode_hash: Some(acc.code_hash),
+                                    },
+                                );
+                            }
+                        }
+                    }
+
+                    // Update storage cache
+                    for storage in &change_set.storage {
+                        if storage.wipe_storage {
+                            let walker = cursor.walk_dup(Some(storage.address), None).unwrap();
+                            for kv in walker {
+                                match kv {
+                                    Ok((k, v)) => {
+                                        PLAIN_STORAGES.remove(&(k, v.key));
+                                    }
+                                    Err(_) => {
+                                        PLAIN_STORAGES.clear();
+                                        break;
+                                    }
+                                }
+                            }
+                        }
+
+                        for (k, v) in storage.storage.clone() {
+                            insert_storage((storage.address, StorageKey::from(k)), v);
+                        }
+                    }
+                }
+            }
+            Err(_) => {
+                PLAIN_ACCOUNTS.clear();
+                PLAIN_STORAGES.clear();
+            }
         }
-
-        insert_storages(storage.address, storage.storage.clone());
     }
 }
 
 /// Clear cached accounts and storages.
 pub(crate) fn clear_plain_state() {
     PLAIN_ACCOUNTS.clear();
-    let mut outer = STORAGES.lock().unwrap();
-    outer.clear();
+    PLAIN_STORAGES.clear();
 }

crates/chain-state/src/cache/plain_state.rsold

@@ -31,7 +31,6 @@ lazy_static! {
     /// Mapping for deleting storages
     static ref PLAIN_STORAGES_MAPPING: Mutex<HashMap<Address, HashSet<B256>>> = Mutex::new(HashMap::new());
 
-     pub(crate) static ref STORAGES: Arc<RwLock<LruMap<Address,HashMap<B256, U256>>>> = Arc::new(RwLock::new(LruMap::new(ByLength::new(ACCOUNT_CACHE_SIZE as u32))));
 }
 
 pub(crate) fn insert_account(k: Address, v: Account) {