Build inlined functions in pprof

CHANGELOG.md

@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- `--max-function-trace-depth` allowing to specify maximum depth of the function tree in function level profiling
+- `--max-function-stack-trace-depth` allowing to specify maximum depth of the function tree in function level profiling
 - `--split-generics` flag allowing to differentiate between non-inlined generics monomorphised with different types
 
 ## [0.3.0] - 2024-04-20

crates/cairo-profiler/src/main.rs

@@ -38,7 +38,7 @@ struct Cli {
     /// Specify maximum depth of function tree in function level profiling.
     /// The is applied per entrypoint - each entrypoint function tree is treated separately.
     #[arg(long, default_value_t = 100)]
-    max_function_trace_depth: usize,
+    max_function_stack_trace_depth: usize,
 
     /// Split non-inlined generic functions based on the type they were monomorphised with.
     /// E.g. treat `function<felt252>` as different from `function<u8>`.
@@ -54,10 +54,10 @@ fn main() -> Result<()> {
     let serialized_trace: CallTrace =
         serde_json::from_str(&data).context("Failed to deserialize call trace")?;
 
-    let compiled_artifacts_path_map = collect_and_compile_all_sierra_programs(&serialized_trace)?;
+    let compiled_artifacts_cache = collect_and_compile_all_sierra_programs(&serialized_trace)?;
     let samples = collect_samples_from_trace(
         &serialized_trace,
-        &compiled_artifacts_path_map,
+        &compiled_artifacts_cache,
         &ProfilerConfig::from(&cli),
     )?;
 

crates/cairo-profiler/src/profile_builder.rs

@@ -10,8 +10,10 @@ use std::collections::{HashMap, HashSet};
 
 pub use perftools::profiles as pprof;
 
-use crate::trace_reader::functions::FunctionName;
-use crate::trace_reader::{ContractCallSample, MeasurementUnit, MeasurementValue};
+use crate::trace_reader::function_name::FunctionName;
+use crate::trace_reader::sample::{
+    Function, InternalFunction, MeasurementUnit, MeasurementValue, Sample,
+};
 
 #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
 struct StringId(u64);
@@ -43,7 +45,7 @@ impl From<FunctionId> for u64 {
 struct ProfilerContext {
     strings: HashMap<String, StringId>,
     functions: HashMap<FunctionName, pprof::Function>,
-    locations: HashMap<FunctionName, pprof::Location>,
+    locations: HashMap<Vec<Function>, pprof::Location>,
 }
 
 impl ProfilerContext {
@@ -66,25 +68,76 @@ impl ProfilerContext {
         }
     }
 
-    fn location_id(&mut self, location: &FunctionName) -> LocationId {
-        if let Some(loc) = self.locations.get(location) {
-            LocationId(loc.id)
-        } else {
-            let line = pprof::Line {
-                function_id: self.function_id(location).into(),
-                line: 0,
-            };
-            let location_data = pprof::Location {
-                id: (self.locations.len() + 1) as u64,
-                mapping_id: 0,
-                address: 0,
-                line: vec![line],
-                is_folded: true,
-            };
-
-            self.locations.insert(location.clone(), location_data);
-            LocationId(self.locations.len() as u64)
+    fn locations_ids(&mut self, call_stack: &[Function]) -> Vec<LocationId> {
+        let mut locations_ids = vec![];
+        // This vector represent stacks of functions corresponding to single locations.
+        // It contains tuples of form (start_index, end_index).
+        // A single stack is `&call_stack[start_index..=end_index]`.
+        let mut function_stacks_indexes = vec![];
+
+        let mut current_function_stack_start_index = 0;
+        for (index, function) in call_stack.iter().enumerate() {
+            match function {
+                Function::InternalFunction(InternalFunction::NonInlined(_))
+                | Function::Entrypoint(_) => {
+                    if index != 0 {
+                        function_stacks_indexes
+                            .push((current_function_stack_start_index, index - 1));
+                    }
+                    current_function_stack_start_index = index;
+                }
+                Function::InternalFunction(InternalFunction::Inlined(_)) => {}
+            }
         }
+        function_stacks_indexes.push((current_function_stack_start_index, call_stack.len() - 1));
+
+        for (start_index, end_index) in function_stacks_indexes {
+            let function_stack = &call_stack[start_index..=end_index];
+            if let Some(location) = self.locations.get(function_stack) {
+                locations_ids.push(LocationId(location.id));
+            } else {
+                let mut location = match &function_stack[0] {
+                    Function::Entrypoint(function_name) | Function::InternalFunction(InternalFunction::NonInlined(function_name)) => {
+                        let line = pprof::Line {
+                            function_id: self.function_id(function_name).into(),
+                            line: 0,
+                        };
+                         pprof::Location {
+                            id: (self.locations.len() + 1) as u64,
+                            mapping_id: 0,
+                            address: 0,
+                            line: vec![line],
+                            is_folded: true,
+                        }
+                    }
+                    Function::InternalFunction(InternalFunction::Inlined(_)) => unreachable!("First function in a call stack corresponding to a single location cannot be inlined")
+                };
+
+                for function in function_stack.get(1..).unwrap_or_default() {
+                    match function {
+                        Function::InternalFunction(InternalFunction::Inlined(function_name)) => {
+                            let line = pprof::Line {
+                                function_id: self.function_id(function_name).into(),
+                                line: 0,
+                            };
+                            location.line.push(line);
+                        }
+                        Function::Entrypoint(_)
+                        | Function::InternalFunction(InternalFunction::NonInlined(_)) => {
+                            unreachable!("Only first function in a call stack corresponding to a single location can be not inlined")
+                        }
+                    }
+                }
+
+                // pprof format represents callstack from the least meaningful elements
+                location.line.reverse();
+                locations_ids.push(LocationId(location.id));
+
+                self.locations.insert(function_stack.to_vec(), location);
+            }
+        }
+
+        locations_ids
     }
 
     fn function_id(&mut self, function_name: &FunctionName) -> FunctionId {
@@ -110,7 +163,6 @@ impl ProfilerContext {
         }
 
         let functions = self.functions.into_values().collect();
-
         let locations = self.locations.into_values().collect();
 
         (string_table, functions, locations)
@@ -142,16 +194,16 @@ fn build_value_types(
 
 fn build_samples(
     context: &mut ProfilerContext,
-    samples: &[ContractCallSample],
+    samples: &[Sample],
     all_measurements_units: &[MeasurementUnit],
 ) -> Vec<pprof::Sample> {
     let samples = samples
         .iter()
         .map(|s| pprof::Sample {
-            location_id: s
-                .call_stack
-                .iter()
-                .map(|loc| context.location_id(loc).into())
+            location_id: context
+                .locations_ids(&s.call_stack)
+                .into_iter()
+                .map(Into::into)
                 .rev() // pprof format represents callstack from the least meaningful element
                 .collect(),
             value: all_measurements_units
@@ -171,13 +223,13 @@ fn build_samples(
     samples
 }
 
-fn collect_all_measurements_units(samples: &[ContractCallSample]) -> Vec<MeasurementUnit> {
+fn collect_all_measurements_units(samples: &[Sample]) -> Vec<MeasurementUnit> {
     let units_set: HashSet<&MeasurementUnit> =
         samples.iter().flat_map(|m| m.measurements.keys()).collect();
     units_set.into_iter().cloned().collect()
 }
 
-pub fn build_profile(samples: &[ContractCallSample]) -> pprof::Profile {
+pub fn build_profile(samples: &[Sample]) -> pprof::Profile {
     let mut context = ProfilerContext::new();
     let all_measurements_units = collect_all_measurements_units(samples);
     let value_types = build_value_types(&all_measurements_units, &mut context);

crates/cairo-profiler/src/profiler_config.rs

@@ -2,29 +2,29 @@ use crate::Cli;
 
 pub struct ProfilerConfig {
     pub show_details: bool,
-    pub max_function_trace_depth: usize,
+    pub max_function_stack_trace_depth: usize,
     pub split_generics: bool,
 }
 
 impl From<&Cli> for ProfilerConfig {
     fn from(cli: &Cli) -> ProfilerConfig {
         ProfilerConfig {
             show_details: cli.show_details,
-            max_function_trace_depth: cli.max_function_trace_depth,
+            max_function_stack_trace_depth: cli.max_function_stack_trace_depth,
             split_generics: cli.split_generics,
         }
     }
 }
 
 pub struct FunctionLevelConfig {
-    pub max_function_trace_depth: usize,
+    pub max_function_stack_trace_depth: usize,
     pub split_generics: bool,
 }
 
 impl From<&ProfilerConfig> for FunctionLevelConfig {
     fn from(profiler_config: &ProfilerConfig) -> FunctionLevelConfig {
         FunctionLevelConfig {
-            max_function_trace_depth: profiler_config.max_function_trace_depth,
+            max_function_stack_trace_depth: profiler_config.max_function_stack_trace_depth,
             split_generics: profiler_config.split_generics,
         }
     }