Add function level trace

commit-id:3a7228bd

Author: ksew1

Cargo.lock

@@ -8,14 +8,15 @@ blockifier.workspace = true
 cairo-lang-sierra.workspace = true
 cairo-lang-sierra-to-casm.workspace = true
 cairo-lang-starknet-classes.workspace = true
+cairo-annotations.workspace = true
 cheatnet = { path = "../cheatnet" }
 console.workspace = true
 data-transformer = { path = "../data-transformer" }
 paste.workspace = true
 ptree.workspace = true
 rayon.workspace = true
+regex.workspace = true
 serde_json.workspace = true
 starknet-types-core.workspace = true
 starknet.workspace = true
-starknet_api.workspace = true
-thiserror.workspace = true
+starknet_api.workspace = true

crates/debugging/Cargo.toml

@@ -122,6 +122,13 @@ impl ContractsDataStore {
     pub fn get_casm_debug_info(&self, class_hash: &ClassHash) -> Option<&CairoProgramDebugInfo> {
         self.casm_debug_infos.get(class_hash)
     }
+
+    /// Checks if the contract with the given [`ClassHash`] is a forked contract.
+    pub fn is_fork(&self, class_hash: &ClassHash) -> bool {
+        // We create contract names only from `ContractsData` and not from `ForkData`,
+        // so if the contract name is not present in `contract_names`, it is a fork
+        !self.contract_names.contains_key(class_hash)
+    }
 }
 
 /// Compile the given [`ContractClass`] to `casm` and return [`CairoProgramDebugInfo`]

crates/debugging/src/contracts_data_store.rs

@@ -1,4 +1,5 @@
 use crate::contracts_data_store::ContractsDataStore;
+use crate::trace::function::FunctionTrace;
 use crate::trace::types::{
     CallerAddress, ContractAddress, ContractName, ContractTrace, Selector, TestName, TraceInfo,
     TransformedCallResult, TransformedCalldata,
@@ -12,6 +13,7 @@ use data_transformer::{reverse_transform_input, reverse_transform_output};
 use starknet::core::types::contract::AbiEntry;
 use starknet_api::core::ClassHash;
 use starknet_api::execution_utils::format_panic_data;
+use std::ops::Not;
 
 pub struct Collector<'a> {
     call_trace: &'a CallTrace,
@@ -38,6 +40,7 @@ impl<'a> Collector<'a> {
     fn collect_contract_trace(&self) -> ContractTrace {
         let components = self.context.components();
         let entry_point = &self.call_trace.entry_point;
+
         let nested_calls = self.collect_nested_calls();
         let contract_name = self.collect_contract_name();
         let abi = self.collect_abi();
@@ -52,6 +55,12 @@ impl<'a> Collector<'a> {
             call_type: components.call_type(entry_point.call_type),
             nested_calls,
             call_result: components.call_result_lazy(|| self.collect_transformed_call_result(abi)),
+            #[expect(clippy::obfuscated_if_else)]
+            function_trace: self
+                .is_fork()
+                .not()
+                .then(|| components.function_trace_lazy(|| self.collect_function_trace()))
+                .unwrap_or_default(),
         };
 
         ContractTrace {
@@ -125,6 +134,18 @@ impl<'a> Collector<'a> {
         })
     }
 
+    fn collect_function_trace(&self) -> FunctionTrace {
+        FunctionTrace::create(
+            *self.class_hash(),
+            self.call_trace,
+            self.contracts_data_store(),
+        )
+    }
+
+    fn is_fork(&self) -> bool {
+        self.contracts_data_store().is_fork(self.class_hash())
+    }
+
     fn class_hash(&self) -> &ClassHash {
         self.call_trace
             .entry_point

crates/debugging/src/trace/collect.rs

@@ -1,3 +1,4 @@
+use crate::trace::function::FunctionTrace;
 use crate::trace::types::{
     CallerAddress, ContractAddress, ContractName, TransformedCallResult, TransformedCalldata,
 };
@@ -42,6 +43,8 @@ pub enum Component {
     CallType,
     /// The result of the call, transformed for display.
     CallResult,
+    /// The function trace, will be only collected if the contract is not a fork.
+    FunctionTrace,
 }
 
 macro_rules! impl_component_container {
@@ -58,7 +61,7 @@ macro_rules! impl_component_container {
                 "` that is computed only if `Component::", stringify!($variant),
                 "` is included in the [`Components`]."
             )]
-            #[derive(Debug, Clone)]
+            #[derive(Debug, Clone, Default)]
             pub struct [<$variant Container>] {
                 value: Option<$ty>,
             }
@@ -122,3 +125,4 @@ impl_component_container!(ContractAddress);
 impl_component_container!(CallerAddress);
 impl_component_container!(CallType);
 impl_component_container!(CallResult, TransformedCallResult);
+impl_component_container!(FunctionTrace);

crates/debugging/src/trace/components.rs

@@ -0,0 +1,7 @@
+mod name;
+mod node;
+mod stack;
+mod trace;
+
+pub use node::FunctionNode;
+pub use trace::FunctionTrace;

crates/debugging/src/trace/function/mod.rs

@@ -0,0 +1,49 @@
+use cairo_lang_sierra::program::{Program, StatementIdx};
+use regex::Regex;
+use std::sync::LazyLock;
+
+/// Represents a function name in the Sierra program.
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+pub enum FunctionName {
+    NonInlined(String),
+    // TODO: Add inlined variant in next PR.
+}
+
+impl FunctionName {
+    /// Creates a [`FunctionName`] from a [`StatementIdx`] index and a [`Program`].
+    pub fn from_program(statement_idx: StatementIdx, program: &Program) -> Self {
+        let function_idx = program
+            .funcs
+            .partition_point(|f| f.entry_point.0 <= statement_idx.0)
+            - 1;
+        let function_name = program.funcs[function_idx].id.to_string();
+        let function_name = remove_loop_suffix(&function_name);
+        let function_name = remove_monomorphization_suffix(&function_name);
+        FunctionName::NonInlined(function_name)
+    }
+
+    /// Returns the function name as a [`&str`].
+    pub fn function_name(&self) -> &str {
+        match self {
+            FunctionName::NonInlined(name) => name,
+        }
+    }
+}
+
+/// Remove suffix in case of loop function e.g. `[expr36]`.
+fn remove_loop_suffix(function_name: &str) -> String {
+    static RE_LOOP_FUNC: LazyLock<Regex> = LazyLock::new(|| {
+        Regex::new(r"\[expr\d*]")
+            .expect("Failed to create regex for normalizing loop function names")
+    });
+    RE_LOOP_FUNC.replace(function_name, "").to_string()
+}
+
+/// Remove parameters from monomorphised Cairo generics e.g. `<felt252>`.
+fn remove_monomorphization_suffix(function_name: &str) -> String {
+    static RE_MONOMORPHIZATION: LazyLock<Regex> = LazyLock::new(|| {
+        Regex::new(r"::<.*>")
+            .expect("Failed to create regex for normalizing monomorphized generic function names")
+    });
+    RE_MONOMORPHIZATION.replace(function_name, "").to_string()
+}

crates/debugging/src/trace/function/name.rs

@@ -0,0 +1,36 @@
+use crate::trace::function::name::FunctionName;
+
+/// Represents a node in the function call tree, where each node corresponds to a function name
+/// and can have child nodes representing nested function calls.
+#[derive(Debug, Clone)]
+pub struct FunctionNode {
+    pub value: FunctionName,
+    pub children: Vec<FunctionNode>,
+}
+
+impl FunctionNode {
+    /// Creates a new [`FunctionNode`] with the given [`FunctionName`].
+    pub fn new(value: FunctionName) -> Self {
+        FunctionNode {
+            value,
+            children: Vec::new(),
+        }
+    }
+
+    /// Adds a path of function names to the current node, creating child nodes as necessary.
+    pub fn add_path(&mut self, path: Vec<FunctionName>) {
+        self.add_path_recursive(&mut path.into_iter());
+    }
+
+    fn add_path_recursive(&mut self, iter: &mut impl Iterator<Item = FunctionName>) {
+        if let Some(next) = iter.next() {
+            if let Some(child) = self.children.iter_mut().find(|c| c.value == next) {
+                child.add_path_recursive(iter);
+            } else {
+                let mut new_child = FunctionNode::new(next);
+                new_child.add_path_recursive(iter);
+                self.children.push(new_child);
+            }
+        }
+    }
+}

crates/debugging/src/trace/function/node.rs

@@ -0,0 +1,50 @@
+use crate::trace::function::name::FunctionName;
+
+/// Represents a call stack for function calls, allowing to track the current function call
+pub struct CallStack {
+    stack: Vec<FunctionName>,
+    previous_stack_lengths: Vec<usize>,
+}
+
+impl CallStack {
+    /// Creates a new empty [`CallStack`]
+    pub fn new() -> Self {
+        Self {
+            stack: Vec::new(),
+            previous_stack_lengths: Vec::new(),
+        }
+    }
+
+    /// Enters a new function call by updating the stack with the new call stack.
+    /// It saves the current stack length to allow returning to it later.
+    ///
+    /// The New call stack is expected to be a prefix of the current stack.
+    pub fn enter_function_call(&mut self, new_call_stack: Vec<FunctionName>) {
+        self.previous_stack_lengths.push(self.stack.len());
+
+        self.stack = new_call_stack;
+    }
+
+    /// Exits the current function call by truncating the stack to the previous length.
+    /// If there is no previous length, it does nothing.
+    pub fn exit_function_call(&mut self) {
+        if let Some(previous_stack_len) = self.previous_stack_lengths.pop() {
+            self.stack.truncate(previous_stack_len);
+        }
+    }
+
+    /// Creates new stack with the given function name.
+    pub fn new_stack(&self, function_name: FunctionName) -> Vec<FunctionName> {
+        let mut stack = self.stack.clone();
+
+        let empty_or_different_function = self.stack.last().is_none_or(|current_function| {
+            current_function.function_name() != function_name.function_name()
+        });
+
+        if empty_or_different_function {
+            stack.push(function_name);
+        }
+
+        stack
+    }
+}