fix: Handle multiple entry points for Brillig call stack resolution after metadata deduplication

acvm-repo/acir/src/circuit/mod.rs

@@ -153,9 +153,28 @@ pub struct ResolvedOpcodeLocation {
 /// map opcodes to debug information related to their context.
 pub enum OpcodeLocation {
     Acir(usize),
+    // TODO(https://github.com/noir-lang/noir/issues/5792): We can not get rid of this enum field entirely just yet as this format is still
+    // used for resolving assert messages which is a breaking serialization change.
     Brillig { acir_index: usize, brillig_index: usize },
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
+pub struct BrilligOpcodeLocation(pub usize);
+
+impl OpcodeLocation {
+    // Utility method to allow easily comparing a resolved Brillig location and a debug Brillig location.
+    // This method is useful when fetching Brillig debug locations as this does not need an ACIR index,
+    // and just need the Brillig index.
+    pub fn to_brillig_location(self) -> Option<BrilligOpcodeLocation> {
+        match self {
+            OpcodeLocation::Brillig { brillig_index, .. } => {
+                Some(BrilligOpcodeLocation(brillig_index))
+            }
+            _ => None,
+        }
+    }
+}
+
 impl std::fmt::Display for OpcodeLocation {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -204,6 +223,23 @@ impl FromStr for OpcodeLocation {
     }
 }
 
+impl FromStr for BrilligOpcodeLocation {
+    type Err = OpcodeLocationFromStrError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let index = s
+            .parse()
+            .map_err(|_| OpcodeLocationFromStrError::InvalidOpcodeLocationString(s.to_string()))?;
+        Ok(BrilligOpcodeLocation(index))
+    }
+}
+
+impl std::fmt::Display for BrilligOpcodeLocation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let index = self.0;
+        write!(f, "{index}")
+    }
+}
+
 impl<F> Circuit<F> {
     pub fn num_vars(&self) -> u32 {
         self.current_witness_index + 1

compiler/noirc_errors/src/debug_info.rs

@@ -1,4 +1,5 @@
 use acvm::acir::circuit::brillig::BrilligFunctionId;
+use acvm::acir::circuit::BrilligOpcodeLocation;
 use acvm::acir::circuit::OpcodeLocation;
 use acvm::compiler::AcirTransformationMap;
 
@@ -99,7 +100,8 @@ pub struct DebugInfo {
     #[serde_as(as = "BTreeMap<DisplayFromStr, _>")]
     pub locations: BTreeMap<OpcodeLocation, Vec<Location>>,
     #[serde_as(as = "BTreeMap<_, BTreeMap<DisplayFromStr, _>>")]
-    pub brillig_locations: BTreeMap<BrilligFunctionId, BTreeMap<OpcodeLocation, Vec<Location>>>,
+    pub brillig_locations:
+        BTreeMap<BrilligFunctionId, BTreeMap<BrilligOpcodeLocation, Vec<Location>>>,
     pub variables: DebugVariables,
     pub functions: DebugFunctions,
     pub types: DebugTypes,
@@ -116,7 +118,10 @@ pub struct OpCodesCount {
 impl DebugInfo {
     pub fn new(
         locations: BTreeMap<OpcodeLocation, Vec<Location>>,
-        brillig_locations: BTreeMap<BrilligFunctionId, BTreeMap<OpcodeLocation, Vec<Location>>>,
+        brillig_locations: BTreeMap<
+            BrilligFunctionId,
+            BTreeMap<BrilligOpcodeLocation, Vec<Location>>,
+        >,
         variables: DebugVariables,
         functions: DebugFunctions,
         types: DebugTypes,

compiler/noirc_evaluator/src/ssa/acir_gen/acir_ir/generated_acir.rs

@@ -1,6 +1,6 @@
 //! `GeneratedAcir` is constructed as part of the `acir_gen` pass to accumulate all of the ACIR
 //! program as it is being converted from SSA form.
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, u32};
 
 use crate::{
     brillig::{brillig_gen::brillig_directive, brillig_ir::artifact::GeneratedBrillig},
@@ -11,7 +11,7 @@ use acvm::acir::{
     circuit::{
         brillig::{BrilligFunctionId, BrilligInputs, BrilligOutputs},
         opcodes::{BlackBoxFuncCall, FunctionInput, Opcode as AcirOpcode},
-        AssertionPayload, OpcodeLocation,
+        AssertionPayload, BrilligOpcodeLocation, OpcodeLocation,
     },
     native_types::Witness,
     BlackBoxFunc,
@@ -53,7 +53,7 @@ pub(crate) struct GeneratedAcir<F: AcirField> {
 
     /// Brillig function id -> Opcodes locations map
     /// This map is used to prevent redundant locations being stored for the same Brillig entry point.
-    pub(crate) brillig_locations: BTreeMap<BrilligFunctionId, OpcodeToLocationsMap>,
+    pub(crate) brillig_locations: BTreeMap<BrilligFunctionId, BrilligOpcodeToLocationsMap>,
 
     /// Source code location of the current instruction being processed
     /// None if we do not know the location
@@ -77,6 +77,8 @@ pub(crate) struct GeneratedAcir<F: AcirField> {
 /// Correspondence between an opcode index (in opcodes) and the source code call stack which generated it
 pub(crate) type OpcodeToLocationsMap = BTreeMap<OpcodeLocation, CallStack>;
 
+pub(crate) type BrilligOpcodeToLocationsMap = BTreeMap<BrilligOpcodeLocation, CallStack>;
+
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
 pub(crate) enum BrilligStdlibFunc {
     Inverse,
@@ -590,6 +592,7 @@ impl<F: AcirField> GeneratedAcir<F> {
             return;
         }
 
+        // TODO(https://github.com/noir-lang/noir/issues/5792)
         for (brillig_index, message) in generated_brillig.assert_messages.iter() {
             self.assertion_payloads.insert(
                 OpcodeLocation::Brillig {
@@ -605,13 +608,10 @@ impl<F: AcirField> GeneratedAcir<F> {
         }
 
         for (brillig_index, call_stack) in generated_brillig.locations.iter() {
-            self.brillig_locations.entry(brillig_function_index).or_default().insert(
-                OpcodeLocation::Brillig {
-                    acir_index: self.opcodes.len() - 1,
-                    brillig_index: *brillig_index,
-                },
-                call_stack.clone(),
-            );
+            self.brillig_locations
+                .entry(brillig_function_index)
+                .or_default()
+                .insert(BrilligOpcodeLocation(*brillig_index), call_stack.clone());
         }
     }
 
@@ -625,6 +625,7 @@ impl<F: AcirField> GeneratedAcir<F> {
             OpcodeLocation::Acir(index) => index,
             _ => panic!("should not have brillig index"),
         };
+
         match &mut self.opcodes[acir_index] {
             AcirOpcode::BrilligCall { id, .. } => *id = brillig_function_index,
             _ => panic!("expected brillig call opcode"),

tooling/debugger/src/context.rs

@@ -442,16 +442,15 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> DebugContext<'a, B> {
         self.debug_artifact.debug_symbols[debug_location.circuit_id as usize]
             .opcode_location(&debug_location.opcode_location)
             .unwrap_or_else(|| {
-                if let Some(brillig_function_id) = debug_location.brillig_function_id {
+                if let (Some(brillig_function_id), Some(brillig_location)) = (
+                    debug_location.brillig_function_id,
+                    debug_location.opcode_location.to_brillig_location(),
+                ) {
                     let brillig_locations = self.debug_artifact.debug_symbols
                         [debug_location.circuit_id as usize]
                         .brillig_locations
                         .get(&brillig_function_id);
-                    brillig_locations
-                        .unwrap()
-                        .get(&debug_location.opcode_location)
-                        .cloned()
-                        .unwrap_or_default()
+                    brillig_locations.unwrap().get(&brillig_location).cloned().unwrap_or_default()
                 } else {
                     vec![]
                 }
@@ -660,8 +659,9 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> DebugContext<'a, B> {
     fn get_current_acir_index(&self) -> Option<usize> {
         self.get_current_debug_location().map(|debug_location| {
             match debug_location.opcode_location {
-                OpcodeLocation::Acir(acir_index) => acir_index,
-                OpcodeLocation::Brillig { acir_index, .. } => acir_index,
+                OpcodeLocation::Acir(acir_index) | OpcodeLocation::Brillig { acir_index, .. } => {
+                    acir_index
+                }
             }
         })
     }
@@ -893,8 +893,19 @@ fn build_source_to_opcode_debug_mappings(
         );
 
         for (brillig_function_id, brillig_locations_map) in &debug_symbols.brillig_locations {
+            let brillig_locations_map = brillig_locations_map
+                .iter()
+                .map(|(key, val)| {
+                    (
+                        // TODO: this is a temporary placeholder until the debugger is updated to handle the new brillig debug locations.
+                        OpcodeLocation::Brillig { acir_index: 0, brillig_index: key.0 },
+                        val.clone(),
+                    )
+                })
+                .collect();
+
             add_opcode_locations_map(
-                brillig_locations_map,
+                &brillig_locations_map,
                 &mut result,
                 &simple_files,
                 circuit_id,

tooling/nargo/src/errors.rs

@@ -158,13 +158,16 @@ fn extract_locations_from_error<F: AcirField>(
                 debug[resolved_location.acir_function_index]
                     .opcode_location(&resolved_location.opcode_location)
                     .unwrap_or_else(|| {
-                        if let Some(brillig_function_id) = brillig_function_id {
+                        if let (Some(brillig_function_id), Some(brillig_location)) = (
+                            brillig_function_id,
+                            &resolved_location.opcode_location.to_brillig_location(),
+                        ) {
                             let brillig_locations = debug[resolved_location.acir_function_index]
                                 .brillig_locations
                                 .get(&brillig_function_id);
                             brillig_locations
                                 .unwrap()
-                                .get(&resolved_location.opcode_location)
+                                .get(brillig_location)
                                 .cloned()
                                 .unwrap_or_default()
                         } else {

tooling/profiler/src/cli/gates_flamegraph_cmd.rs

@@ -22,7 +22,7 @@ pub(crate) struct GatesFlamegraphCommand {
     backend_path: String,
 
     /// Command to get a gates report from the backend. Defaults to "gates"
-    #[clap(long, short, default_value = "gates")]
+    #[clap(long, short = 'g', default_value = "gates")]
     backend_gates_command: String,
 
     #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
@@ -87,6 +87,7 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
                 opcode: AcirOrBrilligOpcode::Acir(opcode),
                 call_stack: vec![OpcodeLocation::Acir(index)],
                 count: gates,
+                brillig_function_id: None,
             })
             .collect();
 

tooling/profiler/src/cli/opcodes_flamegraph_cmd.rs

@@ -1,5 +1,6 @@
 use std::path::{Path, PathBuf};
 
+use acir::circuit::brillig::BrilligFunctionId;
 use acir::circuit::{Circuit, Opcode, OpcodeLocation};
 use clap::Args;
 use color_eyre::eyre::{self, Context};
@@ -20,7 +21,7 @@ pub(crate) struct OpcodesFlamegraphCommand {
     #[clap(long, short)]
     output: String,
 
-    /// Wether to skip brillig functions
+    /// Whether to skip brillig functions
     #[clap(long, short, action)]
     skip_brillig: bool,
 }
@@ -62,6 +63,7 @@ fn run_with_generator<Generator: FlamegraphGenerator>(
                 opcode: AcirOrBrilligOpcode::Acir(opcode.clone()),
                 call_stack: vec![OpcodeLocation::Acir(index)],
                 count: 1,
+                brillig_function_id: None,
             })
             .collect();
 
@@ -101,6 +103,7 @@ fn run_with_generator<Generator: FlamegraphGenerator>(
                         brillig_index,
                     }],
                     count: 1,
+                    brillig_function_id: Some(BrilligFunctionId(brillig_fn_index as u32)),
                 })
                 .collect();
 

tooling/profiler/src/flamegraph.rs

@@ -1,6 +1,7 @@
 use std::path::Path;
 use std::{collections::BTreeMap, io::BufWriter};
 
+use acir::circuit::brillig::BrilligFunctionId;
 use acir::circuit::OpcodeLocation;
 use acir::AcirField;
 use color_eyre::eyre::{self};
@@ -19,6 +20,7 @@ pub(crate) struct Sample<F: AcirField> {
     pub(crate) opcode: AcirOrBrilligOpcode<F>,
     pub(crate) call_stack: Vec<OpcodeLocation>,
     pub(crate) count: usize,
+    pub(crate) brillig_function_id: Option<BrilligFunctionId>,
 }
 
 #[derive(Debug, Default)]
@@ -90,9 +92,24 @@ fn generate_folded_sorted_lines<'files, F: AcirField>(
         let mut location_names: Vec<String> = sample
             .call_stack
             .into_iter()
-            .flat_map(|opcode_location| debug_symbols.locations.get(&opcode_location))
-            .flatten()
-            .map(|location| location_to_callsite_label(*location, files))
+            .flat_map(|opcode_location| {
+                debug_symbols.opcode_location(&opcode_location).unwrap_or_else(|| {
+                    if let (Some(brillig_function_id), Some(brillig_location)) =
+                        (sample.brillig_function_id, opcode_location.to_brillig_location())
+                    {
+                        let brillig_locations =
+                            debug_symbols.brillig_locations.get(&brillig_function_id);
+                        if let Some(brillig_locations) = brillig_locations {
+                            brillig_locations.get(&brillig_location).cloned().unwrap_or_default()
+                        } else {
+                            vec![]
+                        }
+                    } else {
+                        vec![]
+                    }
+                })
+            })
+            .map(|location| location_to_callsite_label(location, files))
             .collect();
 
         if location_names.is_empty() {
@@ -286,11 +303,13 @@ mod tests {
                 opcode: AcirOrBrilligOpcode::Acir(AcirOpcode::AssertZero(Expression::default())),
                 call_stack: vec![OpcodeLocation::Acir(0)],
                 count: 10,
+                brillig_function_id: None,
             },
             Sample {
                 opcode: AcirOrBrilligOpcode::Acir(AcirOpcode::AssertZero(Expression::default())),
                 call_stack: vec![OpcodeLocation::Acir(1)],
                 count: 20,
+                brillig_function_id: None,
             },
             Sample {
                 opcode: AcirOrBrilligOpcode::Acir(AcirOpcode::MemoryInit {
@@ -300,6 +319,7 @@ mod tests {
                 }),
                 call_stack: vec![OpcodeLocation::Acir(2)],
                 count: 30,
+                brillig_function_id: None,
             },
         ];