Display gas report table for fork contracts

Author: ddoktorski

crates/forge-runner/src/gas.rs

@@ -18,6 +18,7 @@ use starknet_api::transaction::fields::GasVectorComputationMode;
 
 pub mod report;
 pub mod stats;
+mod utils;
 
 #[tracing::instrument(skip_all, level = "debug")]
 pub fn calculate_used_gas(

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

@@ -1,4 +1,5 @@
 use crate::gas::stats::GasStats;
+use crate::gas::utils::shorten_felt;
 use cheatnet::trace_data::{CallTrace, CallTraceNode};
 use comfy_table::modifiers::UTF8_ROUND_CORNERS;
 use comfy_table::{Attribute, Cell, Color, Table};
@@ -12,14 +13,20 @@ use std::fmt::Display;
 type ContractName = String;
 type Selector = String;
 
+#[derive(Debug, Clone, PartialOrd, PartialEq, Ord, Eq)]
+pub enum ContractId {
+    TestContract(ContractName),
+    ForkedContract(ClassHash),
+}
+
 #[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>);
+pub struct ReportData(BTreeMap<ContractId, ContractInfo>);
 
 #[derive(Debug, Clone, Default)]
 pub struct ContractInfo {
@@ -58,15 +65,15 @@ impl SingleTestGasInfo {
                     "class_hash should be set in `fn execute_call_entry_point` in cheatnet",
                 );
 
-                let contract_name = get_contract_name(contracts_data, class_hash);
+                let contract_id = get_contract_id(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);
+                report_data.update_entry(contract_id, selector, gas);
                 stack.extend(call.nested_calls.clone());
             }
         }
@@ -80,13 +87,8 @@ impl SingleTestGasInfo {
 }
 
 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();
+    fn update_entry(&mut self, contract_id: ContractId, selector: Selector, gas_used: GasVector) {
+        let contract_info = self.0.entry(contract_id).or_default();
 
         let current_gas = contract_info.gas_used;
         contract_info.gas_used = current_gas.checked_add(gas_used).unwrap_or_else(|| {
@@ -119,11 +121,26 @@ impl Display for ReportData {
     }
 }
 
-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()
+impl Display for ContractId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let name = match self {
+            ContractId::TestContract(name) => format!("{name} Contract"),
+            ContractId::ForkedContract(class_hash) => {
+                format!(
+                    "forked contract\n(class hash: {})",
+                    shorten_felt(class_hash.0)
+                )
+            }
+        };
+        write!(f, "{name}")
+    }
+}
+
+fn get_contract_id(contracts_data: &ContractsDataStore, class_hash: ClassHash) -> ContractId {
+    match contracts_data.get_contract_name(&class_hash) {
+        Some(name) => ContractId::TestContract(name.0.to_string()),
+        None => ContractId::ForkedContract(class_hash),
+    }
 }
 
 fn get_selector(contracts_data: &ContractsDataStore, selector: EntryPointSelector) -> Selector {
@@ -134,13 +151,11 @@ fn get_selector(contracts_data: &ContractsDataStore, selector: EntryPointSelecto
         .clone()
 }
 
-pub fn format_table_output(contract_info: &ContractInfo, name: &ContractName) -> Table {
+pub fn format_table_output(contract_info: &ContractInfo, contract_id: &ContractId) -> Table {
     let mut table = Table::new();
     table.apply_modifier(UTF8_ROUND_CORNERS);
 
-    table.set_header(vec![
-        Cell::new(format!("{name} Contract")).fg(Color::Magenta),
-    ]);
+    table.set_header(vec![Cell::new(contract_id.to_string()).fg(Color::Magenta)]);
     table.add_row(vec![
         Cell::new("Function Name").add_attribute(Attribute::Bold),
         Cell::new("Min").add_attribute(Attribute::Bold),

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

@@ -0,0 +1,27 @@
+use starknet_types_core::felt::Felt;
+
+pub(super) fn shorten_felt(felt: Felt) -> String {
+    let padded = format!("{felt:#066x}");
+    let first = &padded[..4];
+    let last = &padded[padded.len() - 4..];
+    format!("{first}…{last}")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_long() {
+        let felt = Felt::from_hex_unchecked(
+            "0x01c902da594beda43db10142ecf1fc3a098b56e8d95f3cd28587a0c6ba05a451",
+        );
+        assert_eq!("0x01…a451", &shorten_felt(felt));
+    }
+
+    #[test]
+    fn test_short() {
+        let felt = Felt::from_hex_unchecked("0x123");
+        assert_eq!("0x00…0123", &shorten_felt(felt));
+    }
+}

crates/forge/tests/e2e/gas_report.rs

@@ -1,4 +1,6 @@
-use crate::e2e::common::runner::{setup_package, test_runner};
+use crate::e2e::common::runner::{
+    BASE_FILE_PATTERNS, Package, setup_package, setup_package_with_file_patterns, test_runner,
+};
 use indoc::indoc;
 use shared::test_utils::output_assert::assert_stdout_contains;
 
@@ -115,3 +117,50 @@ fn multiple_contracts_and_constructor() {
     "},
     );
 }
+
+#[test]
+fn fork() {
+    let temp =
+        setup_package_with_file_patterns(Package::Name("forking".to_string()), BASE_FILE_PATTERNS);
+
+    let output = test_runner(&temp)
+        .arg("test_track_resources")
+        .arg("--gas-report")
+        .assert()
+        .code(0);
+
+    assert_stdout_contains(
+        output,
+        indoc! {r"
+        [..]Compiling[..]
+        [..]Finished[..]
+
+        Collected 1 test(s) from forking package
+        Running 1 test(s) from src/
+        [PASS] forking::tests::test_track_resources (l1_gas: ~0, l1_data_gas: ~[..], l2_gas: ~[..])
+        ╭-------------------------+-------+-------+-------+---------+---------╮
+        | forked contract         |       |       |       |         |         |
+        | (class hash: 0x06…1550) |       |       |       |         |         |
+        +=====================================================================+
+        | Function Name           | Min   | Max   | Avg   | Std Dev | # Calls |
+        |-------------------------+-------+-------+-------+---------+---------|
+        | get_balance             | 40000 | 40000 | 40000 | 0       | 1       |
+        |-------------------------+-------+-------+-------+---------+---------|
+        | increase_balance        | 40000 | 40000 | 40000 | 0       | 1       |
+        ╰-------------------------+-------+-------+-------+---------+---------╯
+
+        ╭-------------------------+-------+-------+-------+---------+---------╮
+        | forked contract         |       |       |       |         |         |
+        | (class hash: 0x07…af4b) |       |       |       |         |         |
+        +=====================================================================+
+        | Function Name           | Min   | Max   | Avg   | Std Dev | # Calls |
+        |-------------------------+-------+-------+-------+---------+---------|
+        | get_balance             | 13840 | 13840 | 13840 | 0       | 1       |
+        |-------------------------+-------+-------+-------+---------+---------|
+        | increase_balance        | 25840 | 25840 | 25840 | 0       | 1       |
+        ╰-------------------------+-------+-------+-------+---------+---------╯
+
+        Tests: 1 passed, 0 failed, 0 ignored, [..] filtered out
+        "},
+    );
+}