test(concurrency): test update sequencer balance in storage (#1898)

crates/blockifier/src/concurrency/worker_logic.rs

@@ -23,12 +23,12 @@ use crate::transaction::objects::{TransactionExecutionInfo, TransactionExecution
 use crate::transaction::transaction_execution::Transaction;
 use crate::transaction::transactions::ExecutableTransaction;
 
-const EXECUTION_OUTPUTS_UNWRAP_ERROR: &str = "Execution task outputs should not be None.";
-
 #[cfg(test)]
 #[path = "worker_logic_test.rs"]
 pub mod test;
 
+const EXECUTION_OUTPUTS_UNWRAP_ERROR: &str = "Execution task outputs should not be None.";
+
 #[derive(Debug)]
 pub struct ExecutionTaskOutput {
     pub reads: StateMaps,
@@ -198,40 +198,40 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
             &mut execution_output.as_mut().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).result;
 
         let tx_context = Arc::new(self.block_context.to_tx_context(tx));
-        // Fix the sequencer balance.
-        // There is no need to fix the balance when the sequencer transfers fee to itself, since we
-        // use the sequential fee transfer in this case.
         if let Ok(tx_info) = result_tx_info.as_mut() {
             // TODO(Meshi, 01/06/2024): ask the bouncer if there is enough room for the transaction.
+            // Update the sequencer balance (in state + call info).
             if tx_context.tx_info.sender_address()
-                != self.block_context.block_info.sequencer_address
+                == self.block_context.block_info.sequencer_address
             {
-                // Update the sequencer balance in the storage.
-                let mut next_tx_versioned_state = self.state.pin_version(tx_index + 1);
-                let (sequencer_balance_value_low, sequencer_balance_value_high) =
-                    next_tx_versioned_state.get_fee_token_balance(
-                        tx_context.block_context.block_info.sequencer_address,
-                        tx_context.fee_token_address(),
-                    )?;
-                if let Some(fee_transfer_call_info) = tx_info.fee_transfer_call_info.as_mut() {
-                    // Fix the transfer call info.
-                    fill_sequencer_balance_reads(
-                        fee_transfer_call_info,
-                        sequencer_balance_value_low,
-                        sequencer_balance_value_high,
-                    );
-                }
-                add_fee_to_sequencer_balance(
+                // When the sequencer is the sender, we use the sequential (full) fee transfer.
+                return Ok(true);
+            }
+
+            let mut next_tx_versioned_state = self.state.pin_version(tx_index + 1);
+            let (sequencer_balance_value_low, sequencer_balance_value_high) =
+                next_tx_versioned_state.get_fee_token_balance(
+                    tx_context.block_context.block_info.sequencer_address,
                     tx_context.fee_token_address(),
-                    &mut tx_versioned_state,
-                    tx_info.actual_fee,
-                    &self.block_context,
+                )?;
+            if let Some(fee_transfer_call_info) = tx_info.fee_transfer_call_info.as_mut() {
+                // Fix the transfer call info.
+                fill_sequencer_balance_reads(
+                    fee_transfer_call_info,
                     sequencer_balance_value_low,
                     sequencer_balance_value_high,
                 );
-                // Changing the sequencer balance storage cell does not trigger (re-)validation of
-                // the next transactions.
             }
+            add_fee_to_sequencer_balance(
+                tx_context.fee_token_address(),
+                &mut tx_versioned_state,
+                tx_info.actual_fee,
+                &self.block_context,
+                sequencer_balance_value_low,
+                sequencer_balance_value_high,
+            );
+            // Changing the sequencer balance storage cell does not trigger (re-)validation of
+            // the next transactions.
         }
 
         Ok(true)

crates/blockifier/src/concurrency/worker_logic_test.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use num_bigint::BigUint;
 use starknet_api::core::{ContractAddress, PatriciaKey};
 use starknet_api::hash::{StarkFelt, StarkHash};
 use starknet_api::transaction::Fee;
@@ -10,14 +11,18 @@ use crate::abi::abi_utils::get_fee_token_var_address;
 use crate::abi::sierra_types::next_storage_key;
 use crate::concurrency::scheduler::{Task, TransactionStatus};
 use crate::concurrency::test_utils::safe_versioned_state_for_testing;
+use crate::concurrency::worker_logic::add_fee_to_sequencer_balance;
 use crate::context::BlockContext;
+use crate::execution::execution_utils::{felt_to_stark_felt, stark_felt_to_felt};
+use crate::fee::fee_utils::get_sequencer_balance_keys;
 use crate::state::cached_state::StateMaps;
 use crate::state::state_api::StateReader;
 use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::initial_test_state::test_state_reader;
 use crate::test_utils::{
     create_calldata, CairoVersion, NonceManager, BALANCE, MAX_FEE, TEST_ERC20_CONTRACT_ADDRESS,
 };
+use crate::transaction::objects::FeeType;
 use crate::transaction::test_utils::account_invoke_tx;
 use crate::transaction::transaction_execution::Transaction;
 use crate::{invoke_tx_args, nonce, storage_key};
@@ -295,3 +300,62 @@ fn test_worker_validate() {
     let next_task2 = worker_executor.validate(tx_index);
     assert_eq!(next_task2, Task::NoTask);
 }
+use cairo_felt::Felt252;
+use rstest::rstest;
+
+#[rstest]
+#[case::no_overflow(Fee(50_u128), stark_felt!(100_u128), StarkFelt::ZERO)]
+#[case::overflow(Fee(150_u128), stark_felt!(u128::max_value()), stark_felt!(5_u128))]
+#[case::overflow_edge_case(Fee(500_u128), stark_felt!(u128::max_value()), stark_felt!(u128::max_value()-1))]
+pub fn test_add_fee_to_sequencer_balance(
+    #[case] actual_fee: Fee,
+    #[case] sequencer_balance_low: StarkFelt,
+    #[case] sequencer_balance_high: StarkFelt,
+) {
+    let tx_index = 0;
+    let block_context = BlockContext::create_for_account_testing_with_concurrency_mode(true);
+    let account = FeatureContract::Empty(CairoVersion::Cairo1);
+    let safe_versioned_state = safe_versioned_state_for_testing(test_state_reader(
+        &block_context.chain_info,
+        0,
+        &[(account, 1)],
+    ));
+    let mut tx_versioned_state = safe_versioned_state.pin_version(tx_index);
+    let (sequencer_balance_key_low, sequencer_balance_key_high) =
+        get_sequencer_balance_keys(&block_context);
+
+    let fee_token_address = block_context.chain_info.fee_token_address(&FeeType::Strk);
+
+    add_fee_to_sequencer_balance(
+        fee_token_address,
+        &mut tx_versioned_state,
+        actual_fee,
+        &block_context,
+        sequencer_balance_low,
+        sequencer_balance_high,
+    );
+    let next_tx_versioned_state = safe_versioned_state.pin_version(tx_index + 1);
+
+    let new_sequencer_balance_value_low = next_tx_versioned_state
+        .get_storage_at(fee_token_address, sequencer_balance_key_low)
+        .unwrap();
+    let new_sequencer_balance_value_high = next_tx_versioned_state
+        .get_storage_at(fee_token_address, sequencer_balance_key_high)
+        .unwrap();
+    let expected_balance =
+        (stark_felt_to_felt(sequencer_balance_low) + Felt252::from(actual_fee.0)).to_biguint();
+
+    let mask_128_bit = (BigUint::from(1_u8) << 128) - 1_u8;
+    let expected_sequencer_balance_value_low = Felt252::from(&expected_balance & mask_128_bit);
+    let expected_sequencer_balance_value_high =
+        stark_felt_to_felt(sequencer_balance_high) + Felt252::from(&expected_balance >> 128);
+
+    assert_eq!(
+        new_sequencer_balance_value_low,
+        felt_to_stark_felt(&expected_sequencer_balance_value_low)
+    );
+    assert_eq!(
+        new_sequencer_balance_value_high,
+        felt_to_stark_felt(&expected_sequencer_balance_value_high)
+    );
+}