Better TransactionType handling in native_blockifier (#1136)

- Remove `py_enum_name`, it was only used for `tx_type`.
- Push down `tx_type` extraction into `py_tx`, it's only used there and is
  already included inside the `tx` argument.
- Cast string `tx_type` into `TransactionType` as soon as possible.
- Unify tx parsing logic under `py_tx`
- Rename the arguments `py_tx` -> `tx` since `py_tx` shadows the function
  with the same name (also for consistency).
- Add `From` to `Transaction`, for less verbose initialization.
- Throw regular panic instead of `unimplemented!()` --- this was true in
  the past when we didn't implement all tx_types, but now that we do any
  value passed from Python that is not in the match is a bug.

Co-Authored-By: Gilad Chase <[email protected]>

crates/blockifier/src/transaction/errors.rs

@@ -92,3 +92,9 @@ pub enum TransactionPreValidationError {
     #[error(transparent)]
     TransactionFeeError(#[from] TransactionFeeError),
 }
+
+#[derive(Debug, Error)]
+pub enum ParseError {
+    #[error("Unsupported transaction type: {0}")]
+    UnknownTransactionType(String),
+}

crates/blockifier/src/transaction/test_utils.rs

@@ -203,7 +203,7 @@ pub fn create_account_tx_for_validate_test(
             });
             AccountTransaction::Invoke(invoke_tx)
         }
-        TransactionType::L1Handler => unimplemented!(),
+        _ => panic!("{tx_type:?} is not an account transaction."),
     }
 }
 

crates/blockifier/src/transaction/transaction_execution.rs

@@ -16,7 +16,7 @@ use crate::transaction::transactions::{
     InvokeTransaction, L1HandlerTransaction,
 };
 
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum Transaction {
     AccountTransaction(AccountTransaction),
     L1HandlerTransaction(L1HandlerTransaction),

crates/blockifier/src/transaction/transaction_types.rs

@@ -1,10 +1,28 @@
+use std::str::FromStr;
+
 use serde::Deserialize;
 use strum_macros::EnumIter;
 
+use crate::transaction::errors::ParseError;
+
 #[derive(Clone, Copy, Debug, Deserialize, EnumIter, Eq, Hash, PartialEq)]
 pub enum TransactionType {
     Declare,
     DeployAccount,
     InvokeFunction,
     L1Handler,
 }
+
+impl FromStr for TransactionType {
+    type Err = ParseError;
+
+    fn from_str(tx_type: &str) -> Result<Self, Self::Err> {
+        match tx_type {
+            "Declare" | "DECLARE" => Ok(TransactionType::Declare),
+            "DeployAccount" | "DEPLOY_ACCOUNT" => Ok(TransactionType::DeployAccount),
+            "InvokeFunction" | "INVOKE_FUNCTION" => Ok(TransactionType::InvokeFunction),
+            "L1Handler" | "L1_HANDLER" => Ok(TransactionType::L1Handler),
+            unknown_tx_type => Err(ParseError::UnknownTransactionType(unknown_tx_type.to_string())),
+        }
+    }
+}

crates/native_blockifier/src/errors.rs

@@ -1,5 +1,7 @@
 use blockifier::state::errors::StateError;
-use blockifier::transaction::errors::{TransactionExecutionError, TransactionPreValidationError};
+use blockifier::transaction::errors::{
+    ParseError, TransactionExecutionError, TransactionPreValidationError,
+};
 use blockifier::transaction::transaction_types::TransactionType;
 use cairo_vm::types::errors::program_errors::ProgramError;
 use pyo3::create_exception;
@@ -70,6 +72,8 @@ native_blockifier_errors!(
 
 #[derive(Debug, Error)]
 pub enum NativeBlockifierInputError {
+    #[error(transparent)]
+    ParseError(#[from] ParseError),
     #[error(transparent)]
     ProgramError(#[from] ProgramError),
     #[error(transparent)]

crates/native_blockifier/src/py_transaction.rs

@@ -4,12 +4,13 @@ use blockifier::fee::actual_cost::ActualCost;
 use blockifier::transaction::account_transaction::AccountTransaction;
 use blockifier::transaction::objects::ResourcesMapping;
 use blockifier::transaction::transaction_execution::Transaction;
+use blockifier::transaction::transaction_types::TransactionType;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use starknet_api::transaction::{Fee, Resource, ResourceBounds};
 use starknet_api::StarknetApiError;
 
-use crate::errors::NativeBlockifierResult;
+use crate::errors::{NativeBlockifierInputError, NativeBlockifierResult};
 use crate::py_declare::py_declare;
 use crate::py_deploy_account::py_deploy_account;
 use crate::py_invoke_function::py_invoke_function;
@@ -132,32 +133,36 @@ impl From<PyActualCost> for ActualCost {
     }
 }
 
-// Transactions creation.
+// Transaction creation.
 
 pub fn py_account_tx(
-    tx_type: &str,
-    py_tx: &PyAny,
+    tx: &PyAny,
     raw_contract_class: Option<&str>,
 ) -> NativeBlockifierResult<AccountTransaction> {
-    match tx_type {
-        "DECLARE" => {
+    let Transaction::AccountTransaction(account_tx) = py_tx(tx, raw_contract_class)? else {
+        panic!("Not an account transaction.");
+    };
+
+    Ok(account_tx)
+}
+
+pub fn py_tx(tx: &PyAny, raw_contract_class: Option<&str>) -> NativeBlockifierResult<Transaction> {
+    let tx_type: &str = tx.getattr("tx_type")?.getattr("name")?.extract()?;
+    let tx_type: TransactionType =
+        tx_type.parse().map_err(NativeBlockifierInputError::ParseError)?;
+
+    Ok(match tx_type {
+        TransactionType::Declare => {
             let raw_contract_class: &str = raw_contract_class
                 .expect("A contract class must be passed in a Declare transaction.");
-            Ok(AccountTransaction::Declare(py_declare(py_tx, raw_contract_class)?))
+            AccountTransaction::Declare(py_declare(tx, raw_contract_class)?).into()
         }
-        "DEPLOY_ACCOUNT" => Ok(AccountTransaction::DeployAccount(py_deploy_account(py_tx)?)),
-        "INVOKE_FUNCTION" => Ok(AccountTransaction::Invoke(py_invoke_function(py_tx)?)),
-        _ => unimplemented!(),
-    }
-}
-
-pub fn py_tx(
-    tx_type: &str,
-    py_tx: &PyAny,
-    raw_contract_class: Option<&str>,
-) -> NativeBlockifierResult<Transaction> {
-    if tx_type == "L1_HANDLER" {
-        return Ok(Transaction::L1HandlerTransaction(py_l1_handler(py_tx)?));
-    }
-    Ok(Transaction::AccountTransaction(py_account_tx(tx_type, py_tx, raw_contract_class)?))
+        TransactionType::DeployAccount => {
+            AccountTransaction::DeployAccount(py_deploy_account(tx)?).into()
+        }
+        TransactionType::InvokeFunction => {
+            AccountTransaction::Invoke(py_invoke_function(tx)?).into()
+        }
+        TransactionType::L1Handler => py_l1_handler(tx)?.into(),
+    })
 }

crates/native_blockifier/src/py_utils.rs

@@ -108,12 +108,3 @@ where
 {
     Ok(obj.getattr(attr)?.extract()?)
 }
-
-pub fn py_enum_name<T>(obj: &PyAny, attr: &str) -> NativeBlockifierResult<T>
-where
-    T: for<'a> FromPyObject<'a>,
-    T: Clone,
-    T: ToString,
-{
-    py_attr(obj.getattr(attr)?, "name")
-}

crates/native_blockifier/src/py_validator.rs

@@ -16,7 +16,7 @@ use crate::py_transaction::{py_account_tx, PyActualCost};
 use crate::py_transaction_execution_info::{
     PyCallInfo, PyTransactionExecutionInfo, PyVmExecutionResources,
 };
-use crate::py_utils::{py_enum_name, PyFelt};
+use crate::py_utils::PyFelt;
 use crate::state_readers::py_state_reader::PyStateReader;
 use crate::transaction_executor::TransactionExecutor;
 
@@ -105,8 +105,7 @@ impl PyValidator {
         remaining_gas: u64,
         raw_contract_class: Option<&str>,
     ) -> NativeBlockifierResult<(Option<PyCallInfo>, PyActualCost)> {
-        let tx_type: String = py_enum_name(tx, "tx_type")?;
-        let account_tx = py_account_tx(&tx_type, tx, raw_contract_class)?;
+        let account_tx = py_account_tx(tx, raw_contract_class)?;
         let (optional_call_info, actual_cost) =
             self.tx_executor().validate(&account_tx, remaining_gas)?;
         let py_optional_call_info = optional_call_info.map(PyCallInfo::from);
@@ -126,8 +125,7 @@ impl PyValidator {
         raw_contract_class: Option<&str>,
         deploy_account_tx_hash: Option<PyFelt>,
     ) -> NativeBlockifierResult<()> {
-        let tx_type: String = py_enum_name(tx, "tx_type")?;
-        let account_tx = py_account_tx(&tx_type, tx, raw_contract_class)?;
+        let account_tx = py_account_tx(tx, raw_contract_class)?;
         let account_tx_context = account_tx.get_account_tx_context();
         // Deploy account transactions should be fully executed, since the constructor must run
         // before `__validate_deploy__`. The execution already includes all necessary validations,

crates/native_blockifier/src/transaction_executor.rs

@@ -22,7 +22,7 @@ use crate::py_block_executor::{into_block_context, PyGeneralConfig};
 use crate::py_state_diff::{PyBlockInfo, PyStateDiff};
 use crate::py_transaction::py_tx;
 use crate::py_transaction_execution_info::{PyTransactionExecutionInfo, PyVmExecutionResources};
-use crate::py_utils::{py_enum_name, PyFelt};
+use crate::py_utils::PyFelt;
 
 pub struct TransactionExecutor<S: StateReader> {
     pub block_context: BlockContext,
@@ -65,8 +65,7 @@ impl<S: StateReader> TransactionExecutor<S> {
         raw_contract_class: Option<&str>,
         charge_fee: bool,
     ) -> NativeBlockifierResult<(PyTransactionExecutionInfo, PyVmExecutionResources)> {
-        let tx_type: String = py_enum_name(tx, "tx_type")?;
-        let tx: Transaction = py_tx(&tx_type, tx, raw_contract_class)?;
+        let tx: Transaction = py_tx(tx, raw_contract_class)?;
 
         let mut tx_executed_class_hashes = HashSet::<ClassHash>::new();
         let mut transactional_state = CachedState::create_transactional(&mut self.state);