Store execution summary in trace data

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs

@@ -17,14 +17,17 @@ use crate::runtime_extensions::{
         storage::{calculate_variable_address, load, store},
     },
 };
-use crate::state::{CallTrace, CallTraceNode};
+use crate::state::{CallTrace, CallTraceNode, GasReportData};
 use anyhow::{Context, Result, anyhow};
 use blockifier::bouncer::vm_resources_to_sierra_gas;
 use blockifier::context::TransactionContext;
-use blockifier::execution::call_info::CallInfo;
+use blockifier::execution::call_info::{
+    CallInfo, CallSummary, ChargedResources, EventSummary, ExecutionSummary, OrderedEvent,
+};
 use blockifier::execution::contract_class::TrackedResource;
 use blockifier::execution::syscalls::vm_syscall_utils::{SyscallSelector, SyscallUsageMap};
 use blockifier::state::errors::StateError;
+use blockifier::utils::u64_from_usize;
 use cairo_vm::vm::runners::cairo_runner::CairoRunner;
 use cairo_vm::vm::{
     errors::hint_errors::HintError, runners::cairo_runner::ExecutionResources,
@@ -42,10 +45,11 @@ use runtime::{
 };
 use scarb_oracle_hint_service::OracleHintService;
 use starknet::signers::SigningKey;
+use starknet_api::execution_resources::GasAmount;
 use starknet_api::{contract_class::EntryPointType::L1Handler, core::ClassHash};
 use starknet_types_core::felt::Felt;
 use std::cell::RefCell;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 
@@ -738,6 +742,65 @@ pub fn update_top_call_vm_trace(runtime: &mut ForgeRuntime, cairo_runner: &mut C
     }
 }
 
+pub fn compute_and_store_execution_summary(trace: &Rc<RefCell<CallTrace>>) {
+    let execution_summary = if trace.borrow().nested_calls.is_empty() {
+        get_execution_summary_without_nested_calls(trace)
+    } else {
+        let mut nested_calls_summaries = vec![];
+        for nested_call in &trace.borrow().nested_calls {
+            if let CallTraceNode::EntryPointCall(nested_call) = nested_call {
+                compute_and_store_execution_summary(nested_call);
+                nested_calls_summaries.push(
+                    nested_call
+                        .borrow()
+                        .gas_report_data
+                        .as_ref()
+                        .expect("Gas report data must be set after calling `compute_and_store_execution_summary`")
+                        .execution_summary
+                        .clone());
+            }
+        }
+        let mut current_call_summary = get_execution_summary_without_nested_calls(trace)
+            + nested_calls_summaries.into_iter().sum();
+
+        // vm_resources and gas_consumed of a call already contain the resources of its inner calls.
+        current_call_summary.charged_resources.vm_resources =
+            trace.borrow().used_execution_resources.clone();
+        current_call_summary.charged_resources.gas_consumed =
+            GasAmount(trace.borrow().gas_consumed);
+        current_call_summary
+    };
+
+    trace.borrow_mut().gas_report_data = Some(GasReportData::new(execution_summary.clone()));
+}
+
+// Based on blockifier/src/execution/call_info.rs (summarize)
+fn get_execution_summary_without_nested_calls(trace: &Rc<RefCell<CallTrace>>) -> ExecutionSummary {
+    let current_call = trace.borrow();
+    ExecutionSummary {
+        charged_resources: ChargedResources {
+            vm_resources: current_call.used_execution_resources.clone(),
+            gas_consumed: GasAmount(current_call.gas_consumed),
+        },
+        l2_to_l1_payload_lengths: current_call.used_l1_resources.l2_l1_message_sizes.clone(),
+        event_summary: {
+            let mut event_summary = EventSummary {
+                n_events: current_call.events.len(),
+                ..Default::default()
+            };
+            for OrderedEvent { event, .. } in &current_call.events {
+                event_summary.total_event_data_size += u64_from_usize(event.data.0.len());
+                event_summary.total_event_keys += u64_from_usize(event.keys.len());
+            }
+            event_summary
+        },
+        // Fields below are not relevant for partial gas calculation.
+        call_summary: CallSummary::default(),
+        executed_class_hashes: HashSet::default(),
+        visited_storage_entries: HashSet::default(),
+    }
+}
+
 fn add_sierra_gas_resources(top_call: &Rc<RefCell<CallTrace>>) -> u64 {
     let mut gas_consumed = top_call.borrow().gas_consumed;
     for nested_call in &top_call.borrow().nested_calls {

crates/cheatnet/src/state.rs

@@ -10,7 +10,7 @@ use crate::runtime_extensions::forge_runtime_extension::cheatcodes::cheat_execut
 };
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_events::Event;
 use crate::runtime_extensions::forge_runtime_extension::cheatcodes::spy_messages_to_l1::MessageToL1;
-use blockifier::execution::call_info::{OrderedEvent, OrderedL2ToL1Message};
+use blockifier::execution::call_info::{ExecutionSummary, OrderedEvent, OrderedL2ToL1Message};
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::execution::entry_point::CallEntryPoint;
 use blockifier::execution::syscalls::vm_syscall_utils::SyscallUsageMap;
@@ -203,6 +203,18 @@ impl<T> CheatStatus<T> {
     }
 }
 
+#[derive(Debug)]
+pub struct GasReportData {
+    pub execution_summary: ExecutionSummary,
+}
+
+impl GasReportData {
+    #[must_use]
+    pub fn new(execution_summary: ExecutionSummary) -> Self {
+        Self { execution_summary }
+    }
+}
+
 /// Tree structure representing trace of a call.
 #[derive(Debug)]
 pub struct CallTrace {
@@ -221,6 +233,9 @@ pub struct CallTrace {
     pub gas_consumed: u64,
     pub events: Vec<OrderedEvent>,
     pub signature: Vec<Felt>,
+
+    // This is updated only once after the entire test execution.
+    pub gas_report_data: Option<GasReportData>,
 }
 
 impl CairoSerialize for CallTrace {
@@ -253,6 +268,7 @@ impl CallTrace {
             gas_consumed: u64::default(),
             events: vec![],
             signature: vec![],
+            gas_report_data: None,
         }
     }
 

crates/forge-runner/src/running.rs

@@ -23,8 +23,9 @@ use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::CallToBl
 use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::UsedResources;
 use cheatnet::runtime_extensions::cheatable_starknet_runtime_extension::CheatableStarknetRuntimeExtension;
 use cheatnet::runtime_extensions::forge_runtime_extension::{
-    ForgeExtension, ForgeRuntime, add_resources_to_top_call, get_all_used_resources,
-    update_top_call_l1_resources, update_top_call_resources, update_top_call_vm_trace,
+    ForgeExtension, ForgeRuntime, add_resources_to_top_call, compute_and_store_execution_summary,
+    get_all_used_resources, update_top_call_l1_resources, update_top_call_resources,
+    update_top_call_vm_trace,
 };
 use cheatnet::state::{
     BlockInfoReader, CallTrace, CheatnetState, EncounteredErrors, ExtendedStateReader,
@@ -332,6 +333,7 @@ pub fn run_test_case(
 
     update_top_call_resources(&mut forge_runtime, tracked_resource);
     update_top_call_l1_resources(&mut forge_runtime);
+    compute_and_store_execution_summary(&call_trace_ref);
 
     let fuzzer_args = forge_runtime
         .extended_runtime