Parse mapping from sierra debug info

crates/cairo-profiler/src/sierra_loader.rs

@@ -1,5 +1,6 @@
 use anyhow::{anyhow, Context, Result};
-use cairo_lang_sierra::program::{Program, VersionedProgram};
+use cairo_lang_sierra::debug_info::DebugInfo;
+use cairo_lang_sierra::program::{Program, ProgramArtifact, StatementIdx, 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;
@@ -16,6 +17,7 @@ pub struct CompiledArtifactsCache(HashMap<Utf8PathBuf, CompiledArtifacts>);
 pub struct CompiledArtifacts {
     pub sierra_program: SierraProgram,
     pub casm_debug_info: CairoProgramDebugInfo,
+    pub maybe_statements_functions_map: Option<StatementsFunctionsMap>,
 }
 
 pub enum SierraProgram {
@@ -33,6 +35,17 @@ impl SierraProgram {
     }
 }
 
+#[allow(dead_code)]
+#[derive(Default, Clone)]
+pub struct StatementsFunctionsMap(HashMap<StatementIdx, Vec<String>>);
+
+#[allow(dead_code)]
+impl StatementsFunctionsMap {
+    pub fn get(&self, key: StatementIdx) -> Option<&Vec<String>> {
+        self.0.get(&key)
+    }
+}
+
 impl CompiledArtifactsCache {
     pub fn new() -> Self {
         Self(HashMap::new())
@@ -64,6 +77,9 @@ pub fn compile_sierra_and_add_compiled_artifacts_to_cache(
                 .extract_sierra_program()
                 .context("Failed to extract sierra program from contract code")?;
 
+            let maybe_statements_functions_map =
+                maybe_get_statements_functions_map(contract_class.sierra_program_debug_info);
+
             let contract_class = ContractClass {
                 // Debug info is unused in the compilation. This saves us a costly clone.
                 sierra_program_debug_info: None,
@@ -83,16 +99,19 @@ pub fn compile_sierra_and_add_compiled_artifacts_to_cache(
                 CompiledArtifacts {
                     sierra_program: SierraProgram::ContractClass(program),
                     casm_debug_info,
+                    maybe_statements_functions_map,
                 },
             );
 
             return Ok(());
         }
 
         if let Ok(versioned_program) = serde_json::from_str::<VersionedProgram>(&raw_sierra) {
-            let program = versioned_program
+            let ProgramArtifact{ program, debug_info} = versioned_program
                 .into_v1()
-                .context("Failed to extract program artifact from versioned program. Make sure your versioned program is of version 1")?.program;
+                .context("Failed to extract program artifact from versioned program. Make sure your versioned program is of version 1")?;
+
+            let maybe_statements_functions_map = maybe_get_statements_functions_map(debug_info);
 
             let casm = cairo_lang_sierra_to_casm::compiler::compile(
                 &program,
@@ -110,6 +129,7 @@ pub fn compile_sierra_and_add_compiled_artifacts_to_cache(
                 CompiledArtifacts {
                     sierra_program: SierraProgram::VersionedProgram(program),
                     casm_debug_info: casm.debug_info,
+                    maybe_statements_functions_map,
                 },
             );
 
@@ -125,6 +145,30 @@ pub fn compile_sierra_and_add_compiled_artifacts_to_cache(
     Ok(())
 }
 
+fn maybe_get_statements_functions_map(
+    maybe_sierra_program_debug_info: Option<DebugInfo>,
+) -> Option<StatementsFunctionsMap> {
+    maybe_sierra_program_debug_info
+        .and_then(|mut debug_info| {
+            debug_info
+                .annotations
+                .shift_remove("github.com/software-mansion/cairo-profiler")
+        })
+        .map(|mut annotations| {
+            assert!(
+                annotations.get("statements_functions").is_some(),
+                "Wrong debug info annotations format"
+            );
+            annotations["statements_functions"].take()
+        })
+        .map(|statements_functions| {
+            let map =
+                serde_json::from_value::<HashMap<StatementIdx, Vec<String>>>(statements_functions)
+                    .expect("Wrong statements function map format");
+            StatementsFunctionsMap(map)
+        })
+}
+
 pub fn collect_and_compile_all_sierra_programs(
     trace: &CallTrace,
 ) -> Result<CompiledArtifactsCache> {