perf: cache repeated storage trie lookups

crates/stateless-executor/src/io.rs

@@ -1,5 +1,5 @@
 use alloc::{format, string::ToString, vec::Vec};
-use core::iter::once;
+use core::{cell::Cell, iter::once};
 
 use crate::error::StatelessExecutorError;
 use alloy_consensus::Header;
@@ -275,7 +275,7 @@ pub trait WitnessInput<'a> {
             block_hashes.insert(parent_header.number, child_header.parent_hash);
         }
 
-        Ok(WitnessDb { inner: state, block_hashes, bytecode_by_hash })
+        Ok(WitnessDb::new(state, block_hashes, bytecode_by_hash))
     }
 }
 
@@ -287,6 +287,12 @@ pub struct WitnessDb<'a, 'b> {
     inner: &'b EthereumState<'a>,
     block_hashes: HashMap<u64, B256>,
     bytecode_by_hash: HashMap<B256, &'b Bytecode>,
+    /// Most recently hashed address. Account and storage lookups tend to arrive in runs for the
+    /// same account, so a single entry eliminates repeated address keccaks.
+    hashed_address_cache: Cell<Option<(Address, B256)>>,
+    /// Most recently accessed storage trie. Distinct slot reads for one account can reuse both the
+    /// address hash and the storage-tries hash table lookup.
+    storage_trie_cache: Cell<Option<(Address, &'b Mpt<'a>)>>,
 }
 
 impl<'a, 'b> WitnessDb<'a, 'b> {
@@ -295,7 +301,40 @@ impl<'a, 'b> WitnessDb<'a, 'b> {
         block_hashes: HashMap<u64, B256>,
         bytecode_by_hash: HashMap<B256, &'b Bytecode>,
     ) -> Self {
-        Self { inner, block_hashes, bytecode_by_hash }
+        Self {
+            inner,
+            block_hashes,
+            bytecode_by_hash,
+            hashed_address_cache: Cell::new(None),
+            storage_trie_cache: Cell::new(None),
+        }
+    }
+
+    #[inline]
+    fn hashed_address(&self, address: Address) -> B256 {
+        if let Some((cached_address, hashed_address)) = self.hashed_address_cache.get() {
+            if cached_address == address {
+                return hashed_address;
+            }
+        }
+
+        let hashed_address = keccak256(address);
+        self.hashed_address_cache.set(Some((address, hashed_address)));
+        hashed_address
+    }
+
+    #[inline]
+    fn storage_trie(&self, address: Address) -> Option<&'b Mpt<'a>> {
+        if let Some((cached_address, storage_trie)) = self.storage_trie_cache.get() {
+            if cached_address == address {
+                return Some(storage_trie);
+            }
+        }
+
+        let inner: &'b EthereumState<'a> = self.inner;
+        let storage_trie = inner.storage_tries.get(&self.hashed_address(address))?;
+        self.storage_trie_cache.set(Some((address, storage_trie)));
+        Some(storage_trie)
     }
 }
 
@@ -312,7 +351,7 @@ impl DatabaseRef for WitnessDb<'_, '_> {
 
     /// Get basic account information.
     fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
-        let hashed_address = keccak256(address);
+        let hashed_address = self.hashed_address(address);
 
         let account_in_trie = self
             .inner
@@ -343,9 +382,7 @@ impl DatabaseRef for WitnessDb<'_, '_> {
     ///
     /// Returns `U256::ZERO` if the slot is not found in the trie.
     fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
-        let hashed_address = keccak256(address);
-
-        let storage_trie = self.inner.storage_tries.get(&hashed_address).ok_or_else(|| {
+        let storage_trie = self.storage_trie(address).ok_or_else(|| {
             ProviderError::TrieWitnessError(format!("storage trie for {address} not found"))
         })?;