Review changes + detect syscalls by invocation, not counter

crates/cairo-profiler/src/main.rs

@@ -6,7 +6,7 @@ use std::{
 use crate::profiler_config::ProfilerConfig;
 use crate::sierra_loader::collect_and_compile_all_sierra_programs;
 use crate::trace_reader::collect_samples_from_trace;
-use crate::trace_reader::syscall::read_and_parse_versioned_constants_file;
+use crate::versioned_constants_reader::read_and_parse_versioned_constants_file;
 use anyhow::{Context, Result};
 use bytes::{Buf, BytesMut};
 use camino::Utf8PathBuf;
@@ -20,6 +20,7 @@ mod profile_builder;
 mod profiler_config;
 mod sierra_loader;
 mod trace_reader;
+mod versioned_constants_reader;
 
 #[derive(Parser, Debug)]
 #[command(version)]
@@ -51,6 +52,10 @@ struct Cli {
     #[arg(long)]
     show_inlined_functions: bool,
 
+    /// Path to `versioned_constants.json` file, that includes resource cost map
+    /// If not provided, 0.13.2.1 version will be used.
+    /// Versioned files can be found in sequencer repo:
+    /// <https://github.com/starkware-libs/sequencer/blob/main/crates/blockifier/resources/>
     #[arg(long)]
     versioned_constants_path: Option<Utf8PathBuf>,
 }
@@ -60,6 +65,8 @@ fn main() -> Result<()> {
 
     let data = fs::read_to_string(&cli.path_to_trace_data)
         .context("Failed to read call trace from a file")?;
+    let os_resources_map = read_and_parse_versioned_constants_file(&cli.versioned_constants_path)
+        .context("Failed to parse versioned constants file")?;
     let serialized_trace: CallTrace =
         serde_json::from_str(&data).context("Failed to deserialize call trace")?;
 
@@ -76,9 +83,6 @@ fn main() -> Result<()> {
         );
     }
 
-    let os_resources_map = read_and_parse_versioned_constants_file(&cli.versioned_constants_path)
-        .expect("Failed to parse versioned constants file");
-
     let samples = collect_samples_from_trace(
         &serialized_trace,
         &compiled_artifacts_cache,

crates/cairo-profiler/src/trace_reader.rs

@@ -7,14 +7,13 @@ use crate::trace_reader::function_name::FunctionName;
 use crate::trace_reader::function_trace_builder::collect_function_level_profiling_info;
 
 use crate::trace_reader::sample::{FunctionCall, Sample};
-use crate::trace_reader::syscall::collect_syscall_sample;
 
-use trace_data::{CallTrace, CallTraceNode, ExecutionResources, OsResources};
+use crate::versioned_constants_reader::OsResources;
+use trace_data::{CallTrace, CallTraceNode, ExecutionResources};
 
 pub mod function_name;
 mod function_trace_builder;
 pub mod sample;
-pub mod syscall;
 
 pub fn collect_samples_from_trace(
     trace: &CallTrace,
@@ -76,6 +75,7 @@ fn collect_samples<'a>(
             cairo_execution_info.casm_level_info.run_with_call_header,
             &compiled_artifacts.statements_functions_map,
             &FunctionLevelConfig::from(profiler_config),
+            os_resources_map,
         );
 
         let mut function_samples = function_level_profiling_info
@@ -125,19 +125,6 @@ fn collect_samples<'a>(
         &trace.used_l1_resources,
     ));
 
-    call_resources
-        .syscall_counter
-        .iter()
-        .filter(|(_, count)| **count != 0)
-        .for_each(|(syscall, count)| {
-            samples.push(collect_syscall_sample(
-                current_entrypoint_call_stack.clone(),
-                *syscall,
-                *count,
-                os_resources_map,
-            ));
-        });
-
     current_entrypoint_call_stack.pop();
 
     Ok(&trace.cumulative_resources)

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

@@ -8,6 +8,7 @@ use crate::trace_reader::function_trace_builder::inlining::build_original_call_s
 use crate::trace_reader::sample::{
     FunctionCall, InternalFunctionCall, MeasurementUnit, MeasurementValue, Sample,
 };
+use crate::versioned_constants_reader::{map_syscall_name_to_selector, OsResources};
 use cairo_lang_sierra::extensions::core::{CoreConcreteLibfunc, CoreLibfunc, CoreType};
 use cairo_lang_sierra::program::{GenStatement, Program, StatementIdx};
 use cairo_lang_sierra::program_registry::ProgramRegistry;
@@ -53,6 +54,7 @@ pub fn collect_function_level_profiling_info(
     run_with_call_header: bool,
     statements_functions_map: &Option<StatementsFunctionsMap>,
     function_level_config: &FunctionLevelConfig,
+    os_resources_map: &OsResources,
 ) -> FunctionLevelProfilingInfo {
     let sierra_program_registry = &ProgramRegistry::<CoreType, CoreLibfunc>::new(program).unwrap();
 
@@ -81,6 +83,10 @@ pub fn collect_function_level_profiling_info(
     // the number of total steps. The value is different from zero only for functions run with header.
     let mut header_steps = Steps(0);
     let mut end_of_program_reached = false;
+    // Syscalls cannot be mapped using pc offsets
+    // They can be recognised by GenStatement::Invocation but they do not have GenStatement::Return
+    // That's why we must track entry to a syscall, and leave as soon as we're out of given GenStatement::Invocation
+    let mut in_syscall = false;
 
     for step in trace {
         // Skip the header.
@@ -115,10 +121,11 @@ pub fn collect_function_level_profiling_info(
 
         let current_call_stack = build_current_call_stack(
             &call_stack,
-            current_function_name,
+            current_function_name.clone(),
             function_level_config.show_inlined_functions,
             sierra_statement_idx,
             statements_functions_map.as_ref(),
+            in_syscall,
         );
 
         *functions_stack_traces
@@ -131,11 +138,42 @@ pub fn collect_function_level_profiling_info(
 
         match gen_statement {
             GenStatement::Invocation(invocation) => {
-                if matches!(
-                    sierra_program_registry.get_libfunc(&invocation.libfunc_id),
-                    Ok(CoreConcreteLibfunc::FunctionCall(_))
-                ) {
-                    call_stack.enter_function_call(current_call_stack);
+                match sierra_program_registry.get_libfunc(&invocation.libfunc_id) {
+                    Ok(CoreConcreteLibfunc::FunctionCall(_)) => {
+                        call_stack.enter_function_call(current_call_stack);
+                    }
+                    Ok(CoreConcreteLibfunc::StarkNet(_)) => {
+                        if invocation.libfunc_id.debug_name.is_none() {
+                            // this libfunc is not included in the artifact file
+                            // it is likely to be a libfunc from the test itself
+                            continue;
+                        }
+
+                        if !in_syscall {
+                            in_syscall = true;
+                            let mut new_current_call_stack = current_call_stack.clone();
+                            new_current_call_stack.push(FunctionCall::InternalFunctionCall(
+                                InternalFunctionCall::Syscall(FunctionName(
+                                    invocation
+                                        .libfunc_id
+                                        .debug_name
+                                        .clone()
+                                        .unwrap()
+                                        .to_string(),
+                                )),
+                            ));
+
+                            call_stack.enter_function_call(new_current_call_stack);
+                        }
+                    }
+                    _ => {
+                        // If we were in a syscall this is the time we go out of it, as pcs no longer
+                        // belong to GenStatement::Invocation of CoreConcreteLibfunc::StarkNet
+                        if in_syscall {
+                            call_stack.exit_function_call();
+                            in_syscall = false;
+                        }
+                    }
                 }
             }
             GenStatement::Return(_) => {
@@ -146,16 +184,7 @@ pub fn collect_function_level_profiling_info(
         }
     }
 
-    let functions_samples = functions_stack_traces
-        .into_iter()
-        .map(|(call_stack, steps)| Sample {
-            call_stack,
-            measurements: HashMap::from([(
-                MeasurementUnit::from("steps".to_string()),
-                MeasurementValue(i64::try_from(steps.0).unwrap()),
-            )]),
-        })
-        .collect_vec();
+    let functions_samples = stack_trace_to_samples(functions_stack_traces, os_resources_map);
 
     FunctionLevelProfilingInfo {
         functions_samples,
@@ -198,12 +227,14 @@ fn build_current_call_stack(
     show_inlined_functions: bool,
     sierra_statement_idx: StatementIdx,
     statements_functions_map: Option<&StatementsFunctionsMap>,
+    in_syscall: bool,
 ) -> VecWithLimitedCapacity<FunctionCall> {
     let mut current_call_stack = call_stack.current_call_stack().clone();
 
     if current_call_stack.len() == 0
         || *current_call_stack[current_call_stack.len() - 1].function_name()
             != current_function_name
+            && !in_syscall
     {
         current_call_stack.push(FunctionCall::InternalFunctionCall(
             InternalFunctionCall::NonInlined(current_function_name),
@@ -220,3 +251,62 @@ fn build_current_call_stack(
         current_call_stack
     }
 }
+
+fn stack_trace_to_samples(
+    functions_stack_traces: HashMap<Vec<FunctionCall>, Steps>,
+    os_resources_map: &OsResources,
+) -> Vec<Sample> {
+    functions_stack_traces
+        .into_iter()
+        .map(|(call_stack, steps)| {
+            let mut measurements: HashMap<MeasurementUnit, MeasurementValue> = vec![(
+                MeasurementUnit::from("steps".to_string()),
+                MeasurementValue(i64::try_from(steps.0).unwrap()),
+            )]
+            .into_iter()
+            .collect();
+
+            if let Some(FunctionCall::InternalFunctionCall(InternalFunctionCall::Syscall(
+                function_name,
+            ))) = call_stack.last()
+            {
+                let Ok(syscall) = map_syscall_name_to_selector(function_name.0.as_str()) else {
+                    // todo: print the error in debug mode
+                    return Sample {
+                        call_stack,
+                        measurements,
+                    };
+                };
+                let resources = os_resources_map
+                    .execute_syscalls
+                    .get(&syscall)
+                    .unwrap_or_else(|| {
+                        panic!("Missing syscall {syscall:?} from versioned constants file")
+                    });
+
+                if let Some(value) =
+                    measurements.get_mut(&MeasurementUnit::from("steps".to_string()))
+                {
+                    *value += MeasurementValue(i64::try_from(resources.n_steps).unwrap());
+                }
+
+                measurements.insert(
+                    MeasurementUnit::from("memory_holes".to_string()),
+                    MeasurementValue(i64::try_from(resources.n_memory_holes).unwrap()),
+                );
+
+                for (builtin, b_count) in &resources.builtin_instance_counter {
+                    measurements.insert(
+                        MeasurementUnit::from(builtin.to_string()),
+                        MeasurementValue(i64::try_from(*b_count).unwrap()),
+                    );
+                }
+            }
+
+            Sample {
+                call_stack,
+                measurements,
+            }
+        })
+        .collect_vec()
+}

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

@@ -1,5 +1,6 @@
 use crate::trace_reader::function_name::FunctionName;
 use std::collections::HashMap;
+use std::ops::AddAssign;
 use trace_data::{ExecutionResources, L1Resources};
 
 pub(crate) struct Sample {
@@ -53,6 +54,12 @@ impl From<String> for MeasurementUnit {
 #[derive(Debug, Clone)]
 pub struct MeasurementValue(pub i64);
 
+impl AddAssign for MeasurementValue {
+    fn add_assign(&mut self, other: Self) {
+        self.0 += other.0;
+    }
+}
+
 impl Sample {
     pub fn from(
         call_stack: Vec<FunctionCall>,