Handle specially top level call in trace

crates/debugging/src/lib.rs

@@ -6,4 +6,4 @@
 mod trace;
 mod tree;
 
-pub use trace::Trace;
+pub use trace::types::Trace;

crates/debugging/src/trace/collect.rs

@@ -0,0 +1,69 @@
+use crate::Trace;
+use crate::trace::types::{
+    CallerAddress, ContractName, ContractTrace, Selector, StorageAddress, TestName, TraceInfo,
+};
+use cheatnet::runtime_extensions::forge_runtime_extension::contracts_data::ContractsData;
+use cheatnet::state::{CallTrace, CallTraceNode};
+use std::cell::RefCell;
+use std::rc::Rc;
+
+pub fn trace(call_trace: &CallTrace, contracts_data: &ContractsData, test_name: String) -> Trace {
+    Trace {
+        test_name: TestName(test_name),
+        nested_calls: nested_calls(call_trace, contracts_data),
+    }
+}
+
+fn contract_trace(
+    call_trace: &Rc<RefCell<CallTrace>>,
+    contracts_data: &ContractsData,
+) -> ContractTrace {
+    let call_trace = call_trace.borrow();
+    let nested_calls = nested_calls(&call_trace, contracts_data);
+
+    let trace_info = TraceInfo {
+        contract_name: contract_name(&call_trace, contracts_data),
+        entry_point_type: call_trace.entry_point.entry_point_type,
+        calldata: call_trace.entry_point.calldata.clone(),
+        storage_address: StorageAddress(call_trace.entry_point.storage_address),
+        caller_address: CallerAddress(call_trace.entry_point.caller_address),
+        call_type: call_trace.entry_point.call_type,
+        nested_calls,
+        call_result: call_trace.result.clone(),
+    };
+
+    ContractTrace {
+        selector: selector(&call_trace, contracts_data),
+        trace_info,
+    }
+}
+
+fn nested_calls(call_trace: &CallTrace, contracts_data: &ContractsData) -> Vec<ContractTrace> {
+    call_trace
+        .nested_calls
+        .iter()
+        .filter_map(CallTraceNode::extract_entry_point_call)
+        .map(|call_trace| contract_trace(call_trace, contracts_data))
+        .collect()
+}
+
+fn contract_name(call_trace: &CallTrace, contracts_data: &ContractsData) -> ContractName {
+    contracts_data
+        .get_contract_name(
+            &call_trace
+                .entry_point
+                .class_hash
+                .expect("class_hash should be set in `fn execute_call_entry_point` in cheatnet"),
+        )
+        .cloned()
+        .map(ContractName)
+        .expect("contract name should be present in `ContractsData`")
+}
+
+fn selector(call_trace: &CallTrace, contracts_data: &ContractsData) -> Selector {
+    contracts_data
+        .get_function_name(&call_trace.entry_point.entry_point_selector)
+        .cloned()
+        .map(Selector)
+        .expect("selector should be present in `ContractsData`")
+}

crates/debugging/src/trace/mod.rs

@@ -0,0 +1,2 @@
+mod collect;
+pub mod types;

crates/debugging/src/trace/types.rs

@@ -0,0 +1,64 @@
+use crate::trace::collect;
+use crate::tree::TreeSerialize;
+use blockifier::execution::entry_point::CallType;
+use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::CallResult;
+use cheatnet::runtime_extensions::forge_runtime_extension::contracts_data::ContractsData;
+use cheatnet::state::CallTrace;
+use starknet_api::contract_class::EntryPointType;
+use starknet_api::core::ContractAddress;
+use starknet_api::transaction::fields::Calldata;
+use std::fmt;
+use std::fmt::Display;
+
+#[derive(Debug, Clone)]
+pub struct Trace {
+    pub test_name: TestName,
+    pub nested_calls: Vec<ContractTrace>,
+}
+
+#[derive(Debug, Clone)]
+pub struct ContractTrace {
+    pub selector: Selector,
+    pub trace_info: TraceInfo,
+}
+
+#[derive(Debug, Clone)]
+pub struct TraceInfo {
+    pub contract_name: ContractName,
+    pub entry_point_type: EntryPointType,
+    pub calldata: Calldata,
+    pub storage_address: StorageAddress,
+    pub caller_address: CallerAddress,
+    pub call_type: CallType,
+    pub nested_calls: Vec<ContractTrace>,
+    pub call_result: CallResult,
+}
+
+#[derive(Debug, Clone)]
+pub struct Selector(pub String);
+
+#[derive(Debug, Clone)]
+pub struct TestName(pub String);
+
+#[derive(Debug, Clone)]
+pub struct ContractName(pub String);
+
+#[derive(Debug, Clone)]
+pub struct StorageAddress(pub ContractAddress);
+
+#[derive(Debug, Clone)]
+pub struct CallerAddress(pub ContractAddress);
+
+impl Trace {
+    /// Creates a new [`Trace`] from a given `cheatnet` [`CallTrace`], [`ContractsData`] and a test name.
+    #[must_use]
+    pub fn new(call_trace: &CallTrace, contracts_data: &ContractsData, test_name: String) -> Self {
+        collect::trace(call_trace, contracts_data, test_name)
+    }
+}
+
+impl Display for Trace {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.serialize())
+    }
+}

crates/debugging/src/tree/ui/as_tree_node.rs

@@ -1,6 +1,5 @@
-use crate::trace::{Trace, TraceInfo};
+use crate::trace::types::{ContractTrace, Trace, TraceInfo};
 use crate::tree::building::node::Node;
-use crate::tree::ui::display::NodeDisplay;
 
 /// Trait for adding a type to a tree.
 /// Implementations of this trait should only focus on placements of nodes in a tree not display aspects of them.
@@ -10,6 +9,15 @@ pub trait AsTreeNode {
 }
 
 impl AsTreeNode for Trace {
+    fn as_tree_node(&self, parent: &mut Node) {
+        let mut node = parent.child_node(&self.test_name);
+        for nested_call in &self.nested_calls {
+            node.as_tree_node(nested_call);
+        }
+    }
+}
+
+impl AsTreeNode for ContractTrace {
     fn as_tree_node(&self, parent: &mut Node) {
         parent
             .child_node(&self.selector)
@@ -19,7 +27,7 @@ impl AsTreeNode for Trace {
 
 impl AsTreeNode for TraceInfo {
     fn as_tree_node(&self, parent: &mut Node) {
-        parent.as_tree_node(&self.contract_name.as_ref());
+        parent.leaf(&self.contract_name);
         parent.leaf(&self.entry_point_type);
         parent.leaf(&self.calldata);
         parent.leaf(&self.storage_address);
@@ -31,11 +39,3 @@ impl AsTreeNode for TraceInfo {
         }
     }
 }
-
-impl<T: NodeDisplay> AsTreeNode for Option<&T> {
-    fn as_tree_node(&self, parent: &mut Node) {
-        if let Some(value) = *self {
-            parent.leaf(value);
-        }
-    }
-}

crates/debugging/src/tree/ui/display.rs

@@ -1,4 +1,4 @@
-use crate::trace::{CallerAddress, ContractName, Selector, StorageAddress};
+use crate::trace::types::{CallerAddress, ContractName, Selector, StorageAddress, TestName};
 use blockifier::execution::entry_point::CallType;
 use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{
     CallFailure, CallResult,
@@ -22,6 +22,13 @@ pub trait NodeDisplay {
     }
 }
 
+impl NodeDisplay for TestName {
+    const TAG: &'static str = "test name";
+    fn string_pretty(&self) -> String {
+        self.0.clone()
+    }
+}
+
 impl NodeDisplay for ContractName {
     const TAG: &'static str = "contract name";
     fn string_pretty(&self) -> String {
@@ -32,9 +39,7 @@ impl NodeDisplay for ContractName {
 impl NodeDisplay for Selector {
     const TAG: &'static str = "selector";
     fn string_pretty(&self) -> String {
-        self.function_name
-            .as_ref()
-            .map_or_else(|| string_hex(self.selector.0), ToString::to_string)
+        self.0.to_string()
     }
 }
 

crates/forge-runner/src/running.rs

@@ -422,9 +422,10 @@ fn extract_test_case_summary(
                         fuzzer_args: result_with_info.fuzzer_args,
                         test_statistics: (),
                         debugging_trace: cfg!(feature = "debugging").then(|| {
-                            debugging::Trace::from_call_trace(
-                                &result_with_info.call_trace,
+                            debugging::Trace::new(
+                                &result_with_info.call_trace.borrow(),
                                 contracts_data,
+                                case.name.clone(),
                             )
                         }),
                     }

crates/forge-runner/src/test_case_summary.rs

@@ -296,7 +296,7 @@ impl TestCaseSummary<Single> {
             .map(|msg| add_backtrace_footer(msg, contracts_data, encountered_errors));
 
         let debugging_trace = cfg!(feature = "debugging")
-            .then(|| debugging::Trace::from_call_trace(call_trace, contracts_data));
+            .then(|| debugging::Trace::new(&call_trace.borrow(), contracts_data, name.clone()));
 
         match run_result.value {
             RunResultValue::Success(_) => match &test_case.config.expected_result {