From ad54542bb8429c55b2908c205dc0767b65951379 Mon Sep 17 00:00:00 2001
From: Yoni <78365039+Yoni-Starkware@users.noreply.github.com>
Date: Mon, 29 Jul 2024 09:04:24 +0300
Subject: [PATCH] refactor: use TransactionContext::get_fee_token_address
 (#161)

---
 .../blockifier/src/concurrency/worker_logic_test.rs   | 11 +++--------
 .../blockifier/src/transaction/account_transaction.rs |  7 +++----
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/crates/blockifier/src/concurrency/worker_logic_test.rs b/crates/blockifier/src/concurrency/worker_logic_test.rs
index c6e085ff0c..0962f914ee 100644
--- a/crates/blockifier/src/concurrency/worker_logic_test.rs
+++ b/crates/blockifier/src/concurrency/worker_logic_test.rs
@@ -70,20 +70,15 @@ fn verify_sequencer_balance_update<S: StateReader>(
     // We assume the balance is at most 2^128, so the "low" value is sufficient.
     expected_sequencer_balance_low: u128,
 ) {
-    let TransactionContext { block_context, tx_info } = tx_context;
     let tx_version_state = state.pin_version(tx_index);
     let (sequencer_balance_key_low, sequencer_balance_key_high) =
-        get_sequencer_balance_keys(block_context);
+        get_sequencer_balance_keys(&tx_context.block_context);
     for (expected_balance, storage_key) in [
         (felt!(expected_sequencer_balance_low), sequencer_balance_key_low),
         (Felt::ZERO, sequencer_balance_key_high),
     ] {
-        let actual_balance = tx_version_state
-            .get_storage_at(
-                block_context.chain_info.fee_token_address(&tx_info.fee_type()),
-                storage_key,
-            )
-            .unwrap();
+        let actual_balance =
+            tx_version_state.get_storage_at(tx_context.fee_token_address(), storage_key).unwrap();
         assert_eq!(expected_balance, actual_balance);
     }
 }
diff --git a/crates/blockifier/src/transaction/account_transaction.rs b/crates/blockifier/src/transaction/account_transaction.rs
index d98cf2ffda..14fa0ed929 100644
--- a/crates/blockifier/src/transaction/account_transaction.rs
+++ b/crates/blockifier/src/transaction/account_transaction.rs
@@ -352,7 +352,7 @@ impl AccountTransaction {
         let msb_amount = Felt::from(0_u8);
 
         let TransactionContext { block_context, tx_info } = tx_context.as_ref();
-        let storage_address = block_context.chain_info.fee_token_address(&tx_info.fee_type());
+        let storage_address = tx_context.fee_token_address();
         let fee_transfer_call = CallEntryPoint {
             class_hash: None,
             code_address: None,
@@ -388,10 +388,9 @@ impl AccountTransaction {
         tx_context: Arc<TransactionContext>,
         actual_fee: Fee,
     ) -> TransactionExecutionResult<CallInfo> {
-        let TransactionContext { block_context, tx_info } = tx_context.as_ref();
-        let fee_address = block_context.chain_info.fee_token_address(&tx_info.fee_type());
+        let fee_address = tx_context.fee_token_address();
         let (sequencer_balance_key_low, sequencer_balance_key_high) =
-            get_sequencer_balance_keys(block_context);
+            get_sequencer_balance_keys(&tx_context.block_context);
         let mut transfer_state = TransactionalState::create_transactional(state);
 
         // Set the initial sequencer balance to avoid tarnishing the read-set of the transaction.