From 13960a87a6ec6d48e36e2379008ee9597cbf50e9 Mon Sep 17 00:00:00 2001
From: meship-starkware <meshi@starkware.co>
Date: Wed, 25 Dec 2024 14:03:42 +0200
Subject: [PATCH] test(cairo_native): test papyrus state reader get compiled
 class

---
 crates/blockifier/src/blockifier/config.rs    | 10 ++++
 .../papyrus_state_reader/src/papyrus_state.rs |  1 +
 .../src/papyrus_state_test.rs                 | 59 ++++++++++++++++++-
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/crates/blockifier/src/blockifier/config.rs b/crates/blockifier/src/blockifier/config.rs
index aca073f1b4..8955ac507d 100644
--- a/crates/blockifier/src/blockifier/config.rs
+++ b/crates/blockifier/src/blockifier/config.rs
@@ -81,6 +81,16 @@ impl Default for ContractClassManagerConfig {
     }
 }
 
+impl ContractClassManagerConfig {
+    #[cfg(any(test, feature = "testing", feature = "native_blockifier"))]
+    pub fn create_for_testing(
+        run_cairo_native: bool,
+        wait_on_native_compilation: bool,
+    ) -> Self {
+        Self { run_cairo_native, wait_on_native_compilation, contract_cache_size: GLOBAL_CONTRACT_CACHE_SIZE_FOR_TEST }
+    }
+}
+
 impl SerializeConfig for ContractClassManagerConfig {
     fn dump(&self) -> BTreeMap<ParamPath, SerializedParam> {
         BTreeMap::from_iter([
diff --git a/crates/papyrus_state_reader/src/papyrus_state.rs b/crates/papyrus_state_reader/src/papyrus_state.rs
index 1bc09381b8..8cf7dffa05 100644
--- a/crates/papyrus_state_reader/src/papyrus_state.rs
+++ b/crates/papyrus_state_reader/src/papyrus_state.rs
@@ -58,6 +58,7 @@ impl PapyrusReader {
             .get_state_reader()
             .and_then(|sr| sr.get_class_definition_block_number(&class_hash))
             .map_err(|err| StateError::StateReadError(err.to_string()))?;
+
         let class_is_declared: bool = matches!(class_declaration_block_number,
                         Some(block_number) if block_number <= state_number.0);
 
diff --git a/crates/papyrus_state_reader/src/papyrus_state_test.rs b/crates/papyrus_state_reader/src/papyrus_state_test.rs
index e3e006a496..83fce2f317 100644
--- a/crates/papyrus_state_reader/src/papyrus_state_test.rs
+++ b/crates/papyrus_state_reader/src/papyrus_state_test.rs
@@ -1,20 +1,25 @@
+use core::panic;
+
 use assert_matches::assert_matches;
 use blockifier::blockifier::config::ContractClassManagerConfig;
 use blockifier::execution::call_info::CallExecution;
+use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::execution::entry_point::CallEntryPoint;
 use blockifier::retdata;
 use blockifier::state::cached_state::CachedState;
 use blockifier::state::contract_class_manager::ContractClassManager;
 use blockifier::state::state_api::StateReader;
 use blockifier::test_utils::contracts::FeatureContract;
-use blockifier::test_utils::{trivial_external_entry_point_new, CairoVersion};
+use blockifier::test_utils::{trivial_external_entry_point_new, CairoVersion, RunnableCairo1};
 use indexmap::IndexMap;
 use papyrus_storage::class::ClassStorageWriter;
+use papyrus_storage::compiled_class::CasmStorageWriter;
 use papyrus_storage::state::StateStorageWriter;
 use starknet_api::abi::abi_utils::selector_from_name;
 use starknet_api::block::BlockNumber;
 use starknet_api::contract_class::ContractClass;
-use starknet_api::state::{StateDiff, StorageKey};
+use starknet_api::core::Nonce;
+use starknet_api::state::{StateDiff, StorageKey, ThinStateDiff};
 use starknet_api::{calldata, felt};
 
 use crate::papyrus_state::PapyrusReader;
@@ -76,3 +81,53 @@ fn test_entry_point_with_papyrus_state() -> papyrus_storage::StorageResult<()> {
 
     Ok(())
 }
+
+#[cfg(feature = "cairo_native")]
+#[test]
+fn test_get_compiled_class() -> papyrus_storage::StorageResult<()> {
+    use std::sync::Arc;
+
+    use blockifier::state::global_cache::CachedCasm;
+
+    let ((storage_reader, mut storage_writer), _) = papyrus_storage::test_utils::get_test_storage();
+    let test_contract = FeatureContract::TestContract(CairoVersion::Cairo1(RunnableCairo1::Casm));
+    let compiled_class = if let ContractClass::V1((casm_class, _)) = test_contract.get_class() {
+        casm_class
+    } else {
+        panic!("expected a V1 class")
+    };
+    let test_class_hash = test_contract.get_class_hash();
+    let test_compiled_class_hash = test_contract.get_compiled_class_hash();
+    let block_number = BlockNumber::default();
+
+    let thin_state_diff = ThinStateDiff {
+        declared_classes: IndexMap::from([(test_class_hash, test_compiled_class_hash)]),
+        nonces: IndexMap::from([(test_contract.get_instance_address(1), Nonce(1.into()))]),
+        ..Default::default()
+    };
+    let _ = storage_writer
+        .begin_rw_txn()?
+        .append_state_diff(block_number, thin_state_diff)?
+        .append_classes(block_number, &[(test_class_hash, &test_contract.get_sierra())], &[])?
+        .append_casm(&test_class_hash, &compiled_class)?
+        .commit();
+    let contract_manager_config = ContractClassManagerConfig::create_for_testing(true, true);
+
+
+    let papyrus_reader = PapyrusReader::new(
+        storage_reader,
+        block_number,
+        ContractClassManager::start(contract_manager_config),
+    );
+    papyrus_reader.contract_class_manager.set_casm(test_class_hash, CachedCasm::WithSierra(test_contract.get_runnable_class(), Arc::new(test_contract.get_sierra())));
+    let compiled_class = papyrus_reader.get_compiled_class(test_class_hash);
+    assert!(compiled_class.is_ok(), "compilation should have succeeded in this case");
+    assert!(false);
+    assert_matches!(
+        compiled_class,
+        Ok(RunnableCompiledClass::V1Native(_)),
+        "compilation should have succeeded in this case"
+    );
+
+    Ok(())
+}