fix(blockifier_reexecution): fix get_contract_class signature and rem…

…ove duplicate (#1736)
starkware-libs · Nov 3, 2024 · 4df20c4 · 4df20c4
1 parent f98f87e
commit 4df20c4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 25 deletions.
diff --git a/crates/blockifier_reexecution/src/state_reader/reexecution_state_reader.rs b/crates/blockifier_reexecution/src/state_reader/reexecution_state_reader.rs
@@ -1,4 +1,5 @@
 use blockifier::execution::contract_class::ClassInfo;
+use blockifier::state::state_api::StateResult;
 use blockifier::transaction::transaction_execution::Transaction as BlockifierTransaction;
 use papyrus_execution::DEPRECATED_CONTRACT_SIERRA_SIZE;
 use starknet_api::core::ClassHash;
@@ -10,11 +11,10 @@ use crate::state_reader::errors::ReexecutionError;
 use crate::state_reader::test_state_reader::ReexecutionResult;
 
 pub(crate) trait ReexecutionStateReader {
-    fn get_contract_class(&self, class_hash: ClassHash)
-    -> ReexecutionResult<StarknetContractClass>;
+    fn get_contract_class(&self, class_hash: &ClassHash) -> StateResult<StarknetContractClass>;
 
     fn get_class_info(&self, class_hash: ClassHash) -> ReexecutionResult<ClassInfo> {
-        match self.get_contract_class(class_hash)? {
+        match self.get_contract_class(&class_hash)? {
             StarknetContractClass::Sierra(sierra) => {
                 let abi_length = sierra.abi.len();
                 let sierra_length = sierra.sierra_program.len();

diff --git a/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs b/crates/blockifier_reexecution/src/state_reader/test_state_reader.rs
@@ -158,24 +158,6 @@ impl TestStateReader {
         )?)
     }
 
-    pub fn get_contract_class(&self, class_hash: &ClassHash) -> StateResult<StarknetContractClass> {
-        let params = json!({
-            "block_id": self.rpc_state_reader.block_id,
-            "class_hash": class_hash.0.to_string(),
-        });
-        let contract_class: StarknetContractClass = serde_json::from_value(
-            self.rpc_state_reader.send_rpc_request("starknet_getClass", params.clone())?,
-        )
-        .map_err(serde_err_to_state_err)?;
-        // Create a binding to avoid value being dropped.
-        let mut dumper_binding = self.contract_class_mapping_dumper.lock().unwrap();
-        // If dumper exists, insert the contract class to the mapping.
-        if let Some(contract_class_mapping_dumper) = dumper_binding.as_mut() {
-            contract_class_mapping_dumper.insert(*class_hash, contract_class.clone());
-        }
-        Ok(contract_class)
-    }
-
     pub fn get_all_txs_in_block(&self) -> ReexecutionResult<Vec<(Transaction, TransactionHash)>> {
         // TODO(Aviv): Use batch request to get all txs in a block.
         self.get_tx_hashes()?
@@ -262,10 +244,7 @@ impl TestStateReader {
 }
 
 impl ReexecutionStateReader for TestStateReader {
-    fn get_contract_class(
-        &self,
-        class_hash: ClassHash,
-    ) -> ReexecutionResult<StarknetContractClass> {
+    fn get_contract_class(&self, class_hash: &ClassHash) -> StateResult<StarknetContractClass> {
         let params = json!({
             "block_id": self.rpc_state_reader.block_id,
             "class_hash": class_hash.0.to_string(),
@@ -274,6 +253,12 @@ impl ReexecutionStateReader for TestStateReader {
             self.rpc_state_reader.send_rpc_request("starknet_getClass", params.clone())?,
         )
         .map_err(serde_err_to_state_err)?;
+        // Create a binding to avoid value being dropped.
+        let mut dumper_binding = self.contract_class_mapping_dumper.lock().unwrap();
+        // If dumper exists, insert the contract class to the mapping.
+        if let Some(contract_class_mapping_dumper) = dumper_binding.as_mut() {
+            contract_class_mapping_dumper.insert(*class_hash, contract_class.clone());
+        }
         Ok(contract_class)
     }
 }