feat(katana): global compiled class cache (#3004)

crates/katana/executor/Cargo.toml

@@ -10,7 +10,6 @@ version.workspace = true
 katana-cairo = { workspace = true, optional = true }
 katana-primitives.workspace = true
 katana-provider.workspace = true
-katana-trie.workspace = true
 
 thiserror.workspace = true
 tracing.workspace = true

crates/katana/executor/benches/execution.rs

@@ -3,7 +3,8 @@ use std::time::Duration;
 use blockifier::state::cached_state::CachedState;
 use criterion::measurement::WallTime;
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkGroup, Criterion};
-use katana_executor::{ExecutionFlags, StateProviderDb};
+use katana_executor::implementation::blockifier::state::StateProviderDb;
+use katana_executor::ExecutionFlags;
 use katana_primitives::env::{BlockEnv, CfgEnv};
 use katana_primitives::transaction::ExecutableTxWithHash;
 use katana_provider::test_utils;
@@ -45,7 +46,7 @@ fn blockifier(
             || {
                 // setup state
                 let state = provider.latest().expect("failed to get latest state");
-                let state = CachedState::new(StateProviderDb::new(state));
+                let state = CachedState::new(StateProviderDb::new(state, Default::default()));
 
                 (state, &block_context, execution_flags, tx.clone())
             },

crates/katana/executor/src/abstraction/mod.rs

@@ -1,21 +1,15 @@
-mod error;
 mod executor;
 
-pub use error::*;
 pub use executor::*;
-use katana_primitives::class::{ClassHash, CompiledClass, CompiledClassHash, ContractClass};
-use katana_primitives::contract::{ContractAddress, Nonce, StorageKey, StorageValue};
 use katana_primitives::receipt::Receipt;
 use katana_primitives::state::{StateUpdates, StateUpdatesWithClasses};
 use katana_primitives::trace::TxExecInfo;
 use katana_primitives::transaction::TxWithHash;
-use katana_primitives::Felt;
-use katana_provider::traits::contract::ContractClassProvider;
-use katana_provider::traits::state::{StateProofProvider, StateProvider, StateRootProvider};
-use katana_provider::ProviderResult;
-use katana_trie::MultiProof;
+use katana_primitives::{ContractAddress, Felt};
 
-pub type ExecutorResult<T> = Result<T, error::ExecutorError>;
+pub use crate::error::*;
+
+pub type ExecutorResult<T> = Result<T, crate::error::ExecutorError>;
 
 /// See <https://docs.starknet.io/chain-info/#current_limits>.
 #[derive(Debug, Clone, Default)]
@@ -169,88 +163,3 @@ pub struct ResultAndStates {
     pub result: ExecutionResult,
     pub states: StateUpdates,
 }
-
-/// A wrapper around a boxed [StateProvider] for implementing the executor's own state reader
-/// traits.
-#[derive(Debug)]
-pub struct StateProviderDb<'a>(Box<dyn StateProvider + 'a>);
-
-impl<'a> StateProviderDb<'a> {
-    pub fn new(provider: Box<dyn StateProvider + 'a>) -> Self {
-        Self(provider)
-    }
-}
-
-impl<'a> ContractClassProvider for StateProviderDb<'a> {
-    fn class(&self, hash: ClassHash) -> ProviderResult<Option<ContractClass>> {
-        self.0.class(hash)
-    }
-
-    fn compiled_class(&self, hash: ClassHash) -> ProviderResult<Option<CompiledClass>> {
-        self.0.compiled_class(hash)
-    }
-
-    fn compiled_class_hash_of_class_hash(
-        &self,
-        hash: ClassHash,
-    ) -> ProviderResult<Option<CompiledClassHash>> {
-        self.0.compiled_class_hash_of_class_hash(hash)
-    }
-}
-
-impl<'a> StateProvider for StateProviderDb<'a> {
-    fn class_hash_of_contract(
-        &self,
-        address: ContractAddress,
-    ) -> ProviderResult<Option<ClassHash>> {
-        self.0.class_hash_of_contract(address)
-    }
-
-    fn nonce(&self, address: ContractAddress) -> ProviderResult<Option<Nonce>> {
-        self.0.nonce(address)
-    }
-
-    fn storage(
-        &self,
-        address: ContractAddress,
-        storage_key: StorageKey,
-    ) -> ProviderResult<Option<StorageValue>> {
-        self.0.storage(address, storage_key)
-    }
-}
-
-impl<'a> StateProofProvider for StateProviderDb<'a> {
-    fn class_multiproof(&self, classes: Vec<ClassHash>) -> ProviderResult<MultiProof> {
-        self.0.class_multiproof(classes)
-    }
-
-    fn contract_multiproof(&self, addresses: Vec<ContractAddress>) -> ProviderResult<MultiProof> {
-        self.0.contract_multiproof(addresses)
-    }
-
-    fn storage_multiproof(
-        &self,
-        address: ContractAddress,
-        key: Vec<StorageKey>,
-    ) -> ProviderResult<MultiProof> {
-        self.0.storage_multiproof(address, key)
-    }
-}
-
-impl<'a> StateRootProvider for StateProviderDb<'a> {
-    fn classes_root(&self) -> ProviderResult<Felt> {
-        self.0.classes_root()
-    }
-
-    fn contracts_root(&self) -> ProviderResult<Felt> {
-        self.0.contracts_root()
-    }
-
-    fn storage_root(&self, contract: ContractAddress) -> ProviderResult<Option<Felt>> {
-        self.0.storage_root(contract)
-    }
-
-    fn state_root(&self) -> ProviderResult<Felt> {
-        self.0.state_root()
-    }
-}

crates/katana/executor/src/implementation/blockifier/mod.rs

@@ -3,31 +3,39 @@ pub use blockifier;
 use blockifier::bouncer::{Bouncer, BouncerConfig, BouncerWeights};
 
 mod error;
-mod state;
+pub mod state;
 pub mod utils;
 
+use std::collections::HashMap;
 use std::num::NonZeroU128;
+use std::sync::{Arc, LazyLock};
 
 use blockifier::blockifier::block::{BlockInfo, GasPrices};
 use blockifier::context::BlockContext;
+use blockifier::execution::contract_class::ContractClass as BlockifierContractClass;
 use blockifier::state::cached_state::{self, MutRefState};
 use blockifier::state::state_api::StateReader;
 use katana_cairo::starknet_api::block::{BlockNumber, BlockTimestamp};
 use katana_primitives::block::{ExecutableBlock, GasPrices as KatanaGasPrices, PartialHeader};
+use katana_primitives::class::ClassHash;
 use katana_primitives::env::{BlockEnv, CfgEnv};
 use katana_primitives::fee::TxFeeInfo;
 use katana_primitives::transaction::{ExecutableTx, ExecutableTxWithHash, TxWithHash};
 use katana_primitives::Felt;
 use katana_provider::traits::state::StateProvider;
+use parking_lot::Mutex;
 use tracing::info;
 
 use self::state::CachedState;
 use crate::{
     BlockExecutor, BlockLimits, EntryPointCall, ExecutionError, ExecutionFlags, ExecutionOutput,
     ExecutionResult, ExecutionStats, ExecutorError, ExecutorExt, ExecutorFactory, ExecutorResult,
-    ResultAndStates, StateProviderDb,
+    ResultAndStates,
 };
 
+pub static COMPILED_CLASS_CACHE: LazyLock<Arc<Mutex<HashMap<ClassHash, BlockifierContractClass>>>> =
+    LazyLock::new(|| Arc::new(Mutex::new(HashMap::default())));
+
 pub(crate) const LOG_TARGET: &str = "katana::executor::blockifier";
 
 #[derive(Debug)]
@@ -79,7 +87,7 @@ impl ExecutorFactory for BlockifierFactory {
 #[derive(Debug)]
 pub struct StarknetVMProcessor<'a> {
     block_context: BlockContext,
-    state: CachedState<StateProviderDb<'a>>,
+    state: CachedState<'a>,
     transactions: Vec<(TxWithHash, ExecutionResult)>,
     simulation_flags: ExecutionFlags,
     stats: ExecutionStats,
@@ -88,15 +96,15 @@ pub struct StarknetVMProcessor<'a> {
 
 impl<'a> StarknetVMProcessor<'a> {
     pub fn new(
-        state: Box<dyn StateProvider + 'a>,
+        state: impl StateProvider + 'a,
         block_env: BlockEnv,
         cfg_env: CfgEnv,
         simulation_flags: ExecutionFlags,
         limits: BlockLimits,
     ) -> Self {
         let transactions = Vec::new();
         let block_context = utils::block_context_from_envs(&block_env, &cfg_env);
-        let state = state::CachedState::new(StateProviderDb::new(state));
+        let state = state::CachedState::new(state, COMPILED_CLASS_CACHE.clone());
 
         let mut block_max_capacity = BouncerWeights::max();
         block_max_capacity.n_steps = limits.cairo_steps as usize;
@@ -159,7 +167,7 @@ impl<'a> StarknetVMProcessor<'a> {
         F: FnMut(&mut dyn StateReader, (TxWithHash, ExecutionResult)) -> T,
     {
         let block_context = &self.block_context;
-        let state = &mut self.state.0.lock().inner;
+        let state = &mut self.state.inner.lock().cached_state;
         let mut state = cached_state::CachedState::new(MutRefState::new(state));
 
         let mut results = Vec::with_capacity(transactions.len());
@@ -188,7 +196,7 @@ impl<'a> BlockExecutor<'a> for StarknetVMProcessor<'a> {
     ) -> ExecutorResult<(usize, Option<ExecutorError>)> {
         let block_context = &self.block_context;
         let flags = &self.simulation_flags;
-        let mut state = self.state.0.lock();
+        let mut state = self.state.inner.lock();
 
         let mut total_executed = 0;
         for exec_tx in transactions {
@@ -203,7 +211,7 @@ impl<'a> BlockExecutor<'a> for StarknetVMProcessor<'a> {
             let tx = TxWithHash::from(&exec_tx);
             let hash = tx.hash;
             let result = utils::transact(
-                &mut state.inner,
+                &mut state.cached_state,
                 block_context,
                 flags,
                 exec_tx,
@@ -318,8 +326,8 @@ impl ExecutorExt for StarknetVMProcessor<'_> {
 
     fn call(&self, call: EntryPointCall) -> Result<Vec<Felt>, ExecutionError> {
         let block_context = &self.block_context;
-        let mut state = self.state.0.lock();
-        let state = MutRefState::new(&mut state.inner);
+        let mut state = self.state.inner.lock();
+        let state = MutRefState::new(&mut state.cached_state);
         let retdata = utils::call(call, state, block_context, 1_000_000_000)?;
         Ok(retdata)
     }