From 271aa9b7d314a0fc520138551412aa07a81520d7 Mon Sep 17 00:00:00 2001
From: Avi Cohen <avi.cohen@starkware.co>
Date: Tue, 19 Nov 2024 16:05:44 +0200
Subject: [PATCH] feat(cairo_native): add a struct containing all global
 contract caches

---
 crates/blockifier/src/state/global_cache.rs   | 27 +++++++++++++------
 .../papyrus_state_reader/src/papyrus_state.rs |  4 +--
 crates/starknet_batcher/src/block_builder.rs  |  9 +++----
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/crates/blockifier/src/state/global_cache.rs b/crates/blockifier/src/state/global_cache.rs
index 670045fe5f9..6459feed0ac 100644
--- a/crates/blockifier/src/state/global_cache.rs
+++ b/crates/blockifier/src/state/global_cache.rs
@@ -2,32 +2,36 @@ use std::sync::{Arc, Mutex, MutexGuard};
 
 use cached::{Cached, SizedCache};
 use starknet_api::core::ClassHash;
+#[cfg(feature = "cairo_native")]
+use starknet_api::state::ContractClass as SierraContractClass;
 
 use crate::execution::contract_class::RunnableContractClass;
+#[cfg(feature = "cairo_native")]
+use crate::execution::native::contract_class::NativeContractClassV1;
 
 // 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>>>);
 
 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 +40,13 @@ 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_contract_class_cache: GlobalContractCache<RunnableContractClass>,
+    pub cairo_native_contract_class_cache: GlobalContractCache<Option<NativeContractClassV1>>,
+    pub sierra_contract_class_cache: GlobalContractCache<Arc<SierraContractClass>>,
+}
diff --git a/crates/papyrus_state_reader/src/papyrus_state.rs b/crates/papyrus_state_reader/src/papyrus_state.rs
index b957a170b09..3ff2cdc46e3 100644
--- a/crates/papyrus_state_reader/src/papyrus_state.rs
+++ b/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 }
     }
diff --git a/crates/starknet_batcher/src/block_builder.rs b/crates/starknet_batcher/src/block_builder.rs
index a16e974eee6..668ba712604 100644
--- a/crates/starknet_batcher/src/block_builder.rs
+++ b/crates/starknet_batcher/src/block_builder.rs
@@ -4,13 +4,12 @@ use async_trait::async_trait;
 use blockifier::blockifier::block::{BlockInfo, GasPrices};
 use blockifier::blockifier::config::TransactionExecutorConfig;
 use blockifier::blockifier::transaction_executor::{
-    TransactionExecutor,
-    TransactionExecutorError as BlockifierTransactionExecutorError,
-    TransactionExecutorResult,
-    VisitedSegmentsMapping,
+    TransactionExecutor, TransactionExecutorError as BlockifierTransactionExecutorError,
+    TransactionExecutorResult, VisitedSegmentsMapping,
 };
 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;
@@ -276,7 +275,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 {