diff --git a/Cargo.lock b/Cargo.lock index 951f2f568f..43f8c06869 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2071,6 +2071,7 @@ name = "debugging" version = "1.0.0" dependencies = [ "blockifier", + "cairo-annotations", "cairo-lang-sierra", "cairo-lang-sierra-to-casm", "cairo-lang-starknet-classes", @@ -2080,6 +2081,7 @@ dependencies = [ "paste", "ptree", "rayon", + "regex", "serde_json", "starknet", "starknet-types-core", diff --git a/crates/debugging/Cargo.toml b/crates/debugging/Cargo.toml index 7d3b78d1db..e133287cb0 100644 --- a/crates/debugging/Cargo.toml +++ b/crates/debugging/Cargo.toml @@ -8,12 +8,14 @@ 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 diff --git a/crates/debugging/src/contracts_data_store.rs b/crates/debugging/src/contracts_data_store.rs index 07e442d3ee..494032052c 100644 --- a/crates/debugging/src/contracts_data_store.rs +++ b/crates/debugging/src/contracts_data_store.rs @@ -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`] diff --git a/crates/debugging/src/trace/collect.rs b/crates/debugging/src/trace/collect.rs index 73e0f98ce5..9637a24246 100644 --- a/crates/debugging/src/trace/collect.rs +++ b/crates/debugging/src/trace/collect.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::create( + *self.class_hash(), + self.call_trace, + self.contracts_data_store(), + ) + } + fn class_hash(&self) -> &ClassHash { self.call_trace .entry_point diff --git a/crates/debugging/src/trace/components.rs b/crates/debugging/src/trace/components.rs index fbaec149b2..18642f4d12 100644 --- a/crates/debugging/src/trace/components.rs +++ b/crates/debugging/src/trace/components.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); diff --git a/crates/debugging/src/trace/function/mod.rs b/crates/debugging/src/trace/function/mod.rs new file mode 100644 index 0000000000..ad4b5f6498 --- /dev/null +++ b/crates/debugging/src/trace/function/mod.rs @@ -0,0 +1,7 @@ +mod name; +mod node; +mod stack; +mod trace; + +pub use node::FunctionNode; +pub use trace::{FunctionTrace, FunctionTraceError}; diff --git a/crates/debugging/src/trace/function/name.rs b/crates/debugging/src/trace/function/name.rs new file mode 100644 index 0000000000..018fdcf49a --- /dev/null +++ b/crates/debugging/src/trace/function/name.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 = 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. ``. +fn remove_monomorphization_suffix(function_name: &str) -> String { + static RE_MONOMORPHIZATION: LazyLock = LazyLock::new(|| { + Regex::new(r"::<.*>") + .expect("Failed to create regex for normalizing monomorphized generic function names") + }); + RE_MONOMORPHIZATION.replace(function_name, "").to_string() +} diff --git a/crates/debugging/src/trace/function/node.rs b/crates/debugging/src/trace/function/node.rs new file mode 100644 index 0000000000..7250d55c2e --- /dev/null +++ b/crates/debugging/src/trace/function/node.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, +} + +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) { + self.add_path_recursive(&mut path.into_iter()); + } + + fn add_path_recursive(&mut self, iter: &mut impl Iterator) { + 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); + } + } + } +} diff --git a/crates/debugging/src/trace/function/stack.rs b/crates/debugging/src/trace/function/stack.rs new file mode 100644 index 0000000000..b4bc687070 --- /dev/null +++ b/crates/debugging/src/trace/function/stack.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, + previous_stack_lengths: Vec, +} + +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) { + 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 { + 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 + } +} diff --git a/crates/debugging/src/trace/function/trace.rs b/crates/debugging/src/trace/function/trace.rs new file mode 100644 index 0000000000..34c0dc9d6e --- /dev/null +++ b/crates/debugging/src/trace/function/trace.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 { + 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::::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 { + 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>) -> FunctionTrace { + let mut placeholder_root = FunctionNode::new(FunctionName::NonInlined("root".to_string())); + + for stack in stacks { + placeholder_root.add_path(stack); + } + + assert_eq!(placeholder_root.children.len(), 1); + let real_root = placeholder_root + .children + .pop() + .unwrap_or_else(|| unreachable!("assertion above should have prevented this")); + + FunctionTrace { root: real_root } +} diff --git a/crates/debugging/src/trace/mod.rs b/crates/debugging/src/trace/mod.rs index 104ed39ce2..ccafe8650d 100644 --- a/crates/debugging/src/trace/mod.rs +++ b/crates/debugging/src/trace/mod.rs @@ -1,4 +1,5 @@ pub mod collect; pub mod components; pub mod context; +pub mod function; pub mod types; diff --git a/crates/debugging/src/trace/types.rs b/crates/debugging/src/trace/types.rs index b62128256e..ac3f85dd96 100644 --- a/crates/debugging/src/trace/types.rs +++ b/crates/debugging/src/trace/types.rs @@ -3,6 +3,7 @@ use crate::trace::collect::Collector; use crate::trace::components::{ CallResultContainer, CallTypeContainer, CalldataContainer, CallerAddressContainer, ContractAddressContainer, ContractNameContainer, EntryPointTypeContainer, + FunctionTraceContainer, }; use crate::tree::TreeSerialize; use cheatnet::state::CallTrace; @@ -32,6 +33,7 @@ pub struct TraceInfo { pub call_type: CallTypeContainer, pub nested_calls: Vec, pub call_result: CallResultContainer, + pub function_trace: FunctionTraceContainer, } #[derive(Debug, Clone)] diff --git a/crates/debugging/src/tree/building/node.rs b/crates/debugging/src/tree/building/node.rs index b339460646..9e3c7de7e4 100644 --- a/crates/debugging/src/tree/building/node.rs +++ b/crates/debugging/src/tree/building/node.rs @@ -28,6 +28,13 @@ impl<'a> Node<'a> { tree_item.as_tree_node(self); } + /// Calls [`AsTreeNode::as_tree_node`] on the given item if it is not `None`. + pub fn as_tree_node_optional(&mut self, tree_item: Option<&impl AsTreeNode>) { + if let Some(tree_item) = tree_item { + self.as_tree_node(tree_item); + } + } + /// Creates a child node which parent is the current node and returns handle to created node. #[must_use = "if you want to create a leaf node use leaf() instead"] pub fn child_node(&mut self, tree_item: &impl NodeDisplay) -> Node<'_> { diff --git a/crates/debugging/src/tree/ui/as_tree_node.rs b/crates/debugging/src/tree/ui/as_tree_node.rs index b9cd5b5bef..2736fe4b96 100644 --- a/crates/debugging/src/tree/ui/as_tree_node.rs +++ b/crates/debugging/src/tree/ui/as_tree_node.rs @@ -1,3 +1,4 @@ +use crate::trace::function::{FunctionNode, FunctionTrace, FunctionTraceError}; use crate::trace::types::{ContractTrace, Trace, TraceInfo}; use crate::tree::building::node::Node; @@ -34,8 +35,27 @@ impl AsTreeNode for TraceInfo { parent.leaf_optional(self.caller_address.as_option()); parent.leaf_optional(self.call_type.as_option()); parent.leaf_optional(self.call_result.as_option()); + parent.as_tree_node_optional(self.function_trace.as_option()); for nested_call in &self.nested_calls { parent.as_tree_node(nested_call); } } } + +impl AsTreeNode for Result { + fn as_tree_node(&self, parent: &mut Node) { + match self { + Ok(trace) => parent.child_node(trace).as_tree_node(&trace.root), + Err(error) => parent.leaf(error), + } + } +} + +impl AsTreeNode for FunctionNode { + fn as_tree_node(&self, parent: &mut Node) { + let mut node = parent.child_node(self); + for child in &self.children { + node.as_tree_node(child); + } + } +} diff --git a/crates/debugging/src/tree/ui/display.rs b/crates/debugging/src/tree/ui/display.rs index 7ae299c63f..0447bf39f1 100644 --- a/crates/debugging/src/tree/ui/display.rs +++ b/crates/debugging/src/tree/ui/display.rs @@ -1,3 +1,4 @@ +use crate::trace::function::{FunctionNode, FunctionTrace, FunctionTraceError}; use crate::trace::types::{ CallerAddress, ContractAddress, ContractName, Selector, TestName, TransformedCallResult, TransformedCalldata, @@ -17,7 +18,11 @@ pub trait NodeDisplay { fn display(&self) -> String { let tag = console::style(Self::TAG).magenta(); let content = self.string_pretty(); - format!("[{tag}] {content}") + if content.is_empty() { + format!("[{tag}]") + } else { + format!("[{tag}] {content}") + } } } @@ -84,6 +89,27 @@ impl NodeDisplay for TransformedCallResult { } } +impl NodeDisplay for FunctionTrace { + const TAG: &'static str = "function call tree"; + fn string_pretty(&self) -> String { + String::new() + } +} + +impl NodeDisplay for FunctionTraceError { + const TAG: &'static str = "function trace error"; + fn string_pretty(&self) -> String { + self.to_string() + } +} + +impl NodeDisplay for FunctionNode { + const TAG: &'static str = "non inlined"; + fn string_pretty(&self) -> String { + self.value.function_name().to_string() + } +} + /// Helper function to get hex representation /// of a type that can be converted to a [`Felt`]. fn string_hex(data: impl Into) -> String {