--max-function-trace-depth flag (#62)

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `--max-function-trace-depth` allowing to specify maximum depth of the function tree in function level profiling
+
 ## [0.3.0] - 2024-04-20
 
 ### Added

crates/cairo-profiler/src/main.rs

@@ -32,6 +32,13 @@ struct Cli {
  /// Show contract addresses and function selectors in a trace tree
  #[arg(long)]
  show_details: bool,
+
+ /// Specify maximum depth of function tree in function level profiling.
+ /// The is applied per entrypoint - each entrypoint function tree is treated separately.
+ /// Keep in mind recursive functions are also taken into account even though they are later
+ /// aggregated to one function call.
+ #[arg(long, default_value_t = 100)]
+ max_function_trace_depth: usize,
 }
 
 fn main() -> Result<()> {
@@ -47,6 +54,7 @@ fn main() -> Result<()> {
  &serialized_trace,
  &compiled_artifacts_path_map,
  cli.show_details,
+ cli.max_function_trace_depth,
  )?;
 
  let profile = build_profile(&samples);

crates/cairo-profiler/src/trace_reader.rs

@@ -156,6 +156,7 @@ pub fn collect_samples_from_trace(
  trace: &CallTrace,
  compiled_artifacts_path_map: &CompiledArtifactsPathMap,
  show_details: bool,
+ max_function_trace_depth: usize,
 ) -> Result<Vec<ContractCallSample>> {
  let mut samples = vec![];
  let mut current_path = vec![];
@@ -166,6 +167,7 @@ pub fn collect_samples_from_trace(
  trace,
  compiled_artifacts_path_map,
  show_details,
+ max_function_trace_depth,
  )?;
  Ok(samples)
 }
@@ -176,6 +178,7 @@ fn collect_samples<'a>(
  trace: &'a CallTrace,
  compiled_artifacts_path_map: &CompiledArtifactsPathMap,
  show_details: bool,
+ max_function_trace_depth: usize,
 ) -> Result<&'a ExecutionResources> {
  current_call_stack.push(EntryPointId::from(
  trace.entry_point.contract_name.clone(),
@@ -204,6 +207,7 @@ fn collect_samples<'a>(
  compiled_artifacts.sierra.get_program_artifact(),
  &compiled_artifacts.casm_debug_info,
  compiled_artifacts.sierra.was_run_with_header(),
+ max_function_trace_depth,
  )?;
 
  for mut function_stack_trace in profiling_info.functions_stack_traces {
@@ -236,6 +240,7 @@ fn collect_samples<'a>(
  sub_trace,
  compiled_artifacts_path_map,
  show_details,
+ max_function_trace_depth,
  )?;
  }
  }

crates/cairo-profiler/src/trace_reader/function_trace_builder.rs

@@ -10,8 +10,6 @@ use std::collections::HashMap;
 use std::ops::AddAssign;
 use trace_data::TraceEntry;
 
-const MAX_TRACE_DEPTH: u8 = 100;
-
 pub struct ProfilingInfo {
  pub functions_stack_traces: Vec<FunctionStackTrace>,
  pub header_steps: Steps,
@@ -52,6 +50,7 @@ pub fn collect_profiling_info(
  program_artifact: &ProgramArtifact,
  casm_debug_info: &CairoProgramDebugInfo,
  was_run_with_header: bool,
+ max_function_trace_depth: usize,
 ) -> Result<ProfilingInfo> {
  let program = &program_artifact.program;
  let sierra_program_registry = &ProgramRegistry::<CoreType, CoreLibfunc>::new(program).unwrap();
@@ -137,7 +136,7 @@ pub fn collect_profiling_info(
  Ok(CoreConcreteLibfunc::FunctionCall(_))
  ) {
  // Push to the stack.
- if function_stack_depth < MAX_TRACE_DEPTH as usize {
+ if function_stack_depth < max_function_trace_depth {
  function_stack.push((user_function_idx, current_function_steps));
  current_function_steps = Steps(0);
  }
@@ -146,7 +145,7 @@ pub fn collect_profiling_info(
  }
  GenStatement::Return(_) => {
  // Pop from the stack.
- if function_stack_depth <= MAX_TRACE_DEPTH as usize {
+ if function_stack_depth <= max_function_trace_depth {
  let current_stack =
  chain!(function_stack.iter().map(|f| f.0), [user_function_idx])
  .collect_vec();

crates/cairo-profiler/tests/e2e.rs

@@ -23,6 +23,7 @@ fn output_path() {
  assert!(temp_dir.join("my/output/dir/my_file.pb.gz").exists());
 }
 
+#[test_case(&["call.json", "--max-function-trace-depth", "5"]; "with max function trace depth")]
 #[test_case(&["call.json", "--show-details"]; "with details")]
 #[test_case(&["call.json"]; "without details")]
 fn simple_package(args: &[&str]) {