refactor(execution): get_tx_weights for the bouncer

crates/blockifier/src/abi/constants.rs

@@ -29,9 +29,13 @@ pub const CONSUMED_MSG_TO_L2_ENCODED_DATA_SIZE: usize =
     (L1_TO_L2_MSG_HEADER_SIZE + 1) - CONSUMED_MSG_TO_L2_N_TOPICS;
 
 // Transaction resource names.
-pub const L1_GAS_USAGE: &str = "l1_gas_usage";
+pub const L1_GAS_USAGE: &str = "gas_weight";
 pub const BLOB_GAS_USAGE: &str = "l1_blob_gas_usage";
 pub const N_STEPS_RESOURCE: &str = "n_steps";
+pub const N_EVENTS: &str = "n_events";
+pub const MESSAGE_SEGMENT_LENGTH: &str = "message_segment_length";
+pub const STATE_DIFF_SIZE: &str = "state_diff_size";
+pub const N_MEMORY_HOLES: &str = "n_memory_holes";
 
 // Casm hash calculation-related constants.
 pub const CAIRO0_ENTRY_POINT_STRUCT_SIZE: usize = 2;

crates/blockifier/src/blockifier/bouncer.rs

@@ -4,8 +4,7 @@ use cairo_vm::serde::deserialize_program::BuiltinName;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources as VmExecutionResources;
 
 use crate::abi::constants;
-use crate::transaction::objects::{GasVector, ResourcesMapping, TransactionExecutionResult};
-use crate::utils::usize_from_u128;
+use crate::transaction::objects::{ResourcesMapping, TransactionExecutionResult};
 
 #[derive(Clone, Default)]
 pub struct BouncerInfo {
@@ -19,15 +18,12 @@ pub struct BouncerInfo {
 impl BouncerInfo {
     pub fn calculate(
         tx_actual_resources: &ResourcesMapping,
-        tx_starknet_gas_usage: GasVector,
+        gas_weight: usize,
         tx_additional_os_resources: VmExecutionResources,
         message_segment_length: usize,
         state_diff_size: usize,
         n_events: usize,
     ) -> TransactionExecutionResult<Self> {
-        // TODO(Avi, 30/03/2024): Consider removing "l1_gas_usage" from actual resources.
-        let gas_weight = usize_from_u128(tx_starknet_gas_usage.l1_gas)
-            .expect("This conversion should not fail as the value is a converted usize.");
         // TODO(Ayelet, 04/02/2024): Consider defining a constant list.
         let builtin_ordered_list = [
             BuiltinName::output,

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -102,7 +102,7 @@ impl<S: StateReader> TransactionExecutor<S> {
                 let tx_execution_summary = tx_execution_info.summarize();
 
                 // Count message to L1 resources.
-                let (message_segment_length, starknet_gas_usage) =
+                let (message_segment_length, messages_gas_usage) =
                     calc_message_l1_resources(&tx_execution_info, l1_handler_payload_size)?;
 
                 // Count additional OS resources.
@@ -130,7 +130,7 @@ impl<S: StateReader> TransactionExecutor<S> {
                 // Finalize counting logic.
                 let bouncer_info = BouncerInfo::calculate(
                     &tx_execution_info.bouncer_resources,
-                    starknet_gas_usage,
+                    messages_gas_usage,
                     additional_os_resources,
                     message_segment_length,
                     state_diff_size,

crates/blockifier/src/bouncer.rs

@@ -1,15 +1,21 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::vec::IntoIter;
 
+use cairo_vm::serde::deserialize_program::BuiltinName;
+use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
 use serde::Deserialize;
 use starknet_api::core::ClassHash;
 
-use crate::blockifier::transaction_executor::TransactionExecutorResult;
+use crate::abi::constants;
+use crate::blockifier::transaction_executor::{
+    get_casm_hash_calculation_resources, get_particia_update_resources, TransactionExecutorResult,
+};
 use crate::execution::call_info::{CallInfo, MessageL1CostInfo};
-use crate::fee::gas_usage::get_starknet_gas_usage;
+use crate::fee::gas_usage::{get_onchain_data_segment_length, get_starknet_gas_usage};
 use crate::state::cached_state::{StateChangesKeys, StorageEntry, TransactionalState};
 use crate::state::state_api::StateReader;
-use crate::transaction::objects::{GasVector, TransactionExecutionInfo};
+use crate::transaction::objects::TransactionExecutionInfo;
+use crate::utils::usize_from_u128;
 
 #[cfg(test)]
 #[path = "bouncer_test.rs"]
@@ -29,7 +35,11 @@ macro_rules! impl_checked_sub {
     };
 }
 
-#[derive(Clone, Copy, Debug, Default, derive_more::Sub, Deserialize, PartialEq)]
+pub type HashMapWrapper = HashMap<String, usize>;
+
+#[derive(
+    Clone, Copy, Debug, Default, derive_more::Add, derive_more::Sub, Deserialize, PartialEq,
+)]
 /// Represents the execution resources counted throughout block creation.
 pub struct BouncerWeights {
     builtin_count: BuiltinCount,
@@ -40,6 +50,30 @@ pub struct BouncerWeights {
     state_diff_size: usize,
 }
 
+impl From<HashMapWrapper> for BouncerWeights {
+    fn from(mut data: HashMapWrapper) -> Self {
+        Self {
+            gas: data.remove(constants::L1_GAS_USAGE).unwrap_or_default(),
+            n_steps: data.remove(constants::N_STEPS_RESOURCE).unwrap_or_default()
+                + data.remove(constants::N_MEMORY_HOLES).unwrap_or_default(),
+            message_segment_length: data
+                .remove(constants::MESSAGE_SEGMENT_LENGTH)
+                .unwrap_or_default(),
+            state_diff_size: data.remove(constants::STATE_DIFF_SIZE).unwrap_or_default(),
+            n_events: data.remove(constants::N_EVENTS).unwrap_or_default(),
+            builtin_count: BuiltinCount::from(data),
+        }
+    }
+}
+
+impl From<ExecutionResources> for BouncerWeights {
+    fn from(val: ExecutionResources) -> Self {
+        let mut weights = BouncerWeights::from(val.builtin_instance_counter);
+        weights.n_steps = val.n_steps + val.n_memory_holes;
+        weights
+    }
+}
+
 impl BouncerWeights {
     impl_checked_sub!(
         builtin_count,
@@ -51,7 +85,9 @@ impl BouncerWeights {
     );
 }
 
-#[derive(Clone, Copy, Debug, Default, derive_more::Sub, Deserialize, PartialEq)]
+#[derive(
+    Clone, Copy, Debug, Default, derive_more::Add, derive_more::Sub, Deserialize, PartialEq,
+)]
 pub struct BuiltinCount {
     bitwise: usize,
     ecdsa: usize,
@@ -62,6 +98,26 @@ pub struct BuiltinCount {
     range_check: usize,
 }
 
+impl From<HashMapWrapper> for BuiltinCount {
+    fn from(mut data: HashMapWrapper) -> Self {
+        let builtin_count = Self {
+            bitwise: data.remove(BuiltinName::bitwise.name()).unwrap_or_default(),
+            ecdsa: data.remove(BuiltinName::ecdsa.name()).unwrap_or_default(),
+            ec_op: data.remove(BuiltinName::ec_op.name()).unwrap_or_default(),
+            keccak: data.remove(BuiltinName::keccak.name()).unwrap_or_default(),
+            pedersen: data.remove(BuiltinName::pedersen.name()).unwrap_or_default(),
+            poseidon: data.remove(BuiltinName::poseidon.name()).unwrap_or_default(),
+            range_check: data.remove(BuiltinName::range_check.name()).unwrap_or_default(),
+        };
+        assert!(
+            data.is_empty(),
+            "The following keys do not exist in BuiltinCount: {:?} ",
+            data.keys()
+        );
+        builtin_count
+    }
+}
+
 impl BuiltinCount {
     impl_checked_sub!(bitwise, ecdsa, ec_op, keccak, pedersen, poseidon, range_check);
 }
@@ -111,7 +167,56 @@ impl TransactionalBouncer {
         TransactionalBouncer { block: parent, transaction: Bouncer::new(capacity) }
     }
 
-    // TODO update function (in the next PR)
+    // TODO update function (in the next PRs)
+
+    pub fn get_tx_weights<S: StateReader>(
+        &mut self,
+        state: &mut TransactionalState<'_, S>,
+        tx_execution_info: &TransactionExecutionInfo,
+        l1_handler_payload_size: Option<usize>,
+    ) -> TransactionExecutorResult<BouncerWeights> {
+        let mut additional_os_resources = get_casm_hash_calculation_resources(
+            state,
+            &self.block.executed_class_hashes,
+            &self.transaction.executed_class_hashes,
+        )?;
+        additional_os_resources += &get_particia_update_resources(
+            &self.block.visited_storage_entries,
+            &self.transaction.visited_storage_entries,
+        )?;
+        let tx_execution_info_weights = Self::get_tx_execution_info_resources_weights(
+            tx_execution_info,
+            l1_handler_payload_size,
+        )?;
+
+        let mut tx_weights =
+            BouncerWeights::from(additional_os_resources) + tx_execution_info_weights;
+        tx_weights.state_diff_size =
+            get_onchain_data_segment_length(&self.transaction.state_changes_keys.count());
+        Ok(tx_weights)
+    }
+
+    pub fn get_tx_execution_info_resources_weights(
+        tx_execution_info: &TransactionExecutionInfo,
+        l1_handler_payload_size: Option<usize>,
+    ) -> TransactionExecutorResult<BouncerWeights> {
+        let mut execution_info_resources = tx_execution_info.bouncer_resources.0.clone();
+
+        // The blob gas is not being limited by the bouncer, thus we don't use it here.
+        // The gas is determined by the state diff size, which is limited by the bouncer.
+        execution_info_resources.remove(constants::BLOB_GAS_USAGE);
+
+        // TODO(Avi, 30/03/2024): Consider removing "l1_gas_usage" from actual resources.
+        // This value is not used, instead we use the value from calc_message_l1_resources() below.
+        execution_info_resources.remove(constants::L1_GAS_USAGE);
+
+        let mut weights = BouncerWeights::from(execution_info_resources);
+        (weights.message_segment_length, weights.gas) =
+            calc_message_l1_resources(tx_execution_info, l1_handler_payload_size)?;
+        // TODO: consider getting n_events from tx_execution_info.summarize() in the following PRs
+        weights.n_events = tx_execution_info.get_number_of_events();
+        Ok(weights)
+    }
 
     pub fn update_used_state_entries_sets<S: StateReader>(
         &mut self,
@@ -146,7 +251,7 @@ impl TransactionalBouncer {
 pub fn calc_message_l1_resources(
     tx_execution_info: &TransactionExecutionInfo,
     l1_handler_payload_size: Option<usize>,
-) -> TransactionExecutorResult<(usize, GasVector)> {
+) -> TransactionExecutorResult<(usize, usize)> {
     let call_infos: IntoIter<&CallInfo> =
         [&tx_execution_info.validate_call_info, &tx_execution_info.execute_call_info]
             .iter()
@@ -156,11 +261,14 @@ pub fn calc_message_l1_resources(
     let MessageL1CostInfo { l2_to_l1_payload_lengths, message_segment_length } =
         MessageL1CostInfo::calculate(call_infos, l1_handler_payload_size)?;
 
-    let starknet_gas_usage = get_starknet_gas_usage(
+    let tx_starknet_gas_usage = get_starknet_gas_usage(
         message_segment_length,
         &l2_to_l1_payload_lengths,
         l1_handler_payload_size,
     );
 
-    Ok((message_segment_length, starknet_gas_usage))
+    let gas_weight = usize_from_u128(tx_starknet_gas_usage.l1_gas)
+        .expect("This conversion should not fail as the value is a converted usize.");
+
+    Ok((message_segment_length, gas_weight))
 }

crates/blockifier/src/execution/call_info.rs

@@ -205,6 +205,11 @@ impl CallInfo {
         )
     }
 
+    /// Returns the sum of events in CallInfo and its inner_calls
+    pub fn get_number_of_events(&self) -> usize {
+        self.into_iter().map(|call_info| call_info.execution.events.len()).sum()
+    }
+
     pub fn summarize(&self) -> ExecutionSummary {
         let mut executed_class_hashes: HashSet<ClassHash> = HashSet::new();
         let mut visited_storage_entries: HashSet<StorageEntry> = HashSet::new();

crates/blockifier/src/transaction/objects.rs

@@ -221,6 +221,11 @@ impl TransactionExecutionInfo {
             .chain(self.fee_transfer_call_info.iter())
     }
 
+    /// Returns the number of events emitted in this transaction execution.
+    pub fn get_number_of_events(&self) -> usize {
+        self.non_optional_call_infos().map(|call_info| call_info.get_number_of_events()).sum()
+    }
+
     pub fn is_reverted(&self) -> bool {
         self.revert_error.is_some()
     }