feat(cairo_native): add a struct containing all global contract caches (

#2186)

crates/blockifier/src/state/global_cache.rs

@@ -1,33 +1,44 @@
 use std::sync::{Arc, Mutex, MutexGuard};
 
 use cached::{Cached, SizedCache};
+#[cfg(feature = "cairo_native")]
+use cairo_native::executor::AotContractExecutor;
 use starknet_api::core::ClassHash;
+#[cfg(feature = "cairo_native")]
+use starknet_api::state::ContractClass as SierraContractClass;
 
+#[cfg(feature = "cairo_native")]
 use crate::execution::contract_class::RunnableContractClass;
 
-// Note: `ContractClassLRUCache` key-value types must align with `ContractClassMapping`.
-type ContractClassLRUCache = SizedCache<ClassHash, RunnableContractClass>;
-pub type LockedContractClassCache<'a> = MutexGuard<'a, ContractClassLRUCache>;
+type ContractClassLRUCache<T> = SizedCache<ClassHash, T>;
+pub type LockedContractClassCache<'a, T> = MutexGuard<'a, ContractClassLRUCache<T>>;
 #[derive(Debug, Clone)]
 // Thread-safe LRU cache for contract classes, optimized for inter-language sharing when
 // `blockifier` compiles as a shared library.
 // TODO(Yoni, 1/1/2025): consider defining CachedStateReader.
-pub struct GlobalContractCache(pub Arc<Mutex<ContractClassLRUCache>>);
+pub struct GlobalContractCache<T: Clone>(pub Arc<Mutex<ContractClassLRUCache<T>>>);
+
+#[cfg(feature = "cairo_native")]
+#[derive(Debug, Clone)]
+pub enum CachedCairoNative {
+    Compiled(AotContractExecutor),
+    CompilationFailed,
+}
 
 pub const GLOBAL_CONTRACT_CACHE_SIZE_FOR_TEST: usize = 100;
 
-impl GlobalContractCache {
+impl<T: Clone> GlobalContractCache<T> {
     /// Locks the cache for atomic access. Although conceptually shared, writing to this cache is
     /// only possible for one writer at a time.
-    pub fn lock(&self) -> LockedContractClassCache<'_> {
+    pub fn lock(&self) -> LockedContractClassCache<'_, T> {
         self.0.lock().expect("Global contract cache is poisoned.")
     }
 
-    pub fn get(&self, class_hash: &ClassHash) -> Option<RunnableContractClass> {
+    pub fn get(&self, class_hash: &ClassHash) -> Option<T> {
         self.lock().cache_get(class_hash).cloned()
     }
 
-    pub fn set(&self, class_hash: ClassHash, contract_class: RunnableContractClass) {
+    pub fn set(&self, class_hash: ClassHash, contract_class: T) {
         self.lock().cache_set(class_hash, contract_class);
     }
 
@@ -36,6 +47,54 @@ impl GlobalContractCache {
     }
 
     pub fn new(cache_size: usize) -> Self {
-        Self(Arc::new(Mutex::new(ContractClassLRUCache::with_size(cache_size))))
+        Self(Arc::new(Mutex::new(ContractClassLRUCache::<T>::with_size(cache_size))))
+    }
+}
+
+#[cfg(feature = "cairo_native")]
+pub struct GlobalContractCacheManager {
+    pub casm_cache: GlobalContractCache<RunnableContractClass>,
+    pub native_cache: GlobalContractCache<CachedCairoNative>,
+    pub sierra_cache: GlobalContractCache<Arc<SierraContractClass>>,
+}
+
+#[cfg(feature = "cairo_native")]
+impl GlobalContractCacheManager {
+    pub fn get_casm(&self, class_hash: &ClassHash) -> Option<RunnableContractClass> {
+        self.casm_cache.get(class_hash)
+    }
+
+    pub fn set_casm(&self, class_hash: ClassHash, contract_class: RunnableContractClass) {
+        self.casm_cache.set(class_hash, contract_class);
+    }
+
+    pub fn get_native(&self, class_hash: &ClassHash) -> Option<CachedCairoNative> {
+        self.native_cache.get(class_hash)
+    }
+
+    pub fn set_native(&self, class_hash: ClassHash, contract_executor: CachedCairoNative) {
+        self.native_cache.set(class_hash, contract_executor);
+    }
+
+    pub fn get_sierra(&self, class_hash: &ClassHash) -> Option<Arc<SierraContractClass>> {
+        self.sierra_cache.get(class_hash)
+    }
+
+    pub fn set_sierra(&self, class_hash: ClassHash, contract_class: Arc<SierraContractClass>) {
+        self.sierra_cache.set(class_hash, contract_class);
+    }
+
+    pub fn new(cache_size: usize) -> Self {
+        Self {
+            casm_cache: GlobalContractCache::new(cache_size),
+            native_cache: GlobalContractCache::new(cache_size),
+            sierra_cache: GlobalContractCache::new(cache_size),
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.casm_cache.clear();
+        self.native_cache.clear();
+        self.sierra_cache.clear();
     }
 }

crates/native_blockifier/src/py_block_executor.rs

@@ -6,6 +6,7 @@ use blockifier::blockifier::transaction_executor::{TransactionExecutor, Transact
 use blockifier::bouncer::BouncerConfig;
 use blockifier::context::{BlockContext, ChainInfo, FeeTokenAddresses};
 use blockifier::execution::call_info::CallInfo;
+use blockifier::execution::contract_class::RunnableContractClass;
 use blockifier::fee::receipt::TransactionReceipt;
 use blockifier::state::global_cache::GlobalContractCache;
 use blockifier::transaction::objects::{ExecutionResourcesTraits, TransactionExecutionInfo};
@@ -129,7 +130,7 @@ pub struct PyBlockExecutor {
     pub tx_executor: Option<TransactionExecutor<PapyrusReader>>,
     /// `Send` trait is required for `pyclass` compatibility as Python objects must be threadsafe.
     pub storage: Box<dyn Storage + Send>,
-    pub global_contract_cache: GlobalContractCache,
+    pub global_contract_cache: GlobalContractCache<RunnableContractClass>,
 }
 
 #[pymethods]

crates/papyrus_state_reader/src/papyrus_state.rs

@@ -23,14 +23,14 @@ type RawPapyrusReader<'env> = papyrus_storage::StorageTxn<'env, RO>;
 pub struct PapyrusReader {
     storage_reader: StorageReader,
     latest_block: BlockNumber,
-    global_class_hash_to_class: GlobalContractCache,
+    global_class_hash_to_class: GlobalContractCache<RunnableContractClass>,
 }
 
 impl PapyrusReader {
     pub fn new(
         storage_reader: StorageReader,
         latest_block: BlockNumber,
-        global_class_hash_to_class: GlobalContractCache,
+        global_class_hash_to_class: GlobalContractCache<RunnableContractClass>,
     ) -> Self {
         Self { storage_reader, latest_block, global_class_hash_to_class }
     }

crates/starknet_batcher/src/block_builder.rs

@@ -11,6 +11,7 @@ use blockifier::blockifier::transaction_executor::{
 };
 use blockifier::bouncer::{BouncerConfig, BouncerWeights};
 use blockifier::context::{BlockContext, ChainInfo};
+use blockifier::execution::contract_class::RunnableContractClass;
 use blockifier::state::cached_state::CommitmentStateDiff;
 use blockifier::state::errors::StateError;
 use blockifier::state::global_cache::GlobalContractCache;
@@ -295,7 +296,7 @@ impl SerializeConfig for BlockBuilderConfig {
 pub struct BlockBuilderFactory {
     pub block_builder_config: BlockBuilderConfig,
     pub storage_reader: StorageReader,
-    pub global_class_hash_to_class: GlobalContractCache,
+    pub global_class_hash_to_class: GlobalContractCache<RunnableContractClass>,
 }
 
 impl BlockBuilderFactory {