refactor: streamline class info object to match contract class

starkware-libs · Oct 7, 2024 · 5120579 · 5120579
1 parent 85481c0
commit 5120579
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 46 deletions.
diff --git a/crates/blockifier/src/execution/contract_class.rs b/crates/blockifier/src/execution/contract_class.rs
@@ -561,10 +561,9 @@ fn convert_entry_points_v1(external: &[CasmContractEntryPoint]) -> Vec<EntryPoin
 
 #[derive(Clone, Debug)]
 // TODO(Ayelet,10/02/2024): Change to bytes.
-pub struct ClassInfo {
-    contract_class: ContractClass,
-    sierra_program_length: usize,
-    abi_length: usize,
+pub enum ClassInfo {
+    V0 { contract_class: ContractClassV0, abi_length: usize },
+    V1 { contract_class: ContractClassV1, sierra_program_length: usize, abi_length: usize },
 }
 
 impl TryFrom<&starknet_api::contract_class::ClassInfo> for ClassInfo {
@@ -577,7 +576,7 @@ impl TryFrom<&starknet_api::contract_class::ClassInfo> for ClassInfo {
             abi_length,
         } = class_info;
 
-        Ok(Self {
+        Ok(Self::V1 {
             contract_class: casm_contract_class.try_into()?,
             sierra_program_length: *sierra_program_length,
             abi_length: *abi_length,
@@ -595,7 +594,7 @@ impl TryFrom<starknet_api::contract_class::ClassInfo> for ClassInfo {
             abi_length,
         } = class_info;
 
-        Ok(Self {
+        Ok(Self::V1 {
             contract_class: casm_contract_class.try_into()?,
             sierra_program_length,
             abi_length,
@@ -605,19 +604,31 @@ impl TryFrom<starknet_api::contract_class::ClassInfo> for ClassInfo {
 
 impl ClassInfo {
     pub fn bytecode_length(&self) -> usize {
-        self.contract_class.bytecode_length()
+        match self {
+            ClassInfo::V0 { contract_class, .. } => contract_class.bytecode_length(),
+            ClassInfo::V1 { contract_class, .. } => contract_class.bytecode_length(),
+        }
     }
 
     pub fn contract_class(&self) -> ContractClass {
-        self.contract_class.clone()
+        match self {
+            ClassInfo::V0 { contract_class, .. } => ContractClass::V0(contract_class.clone()),
+            ClassInfo::V1 { contract_class, .. } => ContractClass::V1(contract_class.clone()),
+        }
     }
 
     pub fn sierra_program_length(&self) -> usize {
-        self.sierra_program_length
+        match self {
+            ClassInfo::V0 { .. } => 0,
+            ClassInfo::V1 { sierra_program_length, .. } => *sierra_program_length,
+        }
     }
 
     pub fn abi_length(&self) -> usize {
-        self.abi_length
+        match self {
+            ClassInfo::V0 { abi_length, .. } => *abi_length,
+            ClassInfo::V1 { abi_length, .. } => *abi_length,
+        }
     }
 
     pub fn code_size(&self) -> usize {
@@ -626,24 +637,4 @@ impl ClassInfo {
             * eth_gas_constants::WORD_WIDTH
             + self.abi_length()
     }
-
-    pub fn new_v1(
-        contract_class: ContractClassV1,
-        sierra_program_length: usize,
-        abi_length: usize,
-    ) -> Self {
-        Self {
-            contract_class: ContractClass::V1(contract_class),
-            sierra_program_length,
-            abi_length,
-        }
-    }
-
-    pub fn new_v0(contract_class: ContractClassV0, abi_length: usize) -> Self {
-        Self {
-            contract_class: ContractClass::V0(contract_class),
-            sierra_program_length: 0,
-            abi_length,
-        }
-    }
 }
diff --git a/crates/blockifier/src/transaction/test_utils.rs b/crates/blockifier/src/transaction/test_utils.rs
@@ -369,10 +369,9 @@ fn create_all_resource_bounds(
 pub fn calculate_class_info_for_testing(contract_class: ContractClass) -> ClassInfo {
     let abi_length = 100;
     match contract_class {
-        ContractClass::V0(contract_class) => ClassInfo::new_v0(contract_class, abi_length),
+        ContractClass::V0(contract_class) => ClassInfo::V0 { contract_class, abi_length },
         ContractClass::V1(contract_class) => {
-            let sierra_program_length = 100;
-            ClassInfo::new_v1(contract_class, sierra_program_length, abi_length)
+            ClassInfo::V1 { contract_class, sierra_program_length: 100, abi_length }
         }
     }
 }

diff --git a/crates/native_blockifier/src/py_transaction.rs b/crates/native_blockifier/src/py_transaction.rs
@@ -168,17 +168,17 @@ impl PyClassInfo {
             | starknet_api::transaction::DeclareTransaction::V1(_) => {
                 let contract_class =
                     ContractClassV0::try_from_json_string(&py_class_info.raw_contract_class)?;
-                ClassInfo::new_v0(contract_class, py_class_info.abi_length)
+                ClassInfo::V0 { contract_class, abi_length: py_class_info.abi_length }
             }
             starknet_api::transaction::DeclareTransaction::V2(_)
             | starknet_api::transaction::DeclareTransaction::V3(_) => {
                 let contract_class =
                     ContractClassV1::try_from_json_string(&py_class_info.raw_contract_class)?;
-                ClassInfo::new_v1(
+                ClassInfo::V1 {
                     contract_class,
-                    py_class_info.sierra_program_length,
-                    py_class_info.abi_length,
-                )
+                    sierra_program_length: py_class_info.sierra_program_length,
+                    abi_length: py_class_info.abi_length,
+                }
             }
         };
         Ok(class_info)

diff --git a/crates/papyrus_execution/src/lib.rs b/crates/papyrus_execution/src/lib.rs
@@ -737,15 +737,15 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v0 = deprecated_class.try_into().map_err(
+            let contract_class = deprecated_class.try_into().map_err(
                 |e: cairo_vm::types::errors::program_errors::ProgramError| {
                     ExecutionError::TransactionExecutionError {
                         transaction_index,
                         execution_error: e.to_string(),
                     }
                 },
             )?;
-            let class_info = ClassInfo::new_v0(class_v0, abi_length);
+            let class_info = ClassInfo::V0 { contract_class, abi_length };
             BlockifierTransaction::from_api(
                 Transaction::Declare(DeclareTransaction::V0(declare_tx)),
                 tx_hash,
@@ -762,8 +762,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v0 = deprecated_class.try_into().map_err(BlockifierError::new)?;
-            let class_info = ClassInfo::new_v0(class_v0, abi_length);
+            let contract_class = deprecated_class.try_into().map_err(BlockifierError::new)?;
+            let class_info = ClassInfo::V0 { contract_class, abi_length };
             BlockifierTransaction::from_api(
                 Transaction::Declare(DeclareTransaction::V1(declare_tx)),
                 tx_hash,
@@ -781,8 +781,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v1 = compiled_class.try_into().map_err(BlockifierError::new)?;
-            let class_info = ClassInfo::new_v1(class_v1, sierra_program_length, abi_length);
+            let contract_class = compiled_class.try_into().map_err(BlockifierError::new)?;
+            let class_info = ClassInfo::V1 { contract_class, sierra_program_length, abi_length };
             BlockifierTransaction::from_api(
                 Transaction::Declare(DeclareTransaction::V2(declare_tx)),
                 tx_hash,
@@ -800,8 +800,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v1 = compiled_class.try_into().map_err(BlockifierError::new)?;
-            let class_info = ClassInfo::new_v1(class_v1, sierra_program_length, abi_length);
+            let contract_class = compiled_class.try_into().map_err(BlockifierError::new)?;
+            let class_info = ClassInfo::V1 { contract_class, sierra_program_length, abi_length };
             BlockifierTransaction::from_api(
                 Transaction::Declare(DeclareTransaction::V3(declare_tx)),
                 tx_hash,