refactor: streamline class info object to match contract class

starkware-libs · Aug 14, 2024 · ad0e6f0 · ad0e6f0
1 parent 13ae058
commit ad0e6f0
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 98 deletions.
diff --git a/crates/blockifier/src/execution/contract_class.rs b/crates/blockifier/src/execution/contract_class.rs
@@ -34,7 +34,7 @@ use super::execution_utils::poseidon_hash_many_cost;
 use crate::abi::abi_utils::selector_from_name;
 use crate::abi::constants::{self, CONSTRUCTOR_ENTRY_POINT_NAME};
 use crate::execution::entry_point::CallEntryPoint;
-use crate::execution::errors::{ContractClassError, PreExecutionError};
+use crate::execution::errors::PreExecutionError;
 use crate::execution::execution_utils::sn_api_to_cairo_vm_program;
 use crate::fee::eth_gas_constants;
 use crate::transaction::errors::TransactionExecutionError;
@@ -48,8 +48,6 @@ pub mod test;
 // Note: when deserializing from a SN API class JSON string, the ABI field is ignored
 // by serde, since it is not required for execution.
 
-pub type ContractClassResult<T> = Result<T, ContractClassError>;
-
 #[derive(Clone, Debug, Eq, PartialEq, derive_more::From)]
 pub enum ContractClass {
     V0(ContractClassV0),
@@ -492,10 +490,9 @@ fn convert_entry_points_v1(external: Vec<CasmContractEntryPoint>) -> Vec<EntryPo
 
 #[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 {
@@ -509,25 +506,42 @@ impl TryFrom<starknet_api::contract_class::ClassInfo> for ClassInfo {
         } = class_info;
 
         let contract_class: ContractClass = contract_class.try_into()?;
-        Ok(Self { contract_class, sierra_program_length, abi_length })
+        Ok(match contract_class {
+            ContractClass::V0(contract_class) => ClassInfo::V0 { contract_class, abi_length },
+            ContractClass::V1(contract_class) => {
+                ClassInfo::V1 { contract_class, sierra_program_length, abi_length }
+            }
+        })
     }
 }
 
 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 {
@@ -536,24 +550,4 @@ impl ClassInfo {
             * eth_gas_constants::WORD_WIDTH
             + self.abi_length()
     }
-
-    pub fn new(
-        contract_class: &ContractClass,
-        sierra_program_length: usize,
-        abi_length: usize,
-    ) -> ContractClassResult<Self> {
-        let (contract_class_version, condition) = match contract_class {
-            ContractClass::V0(_) => (0, sierra_program_length == 0),
-            ContractClass::V1(_) => (1, sierra_program_length > 0),
-        };
-
-        if condition {
-            Ok(Self { contract_class: contract_class.clone(), sierra_program_length, abi_length })
-        } else {
-            Err(ContractClassError::ContractClassVersionSierraProgramLengthMismatch {
-                contract_class_version,
-                sierra_program_length,
-            })
-        }
-    }
 }
diff --git a/crates/blockifier/src/transaction/test_utils.rs b/crates/blockifier/src/transaction/test_utils.rs
@@ -300,11 +300,13 @@ pub fn l1_resource_bounds(max_amount: u64, max_price: u128) -> ResourceBoundsMap
 }
 
 pub fn calculate_class_info_for_testing(contract_class: ContractClass) -> ClassInfo {
-    let sierra_program_length = match contract_class {
-        ContractClass::V0(_) => 0,
-        ContractClass::V1(_) => 100,
-    };
-    ClassInfo::new(&contract_class, sierra_program_length, 100).unwrap()
+    let abi_length = 100;
+    match contract_class {
+        ContractClass::V0(contract_class) => ClassInfo::V0 { contract_class, abi_length },
+        ContractClass::V1(contract_class) => {
+            ClassInfo::V1 { contract_class, sierra_program_length: 100, abi_length }
+        }
+    }
 }
 
 pub fn emit_n_events_tx(

diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use blockifier::execution::contract_class::{ClassInfo, ContractClass, ContractClassV1};
+use blockifier::execution::contract_class::{ClassInfo, ContractClassV1};
 use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use cairo_lang_starknet_classes::contract_class::ContractClass as CairoLangContractClass;
 use starknet_api::core::CompiledClassHash;
@@ -42,18 +42,15 @@ impl GatewayCompiler {
         let casm_contract_class = self.compile(cairo_lang_contract_class)?;
 
         validate_compiled_class_hash(&casm_contract_class, &tx.compiled_class_hash)?;
-
-        ClassInfo::new(
-            &ContractClass::V1(ContractClassV1::try_from(casm_contract_class).map_err(|e| {
-                error!("Failed to convert CasmContractClass to Blockifier ContractClass: {:?}", e);
-                GatewaySpecError::UnexpectedError { data: "Internal server error.".to_owned() }
-            })?),
-            rpc_contract_class.sierra_program.len(),
-            rpc_contract_class.abi.len(),
-        )
-        .map_err(|e| {
-            error!("Failed to convert Blockifier ContractClass to Blockifier ClassInfo: {:?}", e);
+        let contract_class = ContractClassV1::try_from(casm_contract_class).map_err(|e| {
+            error!("Failed to convert CasmContractClass to Blockifier ContractClass: {:?}", e);
             GatewaySpecError::UnexpectedError { data: "Internal server error.".to_owned() }
+        })?;
+
+        Ok(ClassInfo::V1 {
+            contract_class,
+            sierra_program_length: rpc_contract_class.sierra_program.len(),
+            abi_length: rpc_contract_class.abi.len(),
         })
     }
 

diff --git a/crates/native_blockifier/src/py_transaction.rs b/crates/native_blockifier/src/py_transaction.rs
@@ -1,11 +1,6 @@
 use std::collections::BTreeMap;
 
-use blockifier::execution::contract_class::{
-    ClassInfo,
-    ContractClass,
-    ContractClassV0,
-    ContractClassV1,
-};
+use blockifier::execution::contract_class::{ClassInfo, ContractClassV0, ContractClassV1};
 use blockifier::transaction::account_transaction::AccountTransaction;
 use blockifier::transaction::transaction_execution::Transaction;
 use blockifier::transaction::transaction_types::TransactionType;
@@ -163,21 +158,24 @@ impl PyClassInfo {
         py_class_info: PyClassInfo,
         tx: &starknet_api::transaction::DeclareTransaction,
     ) -> NativeBlockifierResult<ClassInfo> {
-        let contract_class: ContractClass = match tx {
+        let class_info = match tx {
             starknet_api::transaction::DeclareTransaction::V0(_)
             | starknet_api::transaction::DeclareTransaction::V1(_) => {
-                ContractClassV0::try_from_json_string(&py_class_info.raw_contract_class)?.into()
+                let contract_class =
+                    ContractClassV0::try_from_json_string(&py_class_info.raw_contract_class)?;
+                ClassInfo::V0 { contract_class, abi_length: py_class_info.abi_length }
             }
             starknet_api::transaction::DeclareTransaction::V2(_)
             | starknet_api::transaction::DeclareTransaction::V3(_) => {
-                ContractClassV1::try_from_json_string(&py_class_info.raw_contract_class)?.into()
+                let contract_class =
+                    ContractClassV1::try_from_json_string(&py_class_info.raw_contract_class)?;
+                ClassInfo::V1 {
+                    contract_class,
+                    sierra_program_length: py_class_info.sierra_program_length,
+                    abi_length: py_class_info.abi_length,
+                }
             }
         };
-        let class_info = ClassInfo::new(
-            &contract_class,
-            py_class_info.sierra_program_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
@@ -748,19 +748,15 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v0 = BlockifierContractClass::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(&class_v0, DEPRECATED_CONTRACT_SIERRA_SIZE, abi_length)
-                .map_err(|err| ExecutionError::BadDeclareTransaction {
-                    tx: DeclareTransaction::V0(declare_tx.clone()),
-                    err,
-                })?;
+            )?;
+            let class_info = ClassInfo::V0 { contract_class, abi_length };
             BlockifierTransaction::from_api(
                 Transaction::Declare(DeclareTransaction::V0(declare_tx)),
                 tx_hash,
@@ -777,14 +773,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v0 = BlockifierContractClass::V0(
-                deprecated_class.try_into().map_err(BlockifierError::new)?,
-            );
-            let class_info = ClassInfo::new(&class_v0, DEPRECATED_CONTRACT_SIERRA_SIZE, abi_length)
-                .map_err(|err| ExecutionError::BadDeclareTransaction {
-                    tx: DeclareTransaction::V1(declare_tx.clone()),
-                    err,
-                })?;
+            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,
@@ -802,16 +792,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v1 = BlockifierContractClass::V1(
-                compiled_class.try_into().map_err(BlockifierError::new)?,
-            );
-            let class_info =
-                ClassInfo::new(&class_v1, sierra_program_length, abi_length).map_err(|err| {
-                    ExecutionError::BadDeclareTransaction {
-                        tx: DeclareTransaction::V2(declare_tx.clone()),
-                        err,
-                    }
-                })?;
+            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,
@@ -829,16 +811,8 @@ fn to_blockifier_tx(
             abi_length,
             only_query,
         ) => {
-            let class_v1 = BlockifierContractClass::V1(
-                compiled_class.try_into().map_err(BlockifierError::new)?,
-            );
-            let class_info =
-                ClassInfo::new(&class_v1, sierra_program_length, abi_length).map_err(|err| {
-                    ExecutionError::BadDeclareTransaction {
-                        tx: DeclareTransaction::V3(declare_tx.clone()),
-                        err,
-                    }
-                })?;
+            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,