From 5e9d38c3d1a946f2f7e09dc64f8ba964534d1e26 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Est=C3=A9fano=20Bargas?= <estefano.bargas@fing.edu.uy>
Date: Wed, 9 Aug 2023 16:07:20 -0300
Subject: [PATCH] Added revert logic to declare

---
 src/transaction/declare.rs | 25 +++++++++++++++++++++----
 tests/internals.rs         | 14 ++++----------
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/src/transaction/declare.rs b/src/transaction/declare.rs
index abb20c0b7..f6f6fadcc 100644
--- a/src/transaction/declare.rs
+++ b/src/transaction/declare.rs
@@ -2,6 +2,7 @@ use crate::definitions::constants::QUERY_VERSION_BASE;
 use crate::execution::execution_entry_point::ExecutionResult;
 use crate::services::api::contract_classes::deprecated_contract_class::EntryPointType;
 use crate::state::cached_state::CachedState;
+use crate::state::StateDiff;
 use crate::{
     core::{
         contract_address::compute_deprecated_class_hash,
@@ -27,7 +28,7 @@ use crate::{
 use cairo_vm::felt::Felt252;
 use num_traits::Zero;
 
-use super::fee::charge_fee;
+use super::fee::{calculate_tx_fee, charge_fee};
 use super::{verify_version, Transaction};
 
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -268,7 +269,25 @@ impl Declare {
         block_context: &BlockContext,
     ) -> Result<TransactionExecutionInfo, TransactionError> {
         self.handle_nonce(state)?;
-        let mut tx_exec_info = self.apply(state, block_context)?;
+
+        let mut transactional_state = state.create_copy();
+        let mut tx_exec_info = self.apply(&mut transactional_state, block_context)?;
+
+        let actual_fee = calculate_tx_fee(
+            &tx_exec_info.actual_resources,
+            block_context.starknet_os_config.gas_price,
+            block_context,
+        )?;
+
+        if actual_fee <= self.max_fee {
+            state.apply_state_update(&StateDiff::from_cached_state(transactional_state)?)?;
+            state.set_contract_class(&self.class_hash, &self.contract_class)?;
+        } else {
+            tx_exec_info = tx_exec_info.to_revert_error(format!(
+                "Calculated fee ({}) exceeds max fee ({})",
+                actual_fee, self.max_fee
+            ));
+        }
 
         let mut tx_execution_context =
             self.get_execution_context(block_context.invoke_tx_max_n_steps);
@@ -281,8 +300,6 @@ impl Declare {
             self.skip_fee_transfer,
         )?;
 
-        state.set_contract_class(&self.class_hash, &self.contract_class)?;
-
         tx_exec_info.set_fee_info(actual_fee, fee_transfer_info);
 
         Ok(tx_exec_info)
diff --git a/tests/internals.rs b/tests/internals.rs
index 4328faed7..733ce0f55 100644
--- a/tests/internals.rs
+++ b/tests/internals.rs
@@ -1842,16 +1842,10 @@ fn test_state_for_declare_tx() {
     assert_eq!(
         state.cache(),
         &StateCache::new(
-            HashMap::from([
-                (
-                    TEST_ACCOUNT_CONTRACT_ADDRESS.clone(),
-                    felt_to_hash(&TEST_ACCOUNT_CONTRACT_CLASS_HASH)
-                ),
-                (
-                    TEST_ERC20_CONTRACT_ADDRESS.clone(),
-                    felt_to_hash(&TEST_ERC20_CONTRACT_CLASS_HASH)
-                )
-            ]),
+            HashMap::from([(
+                TEST_ERC20_CONTRACT_ADDRESS.clone(),
+                felt_to_hash(&TEST_ERC20_CONTRACT_CLASS_HASH)
+            )]),
             HashMap::new(),
             HashMap::from([(TEST_ACCOUNT_CONTRACT_ADDRESS.clone(), 0.into())]),
             HashMap::from([