Separate brillig/acir opcode locations in DebugInfo

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

@@ -146,6 +146,37 @@ pub enum OpcodeLocation {
     Brillig { acir_index: usize, brillig_index: usize },
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
+/// Opcodes are locatable so that callers can
+/// map opcodes to debug information related to their context.
+pub struct AcirOpcodeLocation(usize);
+impl std::fmt::Display for AcirOpcodeLocation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+/// The implementation of display and FromStr allows serializing and deserializing a OpcodeLocation to a string.
+/// This is useful when used as key in a map that has to be serialized to JSON/TOML, for example when mapping an opcode to its metadata.
+impl FromStr for AcirOpcodeLocation {
+    type Err = OpcodeLocationFromStrError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        fn parse_index(input: &str) -> Result<AcirOpcodeLocation, ParseIntError> {
+            let index = input.parse()?;
+            Ok(AcirOpcodeLocation::new(index))
+        }
+        parse_index(s)
+            .map_err(|_| OpcodeLocationFromStrError::InvalidOpcodeLocationString(s.to_string()))
+    }
+}
+
+impl AcirOpcodeLocation {
+    pub fn new(index: usize) -> Self {
+        AcirOpcodeLocation(index)
+    }
+    pub fn index(&self) -> usize {
+        self.0
+    }
+}
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct BrilligOpcodeLocation(pub usize);
 

acvm-repo/acvm/src/compiler/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use acir::{
-    circuit::{AssertionPayload, Circuit, ExpressionWidth, OpcodeLocation},
+    circuit::{AcirOpcodeLocation, AssertionPayload, Circuit, ExpressionWidth, OpcodeLocation},
     AcirField,
 };
 
@@ -56,6 +56,19 @@ impl AcirTransformationMap {
             },
         )
     }
+
+    pub fn new_acir_locations(
+        &self,
+        old_location: AcirOpcodeLocation,
+    ) -> impl Iterator<Item = AcirOpcodeLocation> + '_ {
+        let old_acir_index = old_location.index();
+
+        self.old_indices_to_new_indices.get(&old_acir_index).into_iter().flat_map(
+            move |new_indices| {
+                new_indices.iter().map(move |new_index| AcirOpcodeLocation::new(*new_index))
+            },
+        )
+    }
 }
 
 fn transform_assert_messages<F: Clone>(

compiler/noirc_errors/src/debug_info.rs

@@ -1,4 +1,5 @@
 use acvm::acir::circuit::brillig::BrilligFunctionId;
+use acvm::acir::circuit::AcirOpcodeLocation;
 use acvm::acir::circuit::BrilligOpcodeLocation;
 use acvm::acir::circuit::OpcodeLocation;
 use acvm::compiler::AcirTransformationMap;
@@ -105,7 +106,7 @@ pub struct DebugInfo {
     /// Serde does not support mapping keys being enums for json, so we indicate
     /// that they should be serialized to/from strings.
     #[serde_as(as = "BTreeMap<DisplayFromStr, _>")]
-    pub location_map: BTreeMap<OpcodeLocation, CallStackId>,
+    pub location_map: BTreeMap<AcirOpcodeLocation, CallStackId>,
     pub variables: DebugVariables,
     pub functions: DebugFunctions,
     pub types: DebugTypes,
@@ -120,7 +121,7 @@ impl DebugInfo {
             BrilligFunctionId,
             BTreeMap<BrilligOpcodeLocation, CallStackId>,
         >,
-        location_map: BTreeMap<OpcodeLocation, CallStackId>,
+        location_map: BTreeMap<AcirOpcodeLocation, CallStackId>,
         location_tree: LocationTree,
         variables: DebugVariables,
         functions: DebugFunctions,
@@ -151,15 +152,22 @@ impl DebugInfo {
         let old_locations = mem::take(&mut self.location_map);
 
         for (old_opcode_location, source_locations) in old_locations {
-            update_map.new_locations(old_opcode_location).for_each(|new_opcode_location| {
+            update_map.new_acir_locations(old_opcode_location).for_each(|new_opcode_location| {
                 self.location_map.insert(new_opcode_location, source_locations);
             });
         }
     }
 
-    pub fn opcode_location(&self, loc: &OpcodeLocation) -> Option<Vec<Location>> {
+    pub fn acir_opcode_location(&self, loc: &AcirOpcodeLocation) -> Option<Vec<Location>> {
         self.location_map
             .get(loc)
             .map(|call_stack_id| self.location_tree.get_call_stack(*call_stack_id))
     }
+
+    pub fn opcode_location(&self, loc: &OpcodeLocation) -> Option<Vec<Location>> {
+        match loc {
+            OpcodeLocation::Brillig { .. } => None, //TODO: need brillig function id in order to look into brillig_locations
+            OpcodeLocation::Acir(loc) => self.acir_opcode_location(&AcirOpcodeLocation::new(*loc)),
+        }
+    }
 }

tooling/debugger/src/context.rs

@@ -892,9 +892,14 @@ fn build_source_to_opcode_debug_mappings(
 
     let mut result: BTreeMap<FileId, Vec<(usize, DebugLocation)>> = BTreeMap::new();
     for (circuit_id, debug_symbols) in debug_artifact.debug_symbols.iter().enumerate() {
+        let location_map = debug_symbols
+            .location_map
+            .iter()
+            .map(|(key, val)| (OpcodeLocation::Acir(key.index()), *val))
+            .collect();
         add_opcode_locations_map(
             debug_symbols,
-            &debug_symbols.location_map,
+            &location_map,
             &mut result,
             &simple_files,
             circuit_id,

tooling/debugger/src/source_code_printer.rs

@@ -222,7 +222,7 @@ fn render_location<'a>(
 mod tests {
     use crate::source_code_printer::render_location;
     use crate::source_code_printer::PrintedLine::Content;
-    use acvm::acir::circuit::OpcodeLocation;
+    use acvm::acir::circuit::AcirOpcodeLocation;
     use fm::FileManager;
     use noirc_artifacts::debug::DebugArtifact;
     use noirc_errors::call_stack::{CallStackId, LocationNodeDebugInfo, LocationTree};
@@ -267,8 +267,8 @@ mod tests {
 
         // We don't care about opcodes in this context,
         // we just use a dummy to construct debug_symbols
-        let mut opcode_locations = BTreeMap::<OpcodeLocation, CallStackId>::new();
-        opcode_locations.insert(OpcodeLocation::Acir(42), CallStackId::new(1));
+        let mut opcode_locations = BTreeMap::<AcirOpcodeLocation, CallStackId>::new();
+        opcode_locations.insert(AcirOpcodeLocation::new(42), CallStackId::new(1));
         let mut location_tree = LocationTree::default();
         location_tree
             .locations

tooling/nargo/src/errors.rs

@@ -2,8 +2,8 @@ use std::collections::BTreeMap;
 
 use acvm::{
     acir::circuit::{
-        brillig::BrilligFunctionId, ErrorSelector, OpcodeLocation, RawAssertionPayload,
-        ResolvedAssertionPayload, ResolvedOpcodeLocation,
+        brillig::BrilligFunctionId, AcirOpcodeLocation, BrilligOpcodeLocation, ErrorSelector,
+        OpcodeLocation, RawAssertionPayload, ResolvedAssertionPayload, ResolvedOpcodeLocation,
     },
     pwg::{ErrorLocation, OpcodeResolutionError},
     AcirField, FieldElement,
@@ -165,28 +165,20 @@ fn extract_locations_from_error<F: AcirField>(
         opcode_locations
             .iter()
             .flat_map(|resolved_location| {
+                let call_stack_id = match resolved_location.opcode_location {
+                    OpcodeLocation::Acir(idx) => {
+                        debug[resolved_location.acir_function_index].location_map
+                            [&AcirOpcodeLocation::new(idx)]
+                    }
+                    // TODO: should we use acir_index here and merge the 2 call stacks?
+                    OpcodeLocation::Brillig { brillig_index, .. } => {
+                        debug[resolved_location.acir_function_index].brillig_locations
+                            [&brillig_function_id.unwrap()][&BrilligOpcodeLocation(brillig_index)]
+                    }
+                };
                 debug[resolved_location.acir_function_index]
-                    .opcode_location(&resolved_location.opcode_location)
-                    .unwrap_or_else(|| {
-                        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);
-                            let call_stack_id = brillig_locations
-                                .unwrap()
-                                .get(brillig_location)
-                                .cloned()
-                                .unwrap_or_default();
-                            debug[resolved_location.acir_function_index]
-                                .location_tree
-                                .get_call_stack(call_stack_id)
-                        } else {
-                            vec![]
-                        }
-                    })
+                    .location_tree
+                    .get_call_stack(call_stack_id)
             })
             .collect(),
     )

tooling/profiler/src/flamegraph.rs

@@ -2,7 +2,7 @@ use std::path::Path;
 use std::{collections::BTreeMap, io::BufWriter};
 
 use acir::circuit::brillig::BrilligFunctionId;
-use acir::circuit::OpcodeLocation;
+use acir::circuit::{AcirOpcodeLocation, OpcodeLocation};
 use color_eyre::eyre::{self};
 use fm::codespan_files::Files;
 use fxhash::FxHashMap as HashMap;
@@ -176,33 +176,38 @@ fn find_callsite_labels<'files>(
     files: &'files impl Files<'files, FileId = fm::FileId>,
 ) -> Vec<String> {
     let mut procedure_id = None;
-    let source_locations = debug_symbols.opcode_location(opcode_location).unwrap_or_else(|| {
-        if let (Some(brillig_function_id), Some(brillig_location)) =
-            (brillig_function_id, opcode_location.to_brillig_location())
-        {
-            let procedure_locs = debug_symbols.brillig_procedure_locs.get(&brillig_function_id);
-            if let Some(procedure_locs) = procedure_locs {
-                for (procedure, range) in procedure_locs.iter() {
-                    if brillig_location.0 >= range.0 && brillig_location.0 <= range.1 {
-                        procedure_id = Some(*procedure);
-                        break;
+    let source_locations = match opcode_location {
+        OpcodeLocation::Acir(idx) => {
+            debug_symbols.acir_opcode_location(&AcirOpcodeLocation::new(*idx)).unwrap_or_default()
+        }
+        OpcodeLocation::Brillig { .. } => {
+            if let (Some(brillig_function_id), Some(brillig_location)) =
+                (brillig_function_id, opcode_location.to_brillig_location())
+            {
+                let procedure_locs = debug_symbols.brillig_procedure_locs.get(&brillig_function_id);
+                if let Some(procedure_locs) = procedure_locs {
+                    for (procedure, range) in procedure_locs.iter() {
+                        if brillig_location.0 >= range.0 && brillig_location.0 <= range.1 {
+                            procedure_id = Some(*procedure);
+                            break;
+                        }
                     }
                 }
-            }
-            let brillig_locations = debug_symbols.brillig_locations.get(&brillig_function_id);
-
-            if let Some(brillig_locations) = brillig_locations {
-                brillig_locations
-                    .get(&brillig_location)
-                    .map(|call_stack| debug_symbols.location_tree.get_call_stack(*call_stack))
-                    .unwrap_or_default()
+                let brillig_locations = debug_symbols.brillig_locations.get(&brillig_function_id);
+
+                if let Some(brillig_locations) = brillig_locations {
+                    brillig_locations
+                        .get(&brillig_location)
+                        .map(|call_stack| debug_symbols.location_tree.get_call_stack(*call_stack))
+                        .unwrap_or_default()
+                } else {
+                    vec![]
+                }
             } else {
                 vec![]
             }
-        } else {
-            vec![]
         }
-    });
+    };
 
     let mut callsite_labels: Vec<_> = source_locations
         .into_iter()
@@ -297,7 +302,7 @@ fn to_folded_sorted_lines(
 #[cfg(test)]
 mod tests {
     use acir::{
-        circuit::{opcodes::BlockId, Opcode as AcirOpcode, OpcodeLocation},
+        circuit::{opcodes::BlockId, AcirOpcodeLocation, Opcode as AcirOpcode, OpcodeLocation},
         native_types::Expression,
         FieldElement,
     };
@@ -365,7 +370,7 @@ mod tests {
             Location::new(find_spans_for(source_code, "whatever()")[2], file_id);
 
         //  let mut opcode_locations = BTreeMap::<OpcodeLocation, Vec<Location>>::new();
-        let mut opcode_locations = BTreeMap::<OpcodeLocation, CallStackId>::new();
+        let mut opcode_locations = BTreeMap::<AcirOpcodeLocation, CallStackId>::new();
         let mut call_stack_hlp = CallStackHelper::default();
         // main::foo::baz::whatever
         let call_stack_id = call_stack_hlp.get_or_insert_locations(&vec![
@@ -374,20 +379,20 @@ mod tests {
             foo_baz_call_location,
             baz_whatever_call_location,
         ]);
-        opcode_locations.insert(OpcodeLocation::Acir(0), call_stack_id);
+        opcode_locations.insert(AcirOpcodeLocation::new(0), call_stack_id);
         // main::bar::whatever
         let call_stack_id = call_stack_hlp.get_or_insert_locations(&vec![
             main_declaration_location,
             main_bar_call_location,
             bar_whatever_call_location,
         ]);
-        opcode_locations.insert(OpcodeLocation::Acir(1), call_stack_id);
+        opcode_locations.insert(AcirOpcodeLocation::new(1), call_stack_id);
         // main::whatever
         let call_stack_id = call_stack_hlp
             .get_or_insert_locations(&vec![main_declaration_location, main_whatever_call_location]);
-        opcode_locations.insert(OpcodeLocation::Acir(2), call_stack_id);
+        opcode_locations.insert(AcirOpcodeLocation::new(2), call_stack_id);
 
-        opcode_locations.insert(OpcodeLocation::Acir(42), CallStackId::new(1));
+        opcode_locations.insert(AcirOpcodeLocation::new(42), CallStackId::new(1));
         let location_tree = call_stack_hlp.to_location_tree();
 
         let debug_info = DebugInfo::new(