fix: Handle multiple entry points for Brillig call stack resolution a…

…fter metadata deduplication (#5788)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

Issue created by #5696.

## Summary\*

After the creation of the brillig locations map in
#5696. we do the following to
attach Brillig locations to our generated acir:
```rust
        if inserted_func_before {
            return;
        }

        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(),
            );
        }
```
This inserts the ACIR index at the location the brillig call was pushed.
However, this is only valid for the first Brillig call. Any later
Brillig calls will also use this acir index to resolve against the debug
metadata.

We just need to avoid trying to match against this acir index as the
Brillig call stacks do not require the acir index in order to be
resolved against the debug metadata.

I also fixed the flamegraph in this PR.

For the following snippet where x == 0 causing a failure:
```rust
fn main(x: Field) {
    unsafe {
        assert(1 == conditional_wrapper(!x as bool));
        assert(1 == conditional_wrapper(x as bool));
    }
}
unconstrained fn conditional_wrapper(x: bool) -> Field {
    conditional(x)
}
unconstrained fn conditional(x: bool) -> Field {
    assert(x);
    1
}
```

Before:
<img width="840" alt="Screenshot 2024-08-21 at 11 45 45 PM"
src="https://github.com/user-attachments/assets/532f218a-9c24-434d-8760-e9bfe6859c01">
After:
<img width="838" alt="Screenshot 2024-08-21 at 11 47 17 PM"
src="https://github.com/user-attachments/assets/82d5d690-3133-4d1e-9818-4f49de2f55d0">

The flame graph also is operating as expected now. Here some example
output from one of the Brillig entry point of the `brillig_calls` test:
<img width="1590" alt="Screenshot 2024-08-21 at 11 48 01 PM"
src="https://github.com/user-attachments/assets/1a840bda-0b8d-4233-88c2-a317bd88cc1e">

## Additional Context

The assert messages still use the `Brillig` variant of `OpcodeLocation`.
This is a breaking serialization change so I have opted for not changing
that here and making a separate struct type for Brillig opcode
locations.

## Documentation\*

Check one:
- [X] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [X] I have tested the changes locally.
- [X] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

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,13 @@ impl FromStr for OpcodeLocation {
     }
 }
 
+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;
 
@@ -98,8 +99,8 @@ pub struct DebugInfo {
     /// that they should be serialized to/from strings.
     #[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 +117,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,
             },
         ];