From 9c10db47f1324f378d52b8df7874d53d0983148c Mon Sep 17 00:00:00 2001
From: AvivYossef-starkware <aviv.y@starkware.co>
Date: Tue, 3 Dec 2024 17:31:19 +0200
Subject: [PATCH] feat(papyrus_storage): add sierra version to papyrus cache

---
 .../blockifier/src/execution/contract_class.rs |  1 +
 .../native_blockifier/src/py_block_executor.rs |  4 ++--
 .../papyrus_state_reader/src/papyrus_state.rs  | 18 +++++++++---------
 crates/starknet_batcher/src/block_builder.rs   |  4 ++--
 4 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/crates/blockifier/src/execution/contract_class.rs b/crates/blockifier/src/execution/contract_class.rs
index 85c135f7a36..976397bd182 100644
--- a/crates/blockifier/src/execution/contract_class.rs
+++ b/crates/blockifier/src/execution/contract_class.rs
@@ -68,6 +68,7 @@ pub enum RunnableCompiledClass {
 }
 
 /// Represents a runnable compiled class for Cairo, with the Sierra version (for Cairo 1).
+#[derive(Clone)]
 pub enum VersionedRunnableCompiledClass {
     Cairo0(RunnableCompiledClass),
     Cairo1((RunnableCompiledClass, SierraVersion)),
diff --git a/crates/native_blockifier/src/py_block_executor.rs b/crates/native_blockifier/src/py_block_executor.rs
index d515e3a2da2..af89ed8a5d8 100644
--- a/crates/native_blockifier/src/py_block_executor.rs
+++ b/crates/native_blockifier/src/py_block_executor.rs
@@ -8,7 +8,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::RunnableCompiledClass;
+use blockifier::execution::contract_class::VersionedRunnableCompiledClass;
 use blockifier::fee::receipt::TransactionReceipt;
 use blockifier::state::global_cache::GlobalContractCache;
 use blockifier::transaction::objects::{ExecutionResourcesTraits, TransactionExecutionInfo};
@@ -138,7 +138,7 @@ pub struct PyBlockExecutor {
     /// `Send` trait is required for `pyclass` compatibility as Python objects must be threadsafe.
     pub storage: Box<dyn Storage + Send>,
     pub contract_class_manager_config: ContractClassManagerConfig,
-    pub global_contract_cache: GlobalContractCache<RunnableCompiledClass>,
+    pub global_contract_cache: GlobalContractCache<VersionedRunnableCompiledClass>,
 }
 
 #[pymethods]
diff --git a/crates/papyrus_state_reader/src/papyrus_state.rs b/crates/papyrus_state_reader/src/papyrus_state.rs
index edc8e465090..503b7e4311a 100644
--- a/crates/papyrus_state_reader/src/papyrus_state.rs
+++ b/crates/papyrus_state_reader/src/papyrus_state.rs
@@ -25,14 +25,14 @@ type RawPapyrusReader<'env> = papyrus_storage::StorageTxn<'env, RO>;
 pub struct PapyrusReader {
     storage_reader: StorageReader,
     latest_block: BlockNumber,
-    global_class_hash_to_class: GlobalContractCache<RunnableCompiledClass>,
+    global_class_hash_to_class: GlobalContractCache<VersionedRunnableCompiledClass>,
 }
 
 impl PapyrusReader {
     pub fn new(
         storage_reader: StorageReader,
         latest_block: BlockNumber,
-        global_class_hash_to_class: GlobalContractCache<RunnableCompiledClass>,
+        global_class_hash_to_class: GlobalContractCache<VersionedRunnableCompiledClass>,
     ) -> Self {
         Self { storage_reader, latest_block, global_class_hash_to_class }
     }
@@ -131,16 +131,16 @@ impl StateReader for PapyrusReader {
 
     fn get_compiled_class(&self, class_hash: ClassHash) -> StateResult<RunnableCompiledClass> {
         // Assumption: the global cache is cleared upon reverted blocks.
-        let contract_class = self.global_class_hash_to_class.get(&class_hash);
+        let versioned_contract_class = self.global_class_hash_to_class.get(&class_hash);
 
-        match contract_class {
-            Some(contract_class) => Ok(contract_class),
+        match versioned_contract_class {
+            Some(contract_class) => Ok(RunnableCompiledClass::from(contract_class)),
             None => {
-                let contract_class_from_db =
-                    RunnableCompiledClass::from(self.get_compiled_class_inner(class_hash)?);
+                let versioned_contract_class_from_db = self.get_compiled_class_inner(class_hash)?;
                 // The class was declared in a previous (finalized) state; update the global cache.
-                self.global_class_hash_to_class.set(class_hash, contract_class_from_db.clone());
-                Ok(contract_class_from_db)
+                self.global_class_hash_to_class
+                    .set(class_hash, versioned_contract_class_from_db.clone());
+                Ok(RunnableCompiledClass::from(versioned_contract_class_from_db))
             }
         }
     }
diff --git a/crates/starknet_batcher/src/block_builder.rs b/crates/starknet_batcher/src/block_builder.rs
index 6b59aa9499a..4ec46dd11ae 100644
--- a/crates/starknet_batcher/src/block_builder.rs
+++ b/crates/starknet_batcher/src/block_builder.rs
@@ -10,7 +10,7 @@ use blockifier::blockifier::transaction_executor::{
 };
 use blockifier::bouncer::{BouncerConfig, BouncerWeights};
 use blockifier::context::{BlockContext, ChainInfo};
-use blockifier::execution::contract_class::RunnableCompiledClass;
+use blockifier::execution::contract_class::VersionedRunnableCompiledClass;
 use blockifier::state::cached_state::CommitmentStateDiff;
 use blockifier::state::errors::StateError;
 use blockifier::state::global_cache::GlobalContractCache;
@@ -290,7 +290,7 @@ impl SerializeConfig for BlockBuilderConfig {
 pub struct BlockBuilderFactory {
     pub block_builder_config: BlockBuilderConfig,
     pub storage_reader: StorageReader,
-    pub global_class_hash_to_class: GlobalContractCache<RunnableCompiledClass>,
+    pub global_class_hash_to_class: GlobalContractCache<VersionedRunnableCompiledClass>,
 }
 
 impl BlockBuilderFactory {