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, Eq, PartialEq, Debug)]
 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/bouncer.rs

@@ -1,14 +1,21 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
+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::{ExecutionSummary, MessageL1CostInfo};
-use crate::fee::gas_usage::{get_message_segment_length, get_messages_gas_usage};
+use crate::fee::gas_usage::{
+    get_message_segment_length, get_messages_gas_usage, get_onchain_data_segment_length,
+};
 use crate::state::cached_state::{StateChangesKeys, StorageEntry, TransactionalState};
 use crate::state::state_api::StateReader;
-use crate::transaction::objects::GasVector;
+use crate::transaction::objects::ResourcesMapping;
 
 #[cfg(test)]
 #[path = "bouncer_test.rs"]
@@ -28,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,
@@ -39,6 +50,16 @@ pub struct BouncerWeights {
     state_diff_size: usize,
 }
 
+impl From<ExecutionResources> for BouncerWeights {
+    fn from(data: ExecutionResources) -> Self {
+        BouncerWeights {
+            n_steps: data.n_steps + data.n_memory_holes,
+            builtin_count: data.builtin_instance_counter.into(),
+            ..Default::default()
+        }
+    }
+}
+
 impl BouncerWeights {
     impl_checked_sub!(
         builtin_count,
@@ -50,7 +71,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,
@@ -61,6 +84,28 @@ 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()).expect("bitwise must be present"),
+            ecdsa: data.remove(BuiltinName::ecdsa.name()).expect("ecdsa must be present"),
+            ec_op: data.remove(BuiltinName::ec_op.name()).expect("ec_op must be present"),
+            keccak: data.remove(BuiltinName::keccak.name()).expect("keccak must be present"),
+            pedersen: data.remove(BuiltinName::pedersen.name()).expect("pedersen must be present"),
+            poseidon: data.remove(BuiltinName::poseidon.name()).expect("poseidon must be present"),
+            range_check: data
+                .remove(BuiltinName::range_check.name())
+                .expect("range_check must be present"),
+        };
+        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);
 }
@@ -110,7 +155,73 @@ impl TransactionalBouncer {
         TransactionalBouncer { bouncer: parent, transactional: 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_summary: &ExecutionSummary,
+        bouncer_resources: &ResourcesMapping,
+        l1_handler_payload_size: Option<usize>,
+    ) -> TransactionExecutorResult<BouncerWeights> {
+        let mut additional_os_resources = get_casm_hash_calculation_resources(
+            state,
+            &self.bouncer.executed_class_hashes,
+            &self.transactional.executed_class_hashes,
+        )?;
+        additional_os_resources += &get_particia_update_resources(
+            &self.bouncer.visited_storage_entries,
+            &self.transactional.visited_storage_entries,
+        )?;
+
+        let execution_info_weights = Self::get_tx_execution_info_resources_weights(
+            tx_execution_summary,
+            bouncer_resources,
+            l1_handler_payload_size,
+        )?;
+
+        let mut tx_weights = BouncerWeights::from(additional_os_resources) + execution_info_weights;
+        tx_weights.state_diff_size =
+            get_onchain_data_segment_length(&self.transactional.state_changes_keys.count());
+        Ok(tx_weights)
+    }
+
+    pub fn get_tx_execution_info_resources_weights(
+        tx_execution_summary: &ExecutionSummary,
+        bouncer_resources: &ResourcesMapping,
+        l1_handler_payload_size: Option<usize>,
+    ) -> TransactionExecutorResult<BouncerWeights> {
+        let mut execution_info_resources = bouncer_resources.0.clone();
+
+        let (message_segment_length, gas_usage) = calculate_message_l1_resources(
+            &tx_execution_summary.l2_to_l1_payload_lengths,
+            l1_handler_payload_size,
+        );
+
+        // 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 weights = BouncerWeights {
+            gas: gas_usage,
+            message_segment_length,
+            n_events: tx_execution_summary.n_events,
+            n_steps: execution_info_resources
+                .remove(constants::N_STEPS_RESOURCE)
+                .expect("n_steps must be present in the execution info")
+                + execution_info_resources
+                    .remove(constants::N_MEMORY_HOLES)
+                    .expect("n_memory_hols must be present in the execution info"),
+            builtin_count: BuiltinCount::from(execution_info_resources),
+            state_diff_size: 0,
+        };
+
+        Ok(weights)
+    }
 
     pub fn update_auxiliary_info<S: StateReader>(
         &mut self,
@@ -144,7 +255,7 @@ impl TransactionalBouncer {
 pub fn calculate_message_l1_resources(
     l2_to_l1_payload_lengths: &[usize],
     l1_handler_payload_size: Option<usize>,
-) -> (usize, GasVector) {
+) -> (usize, usize) {
     let message_segment_length =
         get_message_segment_length(l2_to_l1_payload_lengths, l1_handler_payload_size);
     let gas_usage = get_messages_gas_usage(
@@ -154,5 +265,7 @@ pub fn calculate_message_l1_resources(
         },
         l1_handler_payload_size,
     );
+    let gas_usage: usize =
+        gas_usage.l1_gas.try_into().expect("L1 gas should be smaller than usize");
     (message_segment_length, gas_usage)
 }