From 46fcafe0a5993036516107e210b3c75a24f0357f Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <ayelet@starkware.co>
Date: Thu, 18 Jan 2024 11:59:10 +0200
Subject: [PATCH] Replace PyBouncerInfo's additional_os_resources fiels to
 tx_weights field and _calculate_tx_weights in Rust instead of in Python

---
 .../src/transaction/transaction_utils.rs      | 46 +++++++++++++++++++
 .../src/py_transaction_execution_info.rs      |  2 +
 .../src/transaction_executor.rs               | 11 ++++-
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/crates/blockifier/src/transaction/transaction_utils.rs b/crates/blockifier/src/transaction/transaction_utils.rs
index c0ebea9cb4..ec26e934ec 100644
--- a/crates/blockifier/src/transaction/transaction_utils.rs
+++ b/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 crate::abi::constants;
@@ -73,3 +74,48 @@ pub fn verify_contract_class_version(
         }
     }
 }
+
+// TODO(Ayelet, 01/02/2024): Move to VmExecutionResourcesWrapper when merged.
+pub fn to_dict(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
+}
+
+fn add_counters(x: HashMap<String, usize>, y: HashMap<String, usize>) -> HashMap<String, usize> {
+    let mut result = HashMap::new();
+
+    // Insert values from x into result
+    for (key, value) in x {
+        result.insert(key.clone(), value);
+    }
+
+    // Add values from y to result
+    for (key, value) in y {
+        let entry = result.entry(key.clone()).or_insert(0);
+        *entry += value;
+    }
+
+    result
+}
+
+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.get("l1_gas_usage") {
+        tx_weights.insert("gas_weight".to_string(), value);
+        cairo_resource_usage.remove("l1_gas_usage");
+    }
+    let os_cairo_usage: HashMap<String, usize> = to_dict(additional_os_resources);
+    let cairo_usage = add_counters(cairo_resource_usage, os_cairo_usage);
+    tx_weights.extend(cairo_usage);
+    tx_weights.insert("message_segment_length".to_string(), message_segment_length);
+    tx_weights
+}
diff --git a/crates/native_blockifier/src/py_transaction_execution_info.rs b/crates/native_blockifier/src/py_transaction_execution_info.rs
index c2e52e068b..f4e89237cd 100644
--- a/crates/native_blockifier/src/py_transaction_execution_info.rs
+++ b/crates/native_blockifier/src/py_transaction_execution_info.rs
@@ -190,4 +190,6 @@ pub struct PyBouncerInfo {
     pub state_diff_size: usize,
     #[pyo3(get)]
     pub additional_os_resources: PyVmExecutionResources,
+    #[pyo3(get)]
+    pub tx_weights: HashMap<String, usize>,
 }
diff --git a/crates/native_blockifier/src/transaction_executor.rs b/crates/native_blockifier/src/transaction_executor.rs
index 41b6950e41..1f0107217f 100644
--- a/crates/native_blockifier/src/transaction_executor.rs
+++ b/crates/native_blockifier/src/transaction_executor.rs
@@ -11,6 +11,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;
@@ -85,6 +86,7 @@ impl<S: StateReader> TransactionExecutor<S> {
             .map_err(NativeBlockifierError::from);
         match tx_execution_result {
             Ok(tx_execution_info) => {
+                let actual_resources = tx_execution_info.actual_resources.0.clone();
                 // TODO(Elin, 01/06/2024): consider traversing the calls to collect data once.
                 tx_executed_class_hashes.extend(tx_execution_info.get_executed_class_hashes());
                 tx_visited_storage_entries.extend(tx_execution_info.get_visited_storage_entries());
@@ -100,10 +102,17 @@ impl<S: StateReader> TransactionExecutor<S> {
                     &self.visited_storage_entries,
                     &tx_visited_storage_entries,
                 )?;
+                let message_segment_length = 0;
+                let tx_weights = calculate_tx_weights(
+                    additional_os_resources.clone(),
+                    actual_resources,
+                    message_segment_length,
+                );
                 let py_bouncer_info = PyBouncerInfo {
-                    message_segment_length: 0,
+                    message_segment_length,
                     state_diff_size: 0,
                     additional_os_resources: PyVmExecutionResources::from(additional_os_resources),
+                    tx_weights,
                 };
 
                 self.staged_for_commit_state = Some(