feat(blockifier): add struct StateChangesCountForFee

crates/blockifier/src/fee/gas_usage_test.rs

@@ -70,12 +70,15 @@ fn starknet_resources() -> StarknetResources {
         .map(|call_info| call_info.with_some_class_hash())
         .collect();
     let execution_summary = CallInfo::summarize_many(call_infos.iter());
-    let state_resources = StateResources::new_for_testing(StateChangesCount {
-        n_storage_updates: 7,
-        n_class_hash_updates: 11,
-        n_compiled_class_hash_updates: 13,
-        n_modified_contracts: 17,
-    });
+    let state_resources = StateResources::new_for_testing(
+        StateChangesCount {
+            n_storage_updates: 7,
+            n_class_hash_updates: 11,
+            n_compiled_class_hash_updates: 13,
+            n_modified_contracts: 17,
+        },
+        19,
+    );
     StarknetResources::new(2_usize, 3_usize, 4_usize, state_resources, 6.into(), execution_summary)
 }
 

crates/blockifier/src/fee/receipt.rs

@@ -67,7 +67,12 @@ impl TransactionReceipt {
             calldata_length,
             signature_length,
             code_size,
-            StateResources::new(state_changes, sender_address, tx_context.fee_token_address()),
+            StateResources::new(
+                state_changes,
+                sender_address,
+                tx_context.fee_token_address(),
+                tx_context.block_context.versioned_constants.enable_stateful_compression,
+            ),
             l1_handler_payload_size,
             execution_summary_without_fee_transfer,
         );

crates/blockifier/src/fee/receipt_test.rs

@@ -119,7 +119,7 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(deploy_account_state_changes_count),
+        StateResources::new_for_testing(deploy_account_state_changes_count, 0),
         None,
         ExecutionSummary::default(),
     );
@@ -229,7 +229,7 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         0,
         0,
         0,
-        StateResources::new_for_testing(l2_to_l1_state_changes_count),
+        StateResources::new_for_testing(l2_to_l1_state_changes_count, 0),
         None,
         execution_summary.clone(),
     );
@@ -262,7 +262,7 @@ fn test_calculate_tx_gas_usage_basic<'a>(
 
     assert_eq!(l2_to_l1_messages_gas_usage_vector, manual_gas_computation);
 
-    // Any calculation with storage writings.t
+    // Any calculation with storage writings.
 
     let n_modified_contracts = 7;
     let n_storage_updates = 11;
@@ -276,7 +276,7 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         0,
         0,
         0,
-        StateResources::new_for_testing(storage_writes_state_changes_count),
+        StateResources::new_for_testing(storage_writes_state_changes_count, n_storage_updates / 2),
         None,
         ExecutionSummary::default(),
     );
@@ -305,7 +305,10 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         l1_handler_payload_size,
         signature_length,
         0,
-        StateResources::new_for_testing(combined_state_changes_count),
+        StateResources::new_for_testing(
+            combined_state_changes_count,
+            storage_writes_state_changes_count.n_storage_updates / 2,
+        ),
         Some(l1_handler_payload_size),
         execution_summary.clone(),
     );
@@ -380,7 +383,7 @@ fn test_calculate_tx_gas_usage(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(state_changes_count),
+        StateResources::new_for_testing(state_changes_count, 0),
         None,
         ExecutionSummary::default(),
     );
@@ -437,7 +440,7 @@ fn test_calculate_tx_gas_usage(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(state_changes_count),
+        StateResources::new_for_testing(state_changes_count, 0),
         None,
         // The transfer entrypoint emits an event - pass the call info to count its resources.
         execution_summary,

crates/blockifier/src/fee/resources.rs

@@ -13,7 +13,7 @@ use crate::fee::gas_usage::{
     get_message_segment_length,
     get_onchain_data_segment_length,
 };
-use crate::state::cached_state::{StateChanges, StateChangesCount};
+use crate::state::cached_state::{StateChanges, StateChangesCount, StateChangesCountForFee};
 use crate::transaction::errors::TransactionFeeError;
 use crate::utils::u64_from_usize;
 use crate::versioned_constants::{ArchivalDataGasCosts, VersionedConstants};
@@ -142,34 +142,47 @@ impl StarknetResources {
 #[cfg_attr(feature = "transaction_serde", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct StateResources {
-    state_changes_for_fee: StateChangesCount,
+    pub state_changes_for_fee: StateChangesCountForFee,
 }
 
 impl StateResources {
     pub fn new(
         state_changes: &StateChanges,
         sender_address: Option<ContractAddress>,
         fee_token_address: ContractAddress,
+        enable_stateful_compression: bool,
     ) -> Self {
         Self {
-            state_changes_for_fee: state_changes
-                .count_for_fee_charge(sender_address, fee_token_address),
+            state_changes_for_fee: state_changes.count_for_fee_charge(
+                sender_address,
+                fee_token_address,
+                enable_stateful_compression,
+            ),
         }
     }
 
     #[cfg(any(test, feature = "testing"))]
-    pub fn new_for_testing(state_changes_for_fee: StateChangesCount) -> Self {
-        Self { state_changes_for_fee }
+    pub fn new_for_testing(
+        state_changes_count: StateChangesCount,
+        n_allocated_keys: usize,
+    ) -> Self {
+        Self {
+            state_changes_for_fee: StateChangesCountForFee {
+                state_changes_count,
+                n_allocated_keys,
+            },
+        }
     }
 
     /// Returns the gas cost of the transaction's state changes.
     pub fn to_gas_vector(&self, use_kzg_da: bool) -> GasVector {
         // TODO(Nimrod, 29/3/2024): delete `get_da_gas_cost` and move it's logic here.
-        get_da_gas_cost(&self.state_changes_for_fee, use_kzg_da)
+        // TODO(Yoav): Add the cost of allocating keys.
+        get_da_gas_cost(&self.state_changes_for_fee.state_changes_count, use_kzg_da)
     }
 
     pub fn get_onchain_data_segment_length(&self) -> usize {
-        get_onchain_data_segment_length(&self.state_changes_for_fee)
+        get_onchain_data_segment_length(&self.state_changes_for_fee.state_changes_count)
     }
 }
 

crates/blockifier/src/state/cached_state.rs

@@ -700,7 +700,8 @@ impl StateChanges {
         &self,
         sender_address: Option<ContractAddress>,
         fee_token_address: ContractAddress,
-    ) -> StateChangesCount {
+        enable_stateful_compression: bool,
+    ) -> StateChangesCountForFee {
         let mut modified_contracts = self.state_maps.get_modified_contracts();
 
         // For account transactions, we need to compute the transaction fee before we can execute
@@ -720,11 +721,19 @@ impl StateChanges {
         // block.
         modified_contracts.remove(&fee_token_address);
 
-        StateChangesCount {
-            n_storage_updates,
-            n_class_hash_updates: self.state_maps.class_hashes.len(),
-            n_compiled_class_hash_updates: self.state_maps.compiled_class_hashes.len(),
-            n_modified_contracts: modified_contracts.len(),
+        StateChangesCountForFee {
+            state_changes_count: StateChangesCount {
+                n_storage_updates,
+                n_class_hash_updates: self.state_maps.class_hashes.len(),
+                n_compiled_class_hash_updates: self.state_maps.compiled_class_hashes.len(),
+                n_modified_contracts: modified_contracts.len(),
+            },
+            n_allocated_keys: if enable_stateful_compression {
+                // TODO: Set number of allocated keys.
+                0
+            } else {
+                0
+            },
         }
     }
 }
@@ -738,3 +747,11 @@ pub struct StateChangesCount {
     pub n_compiled_class_hash_updates: usize,
     pub n_modified_contracts: usize,
 }
+
+/// Holds the number of state changes for fee.
+#[cfg_attr(feature = "transaction_serde", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub struct StateChangesCountForFee {
+    pub state_changes_count: StateChangesCount,
+    pub n_allocated_keys: usize,
+}

crates/blockifier/src/state/cached_state_test.rs

@@ -323,25 +323,32 @@ fn create_state_changes_for_test<S: StateReader>(
         let sender_balance_key = get_fee_token_var_address(sender_address);
         state.set_storage_at(fee_token_address, sender_balance_key, felt!("0x1999")).unwrap();
     }
-
     state.get_actual_state_changes().unwrap()
 }
 
 #[rstest]
 fn test_from_state_changes_for_fee_charge(
     #[values(Some(contract_address!("0x102")), None)] sender_address: Option<ContractAddress>,
+    #[values(true, false)] enable_stateful_compression: bool,
 ) {
     let mut state: CachedState<DictStateReader> = CachedState::default();
     let fee_token_address = contract_address!("0x17");
     let state_changes =
         create_state_changes_for_test(&mut state, sender_address, fee_token_address);
-    let state_changes_count = state_changes.count_for_fee_charge(sender_address, fee_token_address);
-    let expected_state_changes_count = StateChangesCount {
+    let state_changes_count = state_changes.count_for_fee_charge(
+        sender_address,
+        fee_token_address,
+        enable_stateful_compression,
+    );
+    let expected_state_changes_count = StateChangesCountForFee {
         // 1 for storage update + 1 for sender balance update if sender is defined.
-        n_storage_updates: 1 + usize::from(sender_address.is_some()),
-        n_class_hash_updates: 1,
-        n_compiled_class_hash_updates: 1,
-        n_modified_contracts: 2,
+        state_changes_count: StateChangesCount {
+            n_storage_updates: 1 + usize::from(sender_address.is_some()),
+            n_class_hash_updates: 1,
+            n_compiled_class_hash_updates: 1,
+            n_modified_contracts: 2,
+        },
+        n_allocated_keys: 0,
     };
     assert_eq!(state_changes_count, expected_state_changes_count);
 }

crates/blockifier/src/transaction/account_transactions_test.rs

@@ -58,7 +58,7 @@ use crate::execution::entry_point::EntryPointExecutionContext;
 use crate::execution::syscalls::SyscallSelector;
 use crate::fee::fee_utils::{get_fee_by_gas_vector, get_sequencer_balance_keys};
 use crate::fee::gas_usage::estimate_minimal_gas_vector;
-use crate::state::cached_state::{StateChangesCount, TransactionalState};
+use crate::state::cached_state::{StateChangesCount, StateChangesCountForFee, TransactionalState};
 use crate::state::state_api::{State, StateReader};
 use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::declare::declare_tx;
@@ -1375,15 +1375,21 @@ fn test_count_actual_storage_changes(
         expected_sequencer_fee_update,
     ]);
 
-    let state_changes_count_1 =
-        state_changes_1.clone().count_for_fee_charge(Some(account_address), fee_token_address);
-    let expected_state_changes_count_1 = StateChangesCount {
-        // See expected storage updates.
-        n_storage_updates: 3,
-        // The contract address (storage update) and the account address (nonce update). Does not
-        // include the fee token address as a modified contract.
-        n_modified_contracts: 2,
-        ..Default::default()
+    let state_changes_count_1 = state_changes_1.clone().count_for_fee_charge(
+        Some(account_address),
+        fee_token_address,
+        block_context.versioned_constants.enable_stateful_compression,
+    );
+    let expected_state_changes_count_1 = StateChangesCountForFee {
+        state_changes_count: StateChangesCount {
+            // See expected storage updates.
+            n_storage_updates: 3,
+            // The contract address (storage update) and the account address (nonce update). Does
+            // not include the fee token address as a modified contract.
+            n_modified_contracts: 2,
+            ..Default::default()
+        },
+        n_allocated_keys: 0,
     };
 
     assert_eq!(expected_modified_contracts, state_changes_1.state_maps.get_modified_contracts());
@@ -1412,15 +1418,21 @@ fn test_count_actual_storage_changes(
     let expected_storage_updates_2 =
         HashMap::from([account_balance_storage_change, expected_sequencer_fee_update]);
 
-    let state_changes_count_2 =
-        state_changes_2.clone().count_for_fee_charge(Some(account_address), fee_token_address);
-    let expected_state_changes_count_2 = StateChangesCount {
-        // See expected storage updates.
-        n_storage_updates: 2,
-        // The account address (nonce update). Does not include the fee token address as a modified
-        // contract.
-        n_modified_contracts: 1,
-        ..Default::default()
+    let state_changes_count_2 = state_changes_2.clone().count_for_fee_charge(
+        Some(account_address),
+        fee_token_address,
+        block_context.versioned_constants.enable_stateful_compression,
+    );
+    let expected_state_changes_count_2 = StateChangesCountForFee {
+        state_changes_count: StateChangesCount {
+            // See expected storage updates.
+            n_storage_updates: 2,
+            // The account address (nonce update). Does not include the fee token address as a
+            // modified contract.
+            n_modified_contracts: 1,
+            ..Default::default()
+        },
+        n_allocated_keys: 0,
     };
 
     assert_eq!(expected_modified_contracts_2, state_changes_2.state_maps.get_modified_contracts());
@@ -1457,16 +1469,21 @@ fn test_count_actual_storage_changes(
         expected_sequencer_fee_update,
     ]);
 
-    let state_changes_count_3 = state_changes_transfer
-        .clone()
-        .count_for_fee_charge(Some(account_address), fee_token_address);
-    let expected_state_changes_count_3 = StateChangesCount {
-        // See expected storage updates.
-        n_storage_updates: 3,
-        // The account address (nonce update). Does not include the fee token address as a modified
-        // contract.
-        n_modified_contracts: 1,
-        ..Default::default()
+    let state_changes_count_3 = state_changes_transfer.clone().count_for_fee_charge(
+        Some(account_address),
+        fee_token_address,
+        block_context.versioned_constants.enable_stateful_compression,
+    );
+    let expected_state_changes_count_3 = StateChangesCountForFee {
+        state_changes_count: StateChangesCount {
+            // See expected storage updates.
+            n_storage_updates: 3,
+            // The account address (nonce update). Does not include the fee token address as a
+            // modified contract.
+            n_modified_contracts: 1,
+            ..Default::default()
+        },
+        n_allocated_keys: 0,
     };
 
     assert_eq!(

crates/blockifier/src/transaction/transactions_test.rs

@@ -470,7 +470,7 @@ fn test_invoke_tx(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(state_changes_for_fee),
+        StateResources::new_for_testing(state_changes_for_fee, 0),
         None,
         ExecutionSummary::default(),
     );
@@ -1504,7 +1504,7 @@ fn test_declare_tx(
         0,
         0,
         class_info.code_size(),
-        StateResources::new_for_testing(state_changes_for_fee),
+        StateResources::new_for_testing(state_changes_for_fee, 0),
         None,
         ExecutionSummary::default(),
     );