Add gas report logic

crates/cheatnet/src/trace_data.rs

@@ -240,8 +240,8 @@ impl GasReportData {
         }
     }
 
-    pub fn get_gas(&self) -> &GasVector {
-        self.partial_gas_usage.get_or_init(|| {
+    pub fn get_gas(&self) -> GasVector {
+        *self.partial_gas_usage.get_or_init(|| {
             self.execution_summary.clone().to_partial_gas_vector(
                 VersionedConstants::latest_constants(),
                 &GasVectorComputationMode::All,

crates/debugging/src/lib.rs

@@ -7,5 +7,6 @@ mod contracts_data_store;
 mod trace;
 mod tree;
 
+pub use contracts_data_store::ContractsDataStore;
 pub use trace::components::{Component, Components};
-pub use trace::{context::Context, types::Trace};
+pub use trace::{context::Context, types::ContractName, types::Trace};

crates/forge-runner/src/forge_config.rs

@@ -34,6 +34,7 @@ pub struct OutputConfig {
     pub trace_args: TraceArgs,
     pub detailed_resources: bool,
     pub execution_data_to_save: ExecutionDataToSave,
+    pub gas_report: bool,
 }
 
 #[derive(Debug, PartialEq, Clone, Default)]

crates/forge-runner/src/gas.rs

@@ -16,6 +16,7 @@ use cheatnet::state::ExtendedStateReader;
 use starknet_api::execution_resources::{GasAmount, GasVector};
 use starknet_api::transaction::fields::GasVectorComputationMode;
 
+pub mod report;
 pub mod stats;
 
 #[tracing::instrument(skip_all, level = "debug")]
@@ -149,16 +150,18 @@ pub fn check_available_gas(
             ..
         } if available_gas.is_some_and(|available_gas| {
             let av_gas = available_gas.to_gas_vector();
-            gas_info.l1_gas > av_gas.l1_gas
-                || gas_info.l1_data_gas > av_gas.l1_data_gas
-                || gas_info.l2_gas > av_gas.l2_gas
+            gas_info.gas_used.l1_gas > av_gas.l1_gas
+                || gas_info.gas_used.l1_data_gas > av_gas.l1_data_gas
+                || gas_info.gas_used.l2_gas > av_gas.l2_gas
         }) =>
         {
             TestCaseSummary::Failed {
                 name,
                 msg: Some(format!(
                     "\n\tTest cost exceeded the available gas. Consumed l1_gas: ~{}, l1_data_gas: ~{}, l2_gas: ~{}",
-                    gas_info.l1_gas, gas_info.l1_data_gas, gas_info.l2_gas
+                    gas_info.gas_used.l1_gas,
+                    gas_info.gas_used.l1_data_gas,
+                    gas_info.gas_used.l2_gas
                 )),
                 fuzzer_args: Vec::default(),
                 test_statistics: (),

crates/forge-runner/src/gas/report.rs

@@ -0,0 +1,119 @@
+use crate::gas::stats::GasStats;
+use cheatnet::trace_data::{CallTrace, CallTraceNode};
+use debugging::ContractsDataStore;
+use starknet_api::core::{ClassHash, EntryPointSelector};
+use starknet_api::execution_resources::GasVector;
+use std::collections::BTreeMap;
+
+type ContractName = String;
+type Selector = String;
+
+#[derive(Debug, Clone)]
+pub struct SingleTestGasInfo {
+    pub gas_used: GasVector,
+    pub report_data: Option<ReportData>,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct ReportData(BTreeMap<ContractName, ContractInfo>);
+
+#[derive(Debug, Clone, Default)]
+pub struct ContractInfo {
+    pub(super) gas_used: GasVector,
+    pub(super) functions: BTreeMap<Selector, SelectorReportData>,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct SelectorReportData {
+    pub(super) gas_stats: GasStats,
+    pub(super) n_calls: u64,
+    pub(super) records: Vec<u64>,
+}
+
+impl SingleTestGasInfo {
+    #[must_use]
+    pub(crate) fn new(gas_used: GasVector) -> Self {
+        Self {
+            gas_used,
+            report_data: None,
+        }
+    }
+
+    pub(crate) fn get_with_report_data(
+        self,
+        trace: &CallTrace,
+        contracts_data: &ContractsDataStore,
+    ) -> Self {
+        let mut report_data = ReportData::default();
+        let mut stack = trace.nested_calls.clone();
+
+        while let Some(call_trace_node) = stack.pop() {
+            if let CallTraceNode::EntryPointCall(call) = call_trace_node {
+                let call = call.borrow();
+                let class_hash = call.entry_point.class_hash.expect(
+                    "class_hash should be set in `fn execute_call_entry_point` in cheatnet",
+                );
+
+                let contract_name = get_contract_name(contracts_data, class_hash);
+                let selector = get_selector(contracts_data, call.entry_point.entry_point_selector);
+                let gas = call
+                    .gas_report_data
+                    .as_ref()
+                    .expect("Gas report data must be updated after test execution")
+                    .get_gas();
+
+                report_data.update_entry(contract_name, selector, gas);
+                stack.extend(call.nested_calls.clone());
+            }
+        }
+        report_data.finalize();
+
+        Self {
+            gas_used: self.gas_used,
+            report_data: Some(report_data),
+        }
+    }
+}
+
+impl ReportData {
+    fn update_entry(
+        &mut self,
+        contract_name: ContractName,
+        selector: Selector,
+        gas_used: GasVector,
+    ) {
+        let contract_info = self.0.entry(contract_name).or_default();
+
+        let current_gas = contract_info.gas_used;
+        contract_info.gas_used = current_gas.checked_add(gas_used).unwrap_or_else(|| {
+            panic!("Gas addition overflow when adding {gas_used:?} to {current_gas:?}.")
+        });
+
+        let entry = contract_info.functions.entry(selector).or_default();
+        entry.records.push(gas_used.l2_gas.0);
+        entry.n_calls += 1;
+    }
+
+    fn finalize(&mut self) {
+        for contract_info in self.0.values_mut() {
+            for gas_info in contract_info.functions.values_mut() {
+                gas_info.gas_stats = GasStats::new(&gas_info.records);
+            }
+        }
+    }
+}
+
+fn get_contract_name(contracts_data: &ContractsDataStore, class_hash: ClassHash) -> ContractName {
+    contracts_data
+        .get_contract_name(&class_hash)
+        .map_or("forked contract", |name| name.0.as_str())
+        .to_string()
+}
+
+fn get_selector(contracts_data: &ContractsDataStore, selector: EntryPointSelector) -> Selector {
+    contracts_data
+        .get_selector(&selector)
+        .expect("`Selector` should be present")
+        .0
+        .clone()
+}

crates/forge-runner/src/messages.rs

@@ -79,7 +79,9 @@ impl TestResultMessage {
             AnyTestCaseSummary::Single(TestCaseSummary::Passed { gas_info, .. }) => {
                 format!(
                     " (l1_gas: ~{}, l1_data_gas: ~{}, l2_gas: ~{})",
-                    gas_info.l1_gas, gas_info.l1_data_gas, gas_info.l2_gas
+                    gas_info.gas_used.l1_gas,
+                    gas_info.gas_used.l1_data_gas,
+                    gas_info.gas_used.l2_gas
                 )
             }
             _ => String::new(),

crates/forge-runner/src/running.rs

@@ -409,6 +409,7 @@ fn extract_test_case_summary(
                 contracts_data,
                 versioned_program_path,
                 trace_args,
+                forge_config.output_config.gas_report,
             ),
             RunResult::Error(run_error) => {
                 let mut message = format!(

crates/forge-runner/src/test_case_summary.rs

@@ -3,6 +3,7 @@ use crate::build_trace_data::build_profiler_call_trace;
 use crate::debugging::{TraceArgs, build_debugging_trace};
 use crate::expected_result::{ExpectedPanicValue, ExpectedTestResult};
 use crate::gas::check_available_gas;
+use crate::gas::report::SingleTestGasInfo;
 use crate::gas::stats::GasStats;
 use crate::package_tests::with_config_resolved::TestCaseWithResolvedConfig;
 use crate::running::{RunCompleted, RunStatus};
@@ -12,6 +13,7 @@ use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::Use
 use cheatnet::runtime_extensions::forge_runtime_extension::contracts_data::ContractsData;
 use conversions::byte_array::ByteArray;
 use conversions::felt::ToShortString;
+use debugging::ContractsDataStore;
 use shared::utils::build_readable_text;
 use starknet_api::execution_resources::GasVector;
 use starknet_types_core::felt::Felt;
@@ -82,7 +84,7 @@ impl TestType for Fuzzing {
 #[derive(Debug, PartialEq, Clone)]
 pub struct Single;
 impl TestType for Single {
-    type GasInfo = GasVector;
+    type GasInfo = SingleTestGasInfo;
     type TestStatistics = ();
     type TraceData = VersionedProfilerCallTrace;
 }
@@ -195,7 +197,7 @@ impl TestCaseSummary<Fuzzing> {
                 let gas_usages: Vec<GasVector> = results
                     .into_iter()
                     .map(|a| match a {
-                        TestCaseSummary::Passed { gas_info, .. } => gas_info,
+                        TestCaseSummary::Passed { gas_info, .. } => gas_info.gas_used,
                         _ => unreachable!(),
                     })
                     .collect();
@@ -278,7 +280,7 @@ impl TestCaseSummary<Single> {
         RunCompleted {
             status,
             call_trace,
-            gas_used: gas_info,
+            gas_used,
             used_resources,
             encountered_errors,
             fuzzer_args,
@@ -288,6 +290,7 @@ impl TestCaseSummary<Single> {
         contracts_data: &ContractsData,
         versioned_program_path: &Utf8Path,
         trace_args: &TraceArgs,
+        gas_report_enabled: bool,
     ) -> Self {
         let name = test_case.name.clone();
 
@@ -299,6 +302,16 @@ impl TestCaseSummary<Single> {
             &fork_data,
         );
 
+        let gas_info = SingleTestGasInfo::new(gas_used);
+        let gas_info = if gas_report_enabled {
+            gas_info.get_with_report_data(
+                &call_trace.borrow(),
+                &ContractsDataStore::new(contracts_data, &fork_data),
+            )
+        } else {
+            gas_info
+        };
+
         match status {
             RunStatus::Success(data) => match &test_case.config.expected_result {
                 ExpectedTestResult::Success => {

crates/forge/src/combine_configs.rs

@@ -21,6 +21,7 @@ pub fn combine_configs(
     save_trace_data: bool,
     build_profile: bool,
     coverage: bool,
+    gas_report: bool,
     max_n_steps: Option<u32>,
     tracked_resource: ForgeTrackedResource,
     contracts_data: ContractsData,
@@ -58,6 +59,7 @@ pub fn combine_configs(
             trace_args,
             detailed_resources: detailed_resources || forge_config_from_scarb.detailed_resources,
             execution_data_to_save,
+            gas_report: gas_report || forge_config_from_scarb.gas_report,
         }),
     }
 }
@@ -76,6 +78,7 @@ mod tests {
             false,
             false,
             false,
+            false,
             None,
             ForgeTrackedResource::CairoSteps,
             ContractsData::default(),
@@ -93,6 +96,7 @@ mod tests {
             false,
             false,
             false,
+            false,
             None,
             ForgeTrackedResource::CairoSteps,
             ContractsData::default(),
@@ -121,6 +125,7 @@ mod tests {
             false,
             false,
             false,
+            false,
             None,
             ForgeTrackedResource::CairoSteps,
             ContractsData::default(),
@@ -149,6 +154,7 @@ mod tests {
                     detailed_resources: false,
                     execution_data_to_save: ExecutionDataToSave::default(),
                     trace_args: TraceArgs::default(),
+                    gas_report: false,
                 }),
             }
         );
@@ -165,6 +171,7 @@ mod tests {
             save_trace_data: true,
             build_profile: true,
             coverage: true,
+            gas_report: true,
             max_n_steps: Some(1_000_000),
             tracked_resource: ForgeTrackedResource::CairoSteps,
         };
@@ -177,6 +184,7 @@ mod tests {
             false,
             false,
             false,
+            false,
             None,
             ForgeTrackedResource::CairoSteps,
             ContractsData::default(),
@@ -210,6 +218,7 @@ mod tests {
                         additional_args: vec![],
                     },
                     trace_args: TraceArgs::default(),
+                    gas_report: true,
                 }),
             }
         );
@@ -226,6 +235,7 @@ mod tests {
             save_trace_data: false,
             build_profile: false,
             coverage: false,
+            gas_report: false,
             max_n_steps: Some(1234),
             tracked_resource: ForgeTrackedResource::CairoSteps,
         };
@@ -237,6 +247,7 @@ mod tests {
             true,
             true,
             true,
+            true,
             Some(1_000_000),
             ForgeTrackedResource::CairoSteps,
             ContractsData::default(),
@@ -271,6 +282,7 @@ mod tests {
                         additional_args: vec![],
                     },
                     trace_args: TraceArgs::default(),
+                    gas_report: true,
                 }),
             }
         );

crates/forge/src/lib.rs

@@ -213,6 +213,10 @@ pub struct TestArgs {
     #[arg(long, value_enum, default_value_t)]
     tracked_resource: ForgeTrackedResource,
 
+    /// Display a table of L2 gas breakdown for each contract and selector
+    #[arg(long)]
+    gas_report: bool,
+
     /// Additional arguments for cairo-coverage or cairo-profiler
     #[arg(last = true)]
     additional_args: Vec<OsString>,

crates/forge/src/run_tests/package.rs

@@ -102,6 +102,7 @@ impl RunForPackageArgs {
             args.save_trace_data,
             args.build_profile,
             args.coverage,
+            args.gas_report,
             args.max_n_steps,
             args.tracked_resource,
             contracts_data,

crates/forge/src/scarb/config.rs

@@ -59,6 +59,9 @@ pub struct ForgeConfigFromScarb {
     /// Generate a coverage report for the executed tests which have passed and are not fuzz tests
     #[serde(default)]
     pub coverage: bool,
+    /// Display a table of L2 gas breakdown for each contract and selector
+    #[serde(default)]
+    pub gas_report: bool,
     /// Fork configuration profiles
     #[serde(default, deserialize_with = "validate_forks")]
     pub fork: Vec<ForkTarget>,