refactor(native_blockifier): pybouncerinfo tx_weights field

crates/blockifier/src/transaction/transaction_utils.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use cairo_vm::vm::runners::builtin_runner::SEGMENT_ARENA_BUILTIN_NAME;
+use cairo_vm::vm::runners::cairo_runner::ExecutionResources as VmExecutionResources;
 use starknet_api::transaction::TransactionVersion;
 
 use super::objects::GasAndBlobGasUsages;
@@ -12,6 +13,7 @@ use crate::fee::os_usage::get_additional_os_resources;
 use crate::transaction::errors::TransactionExecutionError;
 use crate::transaction::objects::{ResourcesMapping, TransactionExecutionResult};
 use crate::transaction::transaction_types::TransactionType;
+use crate::utils::merge_hashmaps;
 
 /// Calculates the total resources needed to include the transaction in a Starknet block as
 /// most-recent (recent w.r.t. application on the given state).
@@ -80,3 +82,34 @@ pub fn verify_contract_class_version(
         }
     }
 }
+
+// TODO(Ayelet, 01/02/2024): Move to VmExecutionResourcesWrapper when merged.
+pub fn vm_execution_resources_to_hash_map(
+    execution_resources: VmExecutionResources,
+) -> HashMap<String, usize> {
+    let mut result = execution_resources.builtin_instance_counter.clone();
+    result.insert(
+        String::from("n_steps"),
+        execution_resources.n_steps + execution_resources.n_memory_holes,
+    );
+    result
+}
+
+// TODO(Arni, 24/01/2024): Add state_diff_size to tx_weights
+pub fn calculate_tx_weights(
+    additional_os_resources: VmExecutionResources,
+    actual_resources: HashMap<String, usize>,
+    message_segment_length: usize,
+) -> HashMap<String, usize> {
+    let mut tx_weights: HashMap<String, usize> = HashMap::new();
+    let mut cairo_resource_usage: HashMap<String, usize> = actual_resources;
+    if let Some(value) = cairo_resource_usage.remove("l1_gas_usage") {
+        tx_weights.insert("gas_weight".to_string(), value);
+    }
+    let os_cairo_usage: HashMap<String, usize> =
+        vm_execution_resources_to_hash_map(additional_os_resources);
+    let cairo_usage = merge_hashmaps(&cairo_resource_usage, &os_cairo_usage);
+    tx_weights.extend(cairo_usage);
+    tx_weights.insert("message_segment_length".to_string(), message_segment_length);
+    tx_weights
+}

crates/blockifier/src/utils.rs

@@ -1,4 +1,6 @@
 use std::collections::HashMap;
+use std::hash::Hash;
+use std::ops::{Add, AddAssign};
 
 #[cfg(test)]
 #[path = "utils_test.rs"]
@@ -19,3 +21,18 @@ where
 pub const fn const_max(a: u128, b: u128) -> u128 {
     [a, b][(a < b) as usize]
 }
+
+pub fn merge_hashmaps<K, V>(x: &HashMap<K, V>, y: &HashMap<K, V>) -> HashMap<K, V>
+where
+    K: Hash + Eq + Clone,
+    V: Add<Output = V> + AddAssign + Default + Clone,
+{
+    let mut result = x.clone();
+    for (key, value) in y {
+        result
+            .entry(key.clone())
+            .and_modify(|v| v.add_assign(value.clone()))
+            .or_insert(value.clone());
+    }
+    result
+}

crates/native_blockifier/src/py_transaction_execution_info.rs

@@ -181,6 +181,7 @@ impl From<VmExecutionResources> for PyVmExecutionResources {
 
 #[pyclass]
 #[derive(Clone, Default)]
+// TODO(Ayelet, 24/01/2024): Consider remove message_segment_length, state_diff_size.
 pub struct PyBouncerInfo {
     #[pyo3(get)]
     // The number of felts needed to store L1<>L2 messages.
@@ -189,5 +190,5 @@ pub struct PyBouncerInfo {
     // The number of felts needed to store the state diff.
     pub state_diff_size: usize,
     #[pyo3(get)]
-    pub additional_os_resources: PyVmExecutionResources,
+    pub tx_weights: HashMap<String, usize>,
 }

crates/native_blockifier/src/transaction_executor.rs

@@ -12,6 +12,7 @@ use blockifier::state::cached_state::{
 use blockifier::state::state_api::{State, StateReader};
 use blockifier::transaction::account_transaction::AccountTransaction;
 use blockifier::transaction::transaction_execution::Transaction;
+use blockifier::transaction::transaction_utils::calculate_tx_weights;
 use blockifier::transaction::transactions::{ExecutableTransaction, ValidatableTransaction};
 use cairo_vm::vm::runners::builtin_runner::HASH_BUILTIN_NAME;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources as VmExecutionResources;
@@ -23,9 +24,7 @@ use crate::errors::{NativeBlockifierError, NativeBlockifierResult};
 use crate::py_block_executor::{into_block_context, PyGeneralConfig};
 use crate::py_state_diff::{PyBlockInfo, PyStateDiff};
 use crate::py_transaction::py_tx;
-use crate::py_transaction_execution_info::{
-    PyBouncerInfo, PyTransactionExecutionInfo, PyVmExecutionResources,
-};
+use crate::py_transaction_execution_info::{PyBouncerInfo, PyTransactionExecutionInfo};
 use crate::py_utils::PyFelt;
 
 pub struct TransactionExecutor<S: StateReader> {
@@ -104,8 +103,6 @@ impl<S: StateReader> TransactionExecutor<S> {
                     MessageL1CostInfo::calculate(call_infos, Some(l1_handler_payload_size))?;
 
                 // TODO(Elin, 01/06/2024): consider moving Bouncer logic to a function.
-                let py_tx_execution_info = PyTransactionExecutionInfo::from(tx_execution_info);
-
                 let mut additional_os_resources = get_casm_hash_calculation_resources(
                     &mut transactional_state,
                     &self.executed_class_hashes,
@@ -115,15 +112,19 @@ impl<S: StateReader> TransactionExecutor<S> {
                     &self.visited_storage_entries,
                     &tx_visited_storage_entries,
                 )?;
-                let py_bouncer_info = PyBouncerInfo {
+                let state_diff_size = 0;
+                let actual_resources = tx_execution_info.actual_resources.0.clone();
+                let tx_weights = calculate_tx_weights(
+                    additional_os_resources,
+                    actual_resources,
                     message_segment_length,
-                    state_diff_size: 0,
-                    additional_os_resources: PyVmExecutionResources::from(additional_os_resources),
-                };
-
+                );
+                let py_bouncer_info =
+                    PyBouncerInfo { message_segment_length, state_diff_size, tx_weights };
                 self.staged_for_commit_state = Some(
                     transactional_state.stage(tx_executed_class_hashes, tx_visited_storage_entries),
                 );
+                let py_tx_execution_info = PyTransactionExecutionInfo::from(tx_execution_info);
                 Ok((py_tx_execution_info, py_bouncer_info))
             }
             Err(error) => {