feat(blockifier): compute allocation cost

crates/blockifier/src/fee/receipt.rs

@@ -101,10 +101,10 @@ impl TransactionReceipt {
             tx_context.tx_info.get_fee_by_gas_vector(&tx_context.block_context.block_info, gas)
         };
 
-        let da_gas = tx_resources
-            .starknet_resources
-            .state
-            .to_gas_vector(tx_context.block_context.block_info.use_kzg_da);
+        let da_gas = tx_resources.starknet_resources.state.to_gas_vector(
+            tx_context.block_context.block_info.use_kzg_da,
+            &tx_context.block_context.versioned_constants.allocation_cost,
+        );
 
         Self { resources: tx_resources, gas, da_gas, fee }
     }

crates/blockifier/src/fee/receipt_test.rs

@@ -136,7 +136,9 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         GasVectorComputationMode::All => GasVector::from_l2_gas(calldata_and_signature_gas_cost),
     };
     let manual_gas_vector = manual_starknet_gas_usage_vector
-        + deploy_account_tx_starknet_resources.state.to_gas_vector(use_kzg_da);
+        + deploy_account_tx_starknet_resources
+            .state
+            .to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost);
 
     let deploy_account_gas_usage_vector = deploy_account_tx_starknet_resources.to_gas_vector(
         &versioned_constants,
@@ -249,10 +251,18 @@ fn test_calculate_tx_gas_usage_basic<'a>(
             .unwrap();
     let manual_sharp_gas_usage = message_segment_length
         * eth_gas_constants::SHARP_GAS_PER_MEMORY_WORD
-        + usize_from_u64(l2_to_l1_starknet_resources.state.to_gas_vector(use_kzg_da).l1_gas.0)
-            .unwrap();
-    let manual_sharp_blob_gas_usage =
-        l2_to_l1_starknet_resources.state.to_gas_vector(use_kzg_da).l1_data_gas;
+        + usize_from_u64(
+            l2_to_l1_starknet_resources
+                .state
+                .to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost)
+                .l1_gas
+                .0,
+        )
+        .unwrap();
+    let manual_sharp_blob_gas_usage = l2_to_l1_starknet_resources
+        .state
+        .to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost)
+        .l1_data_gas;
     let manual_gas_computation = GasVector {
         l1_gas: u64_from_usize(manual_starknet_gas_usage + manual_sharp_gas_usage).into(),
         l1_data_gas: manual_sharp_blob_gas_usage,
@@ -288,7 +298,9 @@ fn test_calculate_tx_gas_usage_basic<'a>(
 
     // Manual calculation.
     // No L2 gas is used, so gas amount does not depend on gas vector computation mode.
-    let manual_gas_computation = storage_writes_starknet_resources.state.to_gas_vector(use_kzg_da);
+    let manual_gas_computation = storage_writes_starknet_resources
+        .state
+        .to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost);
 
     assert_eq!(manual_gas_computation, storage_writings_gas_usage_vector);
 
@@ -334,7 +346,10 @@ fn test_calculate_tx_gas_usage_basic<'a>(
         // the combined calculation got it once.
         + u64_from_usize(fee_balance_discount).into(),
         // Expected blob gas usage is from data availability only.
-        l1_data_gas: combined_cases_starknet_resources.state.to_gas_vector(use_kzg_da).l1_data_gas,
+        l1_data_gas: combined_cases_starknet_resources
+            .state
+            .to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost)
+            .l1_data_gas,
         l2_gas: l1_handler_gas_usage_vector.l2_gas,
     };
 
@@ -378,11 +393,12 @@ fn test_calculate_tx_gas_usage(
         n_modified_contracts,
         n_compiled_class_hash_updates: 0,
     };
+    let n_allocated_keys = 0; // This tx doesn't allocate the account balance.
     let starknet_resources = StarknetResources::new(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(state_changes_count, 0),
+        StateResources::new_for_testing(state_changes_count, n_allocated_keys),
         None,
         ExecutionSummary::default(),
     );
@@ -432,6 +448,7 @@ fn test_calculate_tx_gas_usage(
         n_modified_contracts,
         n_compiled_class_hash_updates: 0,
     };
+    let n_allocated_keys = 1; // Only for the recipient.
     let execution_call_info =
         &tx_execution_info.execute_call_info.expect("Execution call info should exist.");
     let execution_summary =
@@ -440,7 +457,7 @@ fn test_calculate_tx_gas_usage(
         calldata_length,
         signature_length,
         0,
-        StateResources::new_for_testing(state_changes_count, 0),
+        StateResources::new_for_testing(state_changes_count, n_allocated_keys),
         None,
         // The transfer entrypoint emits an event - pass the call info to count its resources.
         execution_summary,

crates/blockifier/src/fee/resources.rs

@@ -16,7 +16,7 @@ use crate::fee::gas_usage::{
 use crate::state::cached_state::{StateChanges, StateChangesCountForFee};
 use crate::transaction::errors::TransactionFeeError;
 use crate::utils::u64_from_usize;
-use crate::versioned_constants::{ArchivalDataGasCosts, VersionedConstants};
+use crate::versioned_constants::{AllocationCost, ArchivalDataGasCosts, VersionedConstants};
 
 pub type TransactionFeeResult<T> = Result<T, TransactionFeeError>;
 
@@ -139,7 +139,7 @@ impl StarknetResources {
     ) -> GasVector {
         [
             self.archival_data.to_gas_vector(versioned_constants, mode),
-            self.state.to_gas_vector(use_kzg_da),
+            self.state.to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost),
             self.messages.to_gas_vector(),
         ]
         .iter()
@@ -186,10 +186,24 @@ impl StateResources {
     }
 
     /// 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.
-        // 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 to_gas_vector(&self, use_kzg_da: bool, allocation_cost: &AllocationCost) -> GasVector {
+        let n_allocated_keys = self.state_changes_for_fee.n_allocated_keys;
+        let allocation_gas_vector = allocation_cost.get_cost(use_kzg_da);
+        let total_allocation_cost =
+            allocation_gas_vector.checked_mul(n_allocated_keys).unwrap_or_else(|| {
+                panic!(
+                    "State resources to gas vector overflowed: tried to multiply \
+                     {allocation_gas_vector:?} by {n_allocated_keys:?}",
+                )
+            });
+        let da_gas_cost =
+            get_da_gas_cost(&self.state_changes_for_fee.state_changes_count, use_kzg_da);
+        total_allocation_cost.checked_add(da_gas_cost).unwrap_or_else(|| {
+            panic!(
+                "State resources to gas vector overflowed: tried to add {total_allocation_cost:?} \
+                 to {da_gas_cost:?}",
+            )
+        })
     }
 
     pub fn get_onchain_data_segment_length(&self) -> usize {

crates/blockifier/src/transaction/execution_flavors_test.rs

@@ -51,6 +51,7 @@ use crate::transaction::test_utils::{
 use crate::transaction::transaction_types::TransactionType;
 use crate::transaction::transactions::ExecutableTransaction;
 use crate::utils::u64_from_usize;
+use crate::versioned_constants::AllocationCost;
 const VALIDATE_GAS_OVERHEAD: GasAmount = GasAmount(21);
 
 struct FlavorTestInitialState {
@@ -713,7 +714,8 @@ fn test_simulate_validate_charge_fee_post_execution(
     #[case] fee_type: FeeType,
     #[case] is_deprecated: bool,
 ) {
-    let block_context = BlockContext::create_for_account_testing();
+    let mut block_context = BlockContext::create_for_account_testing();
+    block_context.versioned_constants.allocation_cost = AllocationCost::ZERO;
     let gas_price = block_context.block_info.gas_prices.get_l1_gas_price_by_fee_type(&fee_type);
     let chain_info = &block_context.chain_info;
     let fee_token_address = chain_info.fee_token_address(&fee_type);

crates/blockifier/src/transaction/post_execution_test.rs

@@ -44,6 +44,7 @@ use crate::transaction::test_utils::{
     TestInitData,
 };
 use crate::transaction::transactions::ExecutableTransaction;
+use crate::versioned_constants::AllocationCost;
 
 fn init_data_by_version(chain_info: &ChainInfo, cairo_version: CairoVersion) -> TestInitData {
     let test_contract = FeatureContract::TestContract(cairo_version);
@@ -86,11 +87,12 @@ fn calldata_for_write_and_transfer(
 fn test_revert_on_overdraft(
     max_fee: Fee,
     default_all_resource_bounds: ValidResourceBounds,
-    block_context: BlockContext,
+    mut block_context: BlockContext,
     #[case] version: TransactionVersion,
     #[case] fee_type: FeeType,
     #[values(CairoVersion::Cairo0)] cairo_version: CairoVersion,
 ) {
+    block_context.versioned_constants.allocation_cost = AllocationCost::ZERO;
     let chain_info = &block_context.chain_info;
     let fee_token_address = chain_info.fee_token_addresses.get_by_fee_type(&fee_type);
     // An address to be written into to observe state changes.
@@ -267,6 +269,7 @@ fn test_revert_on_resource_overuse(
     #[values(CairoVersion::Cairo0)] cairo_version: CairoVersion,
 ) {
     block_context.block_info.use_kzg_da = true;
+    block_context.versioned_constants.allocation_cost = AllocationCost::ZERO;
     let gas_mode = resource_bounds.get_gas_vector_computation_mode();
     let fee_type = if version == TransactionVersion::THREE { FeeType::Strk } else { FeeType::Eth };
     let gas_prices = block_context.block_info.gas_prices.get_gas_prices_by_fee_type(&fee_type);

crates/blockifier/src/transaction/transactions_test.rs

@@ -151,7 +151,7 @@ use crate::transaction::test_utils::{
 };
 use crate::transaction::transaction_types::TransactionType;
 use crate::transaction::transactions::ExecutableTransaction;
-use crate::versioned_constants::VersionedConstants;
+use crate::versioned_constants::{AllocationCost, VersionedConstants};
 use crate::{
     check_tx_execution_error_for_custom_hint,
     check_tx_execution_error_for_invalid_scenario,
@@ -554,7 +554,8 @@ fn test_invoke_tx(
         FeatureContract::ERC20(CairoVersion::Cairo0).get_class_hash(),
     );
 
-    let da_gas = starknet_resources.state.to_gas_vector(use_kzg_da);
+    let da_gas =
+        starknet_resources.state.to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost);
 
     let expected_cairo_resources = get_expected_cairo_resources(
         versioned_constants,
@@ -1288,6 +1289,7 @@ fn test_actual_fee_gt_resource_bounds(
     #[values(CairoVersion::Cairo0, CairoVersion::Cairo1)] account_cairo_version: CairoVersion,
 ) {
     let block_context = &mut block_context;
+    block_context.versioned_constants.allocation_cost = AllocationCost::ZERO;
     block_context.block_info.use_kzg_da = true;
     let mut nonce_manager = NonceManager::default();
     let gas_mode = resource_bounds.get_gas_vector_computation_mode();
@@ -1565,7 +1567,8 @@ fn test_declare_tx(
         )
     };
 
-    let da_gas = starknet_resources.state.to_gas_vector(use_kzg_da);
+    let da_gas =
+        starknet_resources.state.to_gas_vector(use_kzg_da, &versioned_constants.allocation_cost);
     let expected_cairo_resources = get_expected_cairo_resources(
         versioned_constants,
         TransactionType::Declare,
@@ -2293,15 +2296,15 @@ fn test_l1_handler(#[values(false, true)] use_kzg_da: bool) {
     let expected_gas = match use_kzg_da {
         true => GasVector {
             l1_gas: 17988_u32.into(),
-            l1_data_gas: 128_u32.into(),
+            l1_data_gas: 160_u32.into(),
             l2_gas: 0_u32.into(),
         },
-        false => GasVector::from_l1_gas(19131_u32.into()),
+        false => GasVector::from_l1_gas(19682_u32.into()),
     };
 
     let expected_da_gas = match use_kzg_da {
-        true => GasVector::from_l1_data_gas(128_u32.into()),
-        false => GasVector::from_l1_gas(1652_u32.into()),
+        true => GasVector::from_l1_data_gas(160_u32.into()),
+        false => GasVector::from_l1_gas(2203_u32.into()),
     };
 
     let state_changes_count = StateChangesCount {