fix: use different cache instance per precompile

fgimenez · fgimenez · commit 75081c83db8b · 2025-05-13T13:37:51.000+02:00
diff --git a/crates/engine/tree/benches/state_root_task.rs b/crates/engine/tree/benches/state_root_task.rs
@@ -13,7 +13,7 @@ use reth_chain_state::EthPrimitives;
 use reth_chainspec::ChainSpec;
 use reth_db_common::init::init_genesis;
 use reth_engine_tree::tree::{
-    executor::WorkloadExecutor, precompile_cache::PrecompileCache, PayloadProcessor,
+    executor::WorkloadExecutor, precompile_cache::PrecompileCacheMap, PayloadProcessor,
     StateProviderBuilder, TreeConfig,
 };
 use reth_evm::OnStateHook;
@@ -221,7 +221,7 @@ fn bench_state_root(c: &mut Criterion) {
                             WorkloadExecutor::default(),
                             EthEvmConfig::new(factory.chain_spec()),
                             &TreeConfig::default(),
-                            PrecompileCache::default(),
+                            PrecompileCacheMap::default(),
                         );
                         let provider = BlockchainProvider::new(factory).unwrap();
 
diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
@@ -18,7 +18,7 @@ use error::{InsertBlockError, InsertBlockErrorKind, InsertBlockFatalError};
 use instrumented_state::InstrumentedStateProvider;
 use payload_processor::sparse_trie::StateRootComputeOutcome;
 use persistence_state::CurrentPersistenceAction;
-use precompile_cache::{CachedPrecompile, PrecompileCache};
+use precompile_cache::{CachedPrecompile, PrecompileCacheMap};
 use reth_chain_state::{
     CanonicalInMemoryState, ExecutedBlock, ExecutedBlockWithTrieUpdates,
     MemoryOverlayStateProvider, NewCanonicalChain,
@@ -268,8 +268,8 @@ where
     payload_processor: PayloadProcessor<N, C>,
     /// The EVM configuration.
     evm_config: C,
-    /// Precompile cache.
-    precompile_cache: PrecompileCache,
+    /// Precompile cache map.
+    precompile_cache_map: PrecompileCacheMap,
 }
 
 impl<N, P: Debug, T: PayloadTypes + Debug, V: Debug, C: Debug> std::fmt::Debug
@@ -334,13 +334,13 @@ where
     ) -> Self {
         let (incoming_tx, incoming) = std::sync::mpsc::channel();
 
-        let precompile_cache = PrecompileCache::default();
+        let precompile_cache_map = PrecompileCacheMap::default();
 
         let payload_processor = PayloadProcessor::new(
             WorkloadExecutor::default(),
             evm_config.clone(),
             &config,
-            precompile_cache.clone(),
+            precompile_cache_map.clone(),
         );
 
         Self {
@@ -362,7 +362,7 @@ where
             engine_kind,
             payload_processor,
             evm_config,
-            precompile_cache,
+            precompile_cache_map,
         }
     }
 
@@ -2275,7 +2275,7 @@ where
 
     /// Executes a block with the given state provider
     fn execute_block<S: StateProvider>(
-        &self,
+        &mut self,
         state_provider: S,
         block: &RecoveredBlock<N::Block>,
         handle: &PayloadHandle,
@@ -2289,8 +2289,11 @@ where
         let mut executor = self.evm_config.executor_for_block(&mut db, block);
 
         if self.config.precompile_cache_enabled() {
-            executor.evm_mut().precompiles_mut().map_precompiles(|_, precompile| {
-                CachedPrecompile::wrap(precompile, self.precompile_cache.clone())
+            executor.evm_mut().precompiles_mut().map_precompiles(|address, precompile| {
+                CachedPrecompile::wrap(
+                    precompile,
+                    self.precompile_cache_map.entry(*address).or_default().clone(),
+                )
             });
         }
 
diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs
@@ -6,7 +6,6 @@ use crate::tree::{
         prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmTaskEvent},
         sparse_trie::StateRootComputeOutcome,
     },
-    precompile_cache::PrecompileCache,
     sparse_trie::SparseTrieTask,
     StateProviderBuilder, TreeConfig,
 };
@@ -39,6 +38,8 @@ use std::{
     },
 };
 
+use super::precompile_cache::PrecompileCacheMap;
+
 pub mod executor;
 pub mod multiproof;
 pub mod prewarm;
@@ -61,8 +62,8 @@ pub struct PayloadProcessor<N, Evm> {
     evm_config: Evm,
     /// whether precompile cache should be enabled.
     precompile_cache_enabled: bool,
-    /// Precompile cache.
-    precompile_cache: PrecompileCache,
+    /// Precompile cache map.
+    precompile_cache_map: PrecompileCacheMap,
     _marker: std::marker::PhantomData<N>,
 }
 
@@ -72,7 +73,7 @@ impl<N, Evm> PayloadProcessor<N, Evm> {
         executor: WorkloadExecutor,
         evm_config: Evm,
         config: &TreeConfig,
-        precompile_cache: PrecompileCache,
+        precompile_cache_map: PrecompileCacheMap,
     ) -> Self {
         Self {
             executor,
@@ -82,7 +83,7 @@ impl<N, Evm> PayloadProcessor<N, Evm> {
             disable_transaction_prewarming: config.disable_caching_and_prewarming(),
             evm_config,
             precompile_cache_enabled: config.precompile_cache_enabled(),
-            precompile_cache,
+            precompile_cache_map,
             _marker: Default::default(),
         }
     }
@@ -266,7 +267,7 @@ where
             metrics: PrewarmMetrics::default(),
             terminate_execution: Arc::new(AtomicBool::new(false)),
             precompile_cache_enabled: self.precompile_cache_enabled,
-            precompile_cache: self.precompile_cache.clone(),
+            precompile_cache_map: self.precompile_cache_map.clone(),
         };
 
         let prewarm_task = PrewarmCacheTask::new(
@@ -434,13 +435,11 @@ impl ExecutionCache {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
-
     use crate::tree::{
         payload_processor::{
             evm_state_to_hashed_post_state, executor::WorkloadExecutor, PayloadProcessor,
         },
-        precompile_cache::PrecompileCache,
+        precompile_cache::PrecompileCacheMap,
         StateProviderBuilder, TreeConfig,
     };
     use alloy_evm::block::StateChangeSource;
@@ -460,6 +459,7 @@ mod tests {
     use reth_trie::{test_utils::state_root, HashedPostState, TrieInput};
     use revm_primitives::{Address, HashMap, B256, KECCAK_EMPTY, U256};
     use revm_state::{AccountInfo, AccountStatus, EvmState, EvmStorageSlot};
+    use std::sync::Arc;
 
     fn create_mock_state_updates(num_accounts: usize, updates_per_account: usize) -> Vec<EvmState> {
         let mut rng = generators::rng();
@@ -563,7 +563,7 @@ mod tests {
             WorkloadExecutor::default(),
             EthEvmConfig::new(factory.chain_spec()),
             &TreeConfig::default(),
-            PrecompileCache::default(),
+            PrecompileCacheMap::default(),
         );
         let provider = BlockchainProvider::new(factory).unwrap();
         let mut handle = payload_processor.spawn(
diff --git a/crates/engine/tree/src/tree/payload_processor/prewarm.rs b/crates/engine/tree/src/tree/payload_processor/prewarm.rs
@@ -5,7 +5,7 @@ use crate::tree::{
     payload_processor::{
         executor::WorkloadExecutor, multiproof::MultiProofMessage, ExecutionCache,
     },
-    precompile_cache::{CachedPrecompile, PrecompileCache},
+    precompile_cache::{CachedPrecompile, PrecompileCacheMap},
     StateProviderBuilder,
 };
 use alloy_consensus::transaction::Recovered;
@@ -207,7 +207,7 @@ pub(super) struct PrewarmContext<N: NodePrimitives, P, Evm> {
     /// An atomic bool that tells prewarm tasks to not start any more execution.
     pub(super) terminate_execution: Arc<AtomicBool>,
     pub(super) precompile_cache_enabled: bool,
-    pub(super) precompile_cache: PrecompileCache,
+    pub(super) precompile_cache_map: PrecompileCacheMap,
 }
 
 impl<N, P, Evm> PrewarmContext<N, P, Evm>
@@ -230,7 +230,7 @@ where
             metrics,
             terminate_execution,
             precompile_cache_enabled,
-            precompile_cache,
+            mut precompile_cache_map,
         } = self;
 
         let state_provider = match provider.build() {
@@ -261,8 +261,11 @@ where
         let mut evm = evm_config.evm_with_env(state_provider, evm_env);
 
         if precompile_cache_enabled {
-            evm.precompiles_mut().map_precompiles(|_, precompile| {
-                CachedPrecompile::wrap(precompile, precompile_cache.clone())
+            evm.precompiles_mut().map_precompiles(|address, precompile| {
+                CachedPrecompile::wrap(
+                    precompile,
+                    precompile_cache_map.entry(*address).or_default().clone(),
+                )
             });
         }
 
diff --git a/crates/engine/tree/src/tree/precompile_cache.rs b/crates/engine/tree/src/tree/precompile_cache.rs
@@ -2,9 +2,12 @@
 
 use reth_evm::precompiles::{DynPrecompile, Precompile};
 use revm::precompile::{PrecompileOutput, PrecompileResult};
-use revm_primitives::Bytes;
+use revm_primitives::{Address, Bytes, HashMap};
 use std::sync::Arc;
 
+/// Stores caches for each precompile.
+pub type PrecompileCacheMap = HashMap<Address, PrecompileCache>;
+
 /// Cache for precompiles, for each input stores the result.
 #[derive(Debug, Clone)]
 pub struct PrecompileCache(
