Refactor in preparation for inlining functions (#94)

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,8 @@ 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::{Function, InternalFunction, MeasurementUnit, MeasurementValue, Sample};
 
 #[derive(Clone, Copy, Hash, PartialEq, Eq, Debug)]
 struct StringId(u64);
@@ -66,7 +66,17 @@ impl ProfilerContext {
         }
     }
 
-    fn location_id(&mut self, location: &FunctionName) -> LocationId {
+    fn location_id(&mut self, location: &Function) -> LocationId {
+        let location = match location {
+            Function::Entrypoint(function_name)
+            | Function::InternalFunction(InternalFunction::NonInlined(function_name)) => {
+                function_name
+            }
+            Function::InternalFunction(InternalFunction::_Inlined(_)) => {
+                todo!("Unused, logic for it will be added in the next PR")
+            }
+        };
+
         if let Some(loc) = self.locations.get(location) {
             LocationId(loc.id)
         } else {
@@ -142,7 +152,7 @@ fn build_value_types(
 
 fn build_samples(
     context: &mut ProfilerContext,
-    samples: &[ContractCallSample],
+    samples: &[Sample],
     all_measurements_units: &[MeasurementUnit],
 ) -> Vec<pprof::Sample> {
     let samples = samples
@@ -171,13 +181,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,
         }
     }

crates/cairo-profiler/src/sierra_loader.rs

@@ -1,5 +1,5 @@
 use anyhow::{anyhow, Context, Result};
-use cairo_lang_sierra::program::{ProgramArtifact, VersionedProgram};
+use cairo_lang_sierra::program::{Program, VersionedProgram};
 use cairo_lang_sierra_to_casm::compiler::{CairoProgramDebugInfo, SierraToCasmConfig};
 use cairo_lang_sierra_to_casm::metadata::calc_metadata;
 use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
@@ -11,134 +11,143 @@ use trace_data::{CallTrace, CallTraceNode};
 
 /// Map with sierra and casm debug info needed for function level profiling.
 /// All paths in the map are absolute paths.
-pub struct CompiledArtifactsPathMap(HashMap<Utf8PathBuf, CompiledArtifacts>);
+pub struct CompiledArtifactsCache(HashMap<Utf8PathBuf, CompiledArtifacts>);
 
 pub struct CompiledArtifacts {
-    pub sierra: SierraProgramArtifact,
+    pub sierra_program: SierraProgram,
     pub casm_debug_info: CairoProgramDebugInfo,
 }
 
-pub enum SierraProgramArtifact {
-    VersionedProgram(ProgramArtifact),
-    ContractClass(ProgramArtifact),
+pub enum SierraProgram {
+    VersionedProgram(Program),
+    ContractClass(Program),
 }
 
-impl SierraProgramArtifact {
-    pub fn get_program_artifact(&self) -> &ProgramArtifact {
+impl SierraProgram {
+    pub fn get_program(&self) -> &Program {
         match self {
-            SierraProgramArtifact::VersionedProgram(program_artifact)
-            | SierraProgramArtifact::ContractClass(program_artifact) => program_artifact,
+            SierraProgram::VersionedProgram(program) | SierraProgram::ContractClass(program) => {
+                program
+            }
         }
     }
 }
 
-impl CompiledArtifactsPathMap {
+impl CompiledArtifactsCache {
     pub fn new() -> Self {
         Self(HashMap::new())
     }
 
-    pub fn compile_sierra_and_add_compiled_artifacts_to_map(
-        &mut self,
-        sierra_path: &Utf8Path,
-    ) -> Result<()> {
-        let absolute_sierra_path = sierra_path
-            .canonicalize_utf8()
-            .with_context(|| format!("Failed to canonicalize path: {sierra_path}"))?;
-
-        if !self.0.contains_key(&absolute_sierra_path) {
-            let raw_sierra = fs::read_to_string(&absolute_sierra_path)?;
-
-            if let Ok(contract_class) = serde_json::from_str::<ContractClass>(&raw_sierra) {
-                let program_artifact = ProgramArtifact {
-                    program: contract_class
-                        .extract_sierra_program()
-                        .context("Failed to extract sierra program from contract code")?,
-                    debug_info: contract_class.sierra_program_debug_info.clone(),
-                };
-
-                let (_casm_contract_class, casm_debug_info) =
-                    CasmContractClass::from_contract_class_with_debug_info(
-                        contract_class,
-                        false,
-                        usize::MAX,
-                    )?;
-
-                self.0.insert(
-                    absolute_sierra_path,
-                    CompiledArtifacts {
-                        sierra: SierraProgramArtifact::ContractClass(program_artifact),
-                        casm_debug_info,
-                    },
-                );
-
-                return Ok(());
-            }
+    pub fn get_compiled_artifacts_for_path(&self, path: &Utf8Path) -> &CompiledArtifacts {
+        self.0
+            .get(path)
+            .unwrap_or_else(|| panic!("Compiled artifacts not found for path {path}"))
+    }
+}
 
-            if let Ok(versioned_program) = serde_json::from_str::<VersionedProgram>(&raw_sierra) {
-                let program_artifact = versioned_program
-                    .into_v1()
-                    .context("Failed to extract program artifact from versioned program. Make sure your versioned program is of version 1")?;
-
-                let casm = cairo_lang_sierra_to_casm::compiler::compile(
-                    &program_artifact.program,
-                    &calc_metadata(&program_artifact.program, Default::default())
-                        .with_context(|| "Failed calculating Sierra variables.")?,
-                    SierraToCasmConfig {
-                        gas_usage_check: true,
-                        max_bytecode_size: usize::MAX,
-                    },
+pub fn compile_sierra_and_add_compiled_artifacts_to_cache(
+    sierra_path: &Utf8Path,
+    compiled_artifacts_cache: &mut CompiledArtifactsCache,
+) -> Result<()> {
+    let absolute_sierra_path = sierra_path
+        .canonicalize_utf8()
+        .with_context(|| format!("Failed to canonicalize path: {sierra_path}"))?;
+
+    if !compiled_artifacts_cache
+        .0
+        .contains_key(&absolute_sierra_path)
+    {
+        let raw_sierra = fs::read_to_string(&absolute_sierra_path)?;
+
+        if let Ok(contract_class) = serde_json::from_str::<ContractClass>(&raw_sierra) {
+            let program = contract_class
+                .extract_sierra_program()
+                .context("Failed to extract sierra program from contract code")?;
+
+            let contract_class = ContractClass {
+                // Debug info is unused in the compilation. This saves us a costly clone.
+                sierra_program_debug_info: None,
+                ..contract_class
+            };
+
+            let (_casm_contract_class, casm_debug_info) =
+                CasmContractClass::from_contract_class_with_debug_info(
+                    contract_class,
+                    false,
+                    usize::MAX,
                 )
-                .with_context(|| "Compilation failed.")?;
+                .context("Sierra -> CASM compilation failed.")?;
 
-                self.0.insert(
-                    absolute_sierra_path,
-                    CompiledArtifacts {
-                        sierra: SierraProgramArtifact::VersionedProgram(program_artifact),
-                        casm_debug_info: casm.debug_info,
-                    },
-                );
+            compiled_artifacts_cache.0.insert(
+                absolute_sierra_path,
+                CompiledArtifacts {
+                    sierra_program: SierraProgram::ContractClass(program),
+                    casm_debug_info,
+                },
+            );
 
-                return Ok(());
-            }
+            return Ok(());
+        }
 
-            return Err(anyhow!(
-                "Failed to deserialize sierra saved under path: {}",
-                absolute_sierra_path
-            ));
+        if let Ok(versioned_program) = serde_json::from_str::<VersionedProgram>(&raw_sierra) {
+            let program = versioned_program
+                .into_v1()
+                .context("Failed to extract program artifact from versioned program. Make sure your versioned program is of version 1")?.program;
+
+            let casm = cairo_lang_sierra_to_casm::compiler::compile(
+                &program,
+                &calc_metadata(&program, Default::default())
+                    .with_context(|| "Failed calculating Sierra variables.")?,
+                SierraToCasmConfig {
+                    gas_usage_check: true,
+                    max_bytecode_size: usize::MAX,
+                },
+            )
+            .context("Sierra -> CASM compilation failed.")?;
+
+            compiled_artifacts_cache.0.insert(
+                absolute_sierra_path,
+                CompiledArtifacts {
+                    sierra_program: SierraProgram::VersionedProgram(program),
+                    casm_debug_info: casm.debug_info,
+                },
+            );
+
+            return Ok(());
         }
 
-        Ok(())
+        return Err(anyhow!(
+            "Failed to deserialize sierra saved under path: {}",
+            absolute_sierra_path
+        ));
     }
 
-    pub fn get_sierra_casm_artifacts_for_path(&self, path: &Utf8Path) -> &CompiledArtifacts {
-        self.0
-            .get(path)
-            .unwrap_or_else(|| panic!("Compiled artifacts not found for path {path}"))
-    }
+    Ok(())
 }
 
 pub fn collect_and_compile_all_sierra_programs(
     trace: &CallTrace,
-) -> Result<CompiledArtifactsPathMap> {
-    let mut compiled_artifacts_path_map = CompiledArtifactsPathMap::new();
-    collect_compiled_artifacts(trace, &mut compiled_artifacts_path_map)?;
+) -> Result<CompiledArtifactsCache> {
+    let mut compiled_artifacts_cache = CompiledArtifactsCache::new();
+    collect_compiled_artifacts(trace, &mut compiled_artifacts_cache)?;
 
-    Ok(compiled_artifacts_path_map)
+    Ok(compiled_artifacts_cache)
 }
 
 fn collect_compiled_artifacts(
     trace: &CallTrace,
-    compiled_artifacts_path_map: &mut CompiledArtifactsPathMap,
+    compiled_artifacts_cache: &mut CompiledArtifactsCache,
 ) -> Result<()> {
     if let Some(cairo_execution_info) = &trace.cairo_execution_info {
-        compiled_artifacts_path_map.compile_sierra_and_add_compiled_artifacts_to_map(
+        compile_sierra_and_add_compiled_artifacts_to_cache(
             &cairo_execution_info.source_sierra_path,
+            compiled_artifacts_cache,
         )?;
     }
 
     for sub_trace_node in &trace.nested_calls {
         if let CallTraceNode::EntryPointCall(sub_trace) = sub_trace_node {
-            collect_compiled_artifacts(sub_trace, compiled_artifacts_path_map)?;
+            collect_compiled_artifacts(sub_trace, compiled_artifacts_cache)?;
         }
     }