refactor(engine): move CachedStateProvider prewarm to const generic (#22106)

Co-authored-by: Amp <amp@ampcode.com>

crates/engine/tree/src/tree/cached_state.rs

@@ -76,8 +76,16 @@ impl CacheConfig for EpochCacheConfig {
 type FixedCache<K, V, H = DefaultHashBuilder> = fixed_cache::Cache<K, V, H, EpochCacheConfig>;
 
 /// A wrapper of a state provider and a shared cache.
+///
+/// The const generic `PREWARM` controls whether every cache miss is populated. This is only
+/// relevant for pre-warm transaction execution with the intention to pre-populate the cache with
+/// data for regular block execution. During regular block execution the cache doesn't need to be
+/// populated because the actual EVM database [`State`](revm::database::State) also caches
+/// internally during block execution and the cache is then updated after the block with the entire
+/// [`BundleState`] output of that block which contains all accessed accounts, code, storage. See
+/// also [`ExecutionCache::insert_state`].
 #[derive(Debug)]
-pub struct CachedStateProvider<S> {
+pub struct CachedStateProvider<S, const PREWARM: bool = false> {
     /// The state provider
     state_provider: S,
 
@@ -86,15 +94,9 @@ pub struct CachedStateProvider<S> {
 
     /// Metrics for the cached state provider
     metrics: CachedStateMetrics,
-
-    /// If prewarm enabled we populate every cache miss
-    prewarm: bool,
 }
 
-impl<S> CachedStateProvider<S>
-where
-    S: StateProvider,
-{
+impl<S, const PREWARM: bool> CachedStateProvider<S, PREWARM> {
     /// Creates a new [`CachedStateProvider`] from an [`ExecutionCache`], state provider, and
     /// [`CachedStateMetrics`].
     pub const fn new(
@@ -102,27 +104,7 @@ where
         caches: ExecutionCache,
         metrics: CachedStateMetrics,
     ) -> Self {
-        Self { state_provider, caches, metrics, prewarm: false }
-    }
-}
-
-impl<S> CachedStateProvider<S> {
-    /// Enables pre-warm mode so that every cache miss is populated.
-    ///
-    /// This is only relevant for pre-warm transaction execution with the intention to pre-populate
-    /// the cache with data for regular block execution. During regular block execution the
-    /// cache doesn't need to be populated because the actual EVM database
-    /// [`State`](revm::database::State) also caches internally during block execution and the cache
-    /// is then updated after the block with the entire [`BundleState`] output of that block which
-    /// contains all accessed accounts,code,storage. See also [`ExecutionCache::insert_state`].
-    pub const fn prewarm(mut self) -> Self {
-        self.prewarm = true;
-        self
-    }
-
-    /// Returns whether this provider should pre-warm cache misses.
-    const fn is_prewarm(&self) -> bool {
-        self.prewarm
+        Self { state_provider, caches, metrics }
     }
 }
 
@@ -307,9 +289,9 @@ impl<K: PartialEq, V> StatsHandler<K, V> for CacheStatsHandler {
     }
 }
 
-impl<S: AccountReader> AccountReader for CachedStateProvider<S> {
+impl<S: AccountReader, const PREWARM: bool> AccountReader for CachedStateProvider<S, PREWARM> {
     fn basic_account(&self, address: &Address) -> ProviderResult<Option<Account>> {
-        if self.is_prewarm() {
+        if PREWARM {
             match self.caches.get_or_try_insert_account_with(*address, || {
                 self.state_provider.basic_account(address)
             })? {
@@ -334,13 +316,13 @@ pub enum CachedStatus<T> {
     Cached(T),
 }
 
-impl<S: StateProvider> StateProvider for CachedStateProvider<S> {
+impl<S: StateProvider, const PREWARM: bool> StateProvider for CachedStateProvider<S, PREWARM> {
     fn storage(
         &self,
         account: Address,
         storage_key: StorageKey,
     ) -> ProviderResult<Option<StorageValue>> {
-        if self.is_prewarm() {
+        if PREWARM {
             match self.caches.get_or_try_insert_storage_with(account, storage_key, || {
                 self.state_provider.storage(account, storage_key).map(Option::unwrap_or_default)
             })? {
@@ -360,9 +342,9 @@ impl<S: StateProvider> StateProvider for CachedStateProvider<S> {
     }
 }
 
-impl<S: BytecodeReader> BytecodeReader for CachedStateProvider<S> {
+impl<S: BytecodeReader, const PREWARM: bool> BytecodeReader for CachedStateProvider<S, PREWARM> {
     fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
-        if self.is_prewarm() {
+        if PREWARM {
             match self.caches.get_or_try_insert_code_with(*code_hash, || {
                 self.state_provider.bytecode_by_hash(code_hash)
             })? {
@@ -378,7 +360,9 @@ impl<S: BytecodeReader> BytecodeReader for CachedStateProvider<S> {
     }
 }
 
-impl<S: StateRootProvider> StateRootProvider for CachedStateProvider<S> {
+impl<S: StateRootProvider, const PREWARM: bool> StateRootProvider
+    for CachedStateProvider<S, PREWARM>
+{
     fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
         self.state_provider.state_root(hashed_state)
     }
@@ -402,7 +386,9 @@ impl<S: StateRootProvider> StateRootProvider for CachedStateProvider<S> {
     }
 }
 
-impl<S: StateProofProvider> StateProofProvider for CachedStateProvider<S> {
+impl<S: StateProofProvider, const PREWARM: bool> StateProofProvider
+    for CachedStateProvider<S, PREWARM>
+{
     fn proof(
         &self,
         input: TrieInput,
@@ -429,7 +415,9 @@ impl<S: StateProofProvider> StateProofProvider for CachedStateProvider<S> {
     }
 }
 
-impl<S: StorageRootProvider> StorageRootProvider for CachedStateProvider<S> {
+impl<S: StorageRootProvider, const PREWARM: bool> StorageRootProvider
+    for CachedStateProvider<S, PREWARM>
+{
     fn storage_root(
         &self,
         address: Address,
@@ -457,7 +445,7 @@ impl<S: StorageRootProvider> StorageRootProvider for CachedStateProvider<S> {
     }
 }
 
-impl<S: BlockHashReader> BlockHashReader for CachedStateProvider<S> {
+impl<S: BlockHashReader, const PREWARM: bool> BlockHashReader for CachedStateProvider<S, PREWARM> {
     fn block_hash(&self, number: alloy_primitives::BlockNumber) -> ProviderResult<Option<B256>> {
         self.state_provider.block_hash(number)
     }
@@ -471,7 +459,9 @@ impl<S: BlockHashReader> BlockHashReader for CachedStateProvider<S> {
     }
 }
 
-impl<S: HashedPostStateProvider> HashedPostStateProvider for CachedStateProvider<S> {
+impl<S: HashedPostStateProvider, const PREWARM: bool> HashedPostStateProvider
+    for CachedStateProvider<S, PREWARM>
+{
     fn hashed_post_state(&self, bundle_state: &reth_revm::db::BundleState) -> HashedPostState {
         self.state_provider.hashed_post_state(bundle_state)
     }
@@ -868,7 +858,7 @@ mod tests {
 
         let caches = ExecutionCache::new(1000);
         let state_provider =
-            CachedStateProvider::new(provider, caches, CachedStateMetrics::zeroed());
+            CachedStateProvider::<_, false>::new(provider, caches, CachedStateMetrics::zeroed());
 
         let res = state_provider.storage(address, storage_key);
         assert!(res.is_ok());
@@ -888,7 +878,7 @@ mod tests {
 
         let caches = ExecutionCache::new(1000);
         let state_provider =
-            CachedStateProvider::new(provider, caches, CachedStateMetrics::zeroed());
+            CachedStateProvider::<_, false>::new(provider, caches, CachedStateMetrics::zeroed());
 
         let res = state_provider.storage(address, storage_key);
         assert!(res.is_ok());

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -297,7 +297,7 @@ where
                 let provider = provider_builder.build().expect("failed to build provider");
                 let provider = if let Some(saved_cache) = saved_cache {
                     let (cache, metrics, _disable_metrics) = saved_cache.split();
-                    Box::new(CachedStateProvider::new(provider, cache, metrics))
+                    Box::new(CachedStateProvider::<_, false>::new(provider, cache, metrics))
                         as Box<dyn StateProvider>
                 } else {
                     Box::new(provider)

crates/engine/tree/src/tree/payload_processor/prewarm.rs

@@ -535,11 +535,11 @@ where
         if let Some(saved_cache) = saved_cache {
             let caches = saved_cache.cache().clone();
             let cache_metrics = saved_cache.metrics().clone();
-            state_provider = Box::new(
-                CachedStateProvider::new(state_provider, caches, cache_metrics)
-                    // ensure we pre-warm the cache
-                    .prewarm(),
-            );
+            state_provider = Box::new(CachedStateProvider::<_, true>::new(
+                state_provider,
+                caches,
+                cache_metrics,
+            ));
         }
 
         let state_provider = StateProviderDatabase::new(state_provider);
@@ -749,7 +749,8 @@ where
         let saved_cache = saved_cache.expect("BAL prewarm should only run with cache");
         let caches = saved_cache.cache().clone();
         let cache_metrics = saved_cache.metrics().clone();
-        let state_provider = CachedStateProvider::new(state_provider, caches, cache_metrics);
+        let state_provider =
+            CachedStateProvider::<_, false>::new(state_provider, caches, cache_metrics);
 
         let start = Instant::now();
 

crates/engine/tree/src/tree/payload_validator.rs

@@ -443,8 +443,11 @@ where
         // Use cached state provider before executing, used in execution after prewarming threads
         // complete
         if let Some((caches, cache_metrics)) = handle.caches().zip(handle.cache_metrics()) {
-            state_provider =
-                Box::new(CachedStateProvider::new(state_provider, caches, cache_metrics));
+            state_provider = Box::new(CachedStateProvider::<_, false>::new(
+                state_provider,
+                caches,
+                cache_metrics,
+            ));
         };
 
         if self.config.state_provider_metrics() {