From 65e65506fa8be72593f45a4235819de8acc818e0 Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Thu, 15 Aug 2024 11:19:19 +0300
Subject: [PATCH] refactor: use snapi class info in process declare tx

---
 crates/gateway/src/compilation.rs               | 17 +++++------------
 crates/gateway/src/compilation_test.rs          |  6 ++----
 crates/gateway/src/gateway.rs                   | 10 +++++++---
 .../src/stateful_transaction_validator_test.rs  | 10 +++++++---
 4 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/crates/gateway/src/compilation.rs b/crates/gateway/src/compilation.rs
index fbd687e8026..950ba41dca2 100644
--- a/crates/gateway/src/compilation.rs
+++ b/crates/gateway/src/compilation.rs
@@ -1,8 +1,8 @@
 use std::sync::Arc;
 
-use blockifier::execution::contract_class::{ClassInfo, ContractClass, ContractClassV1};
 use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use cairo_lang_starknet_classes::contract_class::ContractClass as CairoLangContractClass;
+use starknet_api::contract_class::ClassInfo;
 use starknet_api::core::CompiledClassHash;
 use starknet_api::rpc_transaction::RpcDeclareTransaction;
 use starknet_sierra_compile::cairo_lang_compiler::CairoLangSierraToCasmCompiler;
@@ -43,17 +43,10 @@ impl GatewayCompiler {
 
         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);
-            GatewaySpecError::UnexpectedError { data: "Internal server error.".to_owned() }
+        Ok(ClassInfo {
+            casm_contract_class,
+            sierra_program_length: rpc_contract_class.sierra_program.len(),
+            abi_length: rpc_contract_class.abi.len(),
         })
     }
 
diff --git a/crates/gateway/src/compilation_test.rs b/crates/gateway/src/compilation_test.rs
index fd5018ee522..c2f770a9152 100644
--- a/crates/gateway/src/compilation_test.rs
+++ b/crates/gateway/src/compilation_test.rs
@@ -1,5 +1,4 @@
 use assert_matches::assert_matches;
-use blockifier::execution::contract_class::ContractClass;
 use cairo_lang_sierra_to_casm::compiler::CompilationError;
 use cairo_lang_starknet_classes::allowed_libfuncs::AllowedLibfuncsError;
 use cairo_lang_starknet_classes::casm_contract_class::StarknetSierraCompilationError;
@@ -103,7 +102,6 @@ fn test_process_declare_tx_success(
     let declare_tx = RpcDeclareTransaction::V3(declare_tx_v3);
 
     let class_info = gateway_compiler.process_declare_tx(&declare_tx).unwrap();
-    assert_matches!(class_info.contract_class(), ContractClass::V1(_));
-    assert_eq!(class_info.sierra_program_length(), sierra_program_length);
-    assert_eq!(class_info.abi_length(), abi_length);
+    assert_eq!(class_info.sierra_program_length, sierra_program_length);
+    assert_eq!(class_info.abi_length, abi_length);
 }
diff --git a/crates/gateway/src/gateway.rs b/crates/gateway/src/gateway.rs
index e54d61606b6..9bb1739e431 100644
--- a/crates/gateway/src/gateway.rs
+++ b/crates/gateway/src/gateway.rs
@@ -6,6 +6,7 @@ use async_trait::async_trait;
 use axum::extract::State;
 use axum::routing::{get, post};
 use axum::{Json, Router};
+use blockifier::execution::contract_class::ClassInfo;
 use starknet_api::executable_transaction::Transaction;
 use starknet_api::rpc_transaction::RpcTransaction;
 use starknet_api::transaction::TransactionHash;
@@ -130,9 +131,12 @@ fn process_tx(
 
     // Compile Sierra to Casm.
     let optional_class_info = match &tx {
-        RpcTransaction::Declare(declare_tx) => {
-            Some(gateway_compiler.process_declare_tx(declare_tx)?)
-        }
+        RpcTransaction::Declare(declare_tx) => Some(
+            ClassInfo::try_from(gateway_compiler.process_declare_tx(declare_tx)?).map_err(|e| {
+                error!("Failed to convert Starknet API ClassInfo to Blockifier ClassInfo: {:?}", e);
+                GatewaySpecError::UnexpectedError { data: "Internal server error.".to_owned() }
+            })?,
+        ),
         _ => None,
     };
 
diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs
index 4f9354e4c70..525285d7d88 100644
--- a/crates/gateway/src/stateful_transaction_validator_test.rs
+++ b/crates/gateway/src/stateful_transaction_validator_test.rs
@@ -3,6 +3,7 @@ use blockifier::blockifier::stateful_validator::{
     StatefulValidatorResult as BlockifierStatefulValidatorResult,
 };
 use blockifier::context::BlockContext;
+use blockifier::execution::contract_class::ClassInfo;
 use blockifier::test_utils::CairoVersion;
 use blockifier::transaction::errors::{TransactionFeeError, TransactionPreValidationError};
 use mempool_test_utils::invoke_tx_args;
@@ -83,9 +84,12 @@ fn test_stateful_tx_validator(
 ) {
     let optional_class_info = match &external_tx {
         RpcTransaction::Declare(declare_tx) => Some(
-            GatewayCompiler::new_cairo_lang_compiler(SierraToCasmCompilationConfig::default())
-                .process_declare_tx(declare_tx)
-                .unwrap(),
+            ClassInfo::try_from(
+                GatewayCompiler::new_cairo_lang_compiler(SierraToCasmCompilationConfig::default())
+                    .process_declare_tx(declare_tx)
+                    .unwrap(),
+            )
+            .unwrap(),
         ),
         _ => None,
     };