Add function level trace

commit-id:3a7228bd

Author: ksew1

Cargo.lock

@@ -8,14 +8,16 @@ 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
+thiserror.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, FunctionTraceError};
 use crate::trace::types::{
     CallerAddress, ContractAddress, ContractName, ContractTrace, Selector, TestName, TraceInfo,
     TransformedCallResult, TransformedCalldata,
@@ -38,6 +39,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 +54,7 @@ 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)),
+            function_trace: components.function_trace_lazy(|| self.collect_function_trace()),
         };
 
         ContractTrace {
@@ -125,6 +128,14 @@ impl<'a> Collector<'a> {
         })
     }
 
+    fn collect_function_trace(&self) -> Result<FunctionTrace, FunctionTraceError> {
+        FunctionTrace::create(
+            *self.class_hash(),
+            self.call_trace,
+            self.contracts_data_store(),
+        )
+    }
+
     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, FunctionTraceError};
 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, Result<FunctionTrace, FunctionTraceError>);

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, FunctionTraceError};

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
+    }
+}

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

@@ -0,0 +1,131 @@
+use crate::contracts_data_store::ContractsDataStore;
+use crate::trace::function::name::FunctionName;
+use crate::trace::function::node::FunctionNode;
+use crate::trace::function::stack::CallStack;
+use cairo_annotations::map_pcs_to_sierra_statement_ids;
+use cairo_annotations::trace_data::{CasmLevelInfo, TraceEntry};
+use cairo_lang_sierra::extensions::core::{CoreConcreteLibfunc, CoreLibfunc, CoreType};
+use cairo_lang_sierra::program::{GenStatement, StatementIdx};
+use cairo_lang_sierra::program_registry::ProgramRegistry;
+use cheatnet::state::CallTrace;
+use starknet_api::core::ClassHash;
+
+#[derive(Debug, Clone, thiserror::Error)]
+pub enum FunctionTraceError {
+    #[error("function trace is not supported for forked contracts")]
+    ForkContract,
+}
+/// `FunctionTrace` represents a trace of function calls in a Sierra program.
+/// It captures the structure of function calls and returns, allowing for analysis of the
+/// execution flow within a contract.
+#[derive(Debug, Clone)]
+pub struct FunctionTrace {
+    pub root: FunctionNode,
+}
+
+impl FunctionTrace {
+    /// Creates a new [`FunctionTrace`] from the given [`ClassHash`], [`CallTrace`], and [`ContractsDataStore`].
+    pub fn create(
+        class_hash: ClassHash,
+        call_trace: &CallTrace,
+        contracts_data_store: &ContractsDataStore,
+    ) -> Result<Self, FunctionTraceError> {
+        if contracts_data_store.is_fork(&class_hash) {
+            return Err(FunctionTraceError::ForkContract);
+        }
+        let program_artifact = contracts_data_store
+            .get_program_artifact(&class_hash)
+            .expect("program artifact should be present for local contracts");
+
+        let program = &program_artifact.program;
+
+        let sierra_program_registry = ProgramRegistry::<CoreType, CoreLibfunc>::new(program)
+            .expect("failed to create Sierra program registry");
+
+        let mut call_stack = CallStack::new();
+        let mut stacks = Vec::new();
+
+        for statement_idx in get_sierra_statements(class_hash, call_trace, contracts_data_store) {
+            let function_name = FunctionName::from_program(statement_idx, program);
+
+            let stack = call_stack.new_stack(function_name);
+
+            stacks.push(stack.clone());
+
+            let statement = program
+                .statements
+                .get(statement_idx.0)
+                .expect("statement should be present");
+
+            match statement {
+                GenStatement::Invocation(invocation) => {
+                    let libfunc = sierra_program_registry.get_libfunc(&invocation.libfunc_id);
+
+                    if let Ok(CoreConcreteLibfunc::FunctionCall(_)) = &libfunc {
+                        call_stack.enter_function_call(stack);
+                    }
+                }
+                GenStatement::Return(_) => {
+                    call_stack.exit_function_call();
+                }
+            }
+        }
+
+        Ok(build_function_trace(stacks))
+    }
+}
+
+/// Retrieves vector of [`StatementIdx`] from the given [`CallTrace`].
+fn get_sierra_statements(
+    class_hash: ClassHash,
+    call_trace: &CallTrace,
+    contracts_data_store: &ContractsDataStore,
+) -> Vec<StatementIdx> {
+    let casm_debug_info = contracts_data_store
+        .get_casm_debug_info(&class_hash)
+        .expect("Cairo program debug info should be present");
+
+    let casm_level_info = build_casm_level_info(call_trace);
+
+    map_pcs_to_sierra_statement_ids(casm_debug_info, &casm_level_info)
+        .into_iter()
+        .filter_map(Option::from)
+        .collect()
+}
+
+/// Builds a [`CasmLevelInfo`] from the given [`CallTrace`].
+fn build_casm_level_info(call_trace: &CallTrace) -> CasmLevelInfo {
+    let vm_trace = call_trace
+        .vm_trace
+        .as_ref()
+        .expect("vm trace should be present")
+        .iter()
+        .map(|value| TraceEntry {
+            pc: value.pc,
+            ap: value.ap,
+            fp: value.fp,
+        })
+        .collect();
+
+    CasmLevelInfo {
+        run_with_call_header: false,
+        vm_trace,
+    }
+}
+
+/// Builds a [`FunctionTrace`] from the given stacks of function names.
+fn build_function_trace(stacks: Vec<Vec<FunctionName>>) -> FunctionTrace {
+    let mut dummy_root = FunctionNode::new(FunctionName::NonInlined("root".to_string()));
+
+    for stack in stacks {
+        dummy_root.add_path(stack);
+    }
+
+    assert_eq!(dummy_root.children.len(), 1);
+    let real_root = dummy_root
+        .children
+        .pop()
+        .unwrap_or_else(|| unreachable!("assertion above should have prevented this"));
+
+    FunctionTrace { root: real_root }
+}