feat: location tree for debug_info

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_driver/src/debug.rs

@@ -21,10 +21,14 @@ pub(crate) fn filter_relevant_files(
     let mut files_with_debug_symbols: BTreeSet<FileId> = debug_symbols
         .iter()
         .flat_map(|function_symbols| {
-            function_symbols
-                .locations
-                .values()
-                .flat_map(|call_stack| call_stack.iter().map(|location| location.file))
+            function_symbols.location_map.values().flat_map(|call_stack_id| {
+                function_symbols
+                    .location_tree
+                    .get_call_stack(*call_stack_id)
+                    .iter()
+                    .map(|location| location.file)
+                    .collect::<Vec<_>>()
+            })
         })
         .collect();
 
@@ -33,9 +37,14 @@ pub(crate) fn filter_relevant_files(
         .flat_map(|function_symbols| {
             let brillig_location_maps =
                 function_symbols.brillig_locations.values().flat_map(|brillig_location_map| {
-                    brillig_location_map
-                        .values()
-                        .flat_map(|call_stack| call_stack.iter().map(|location| location.file))
+                    brillig_location_map.values().flat_map(|call_stack_id| {
+                        function_symbols
+                            .location_tree
+                            .get_call_stack(*call_stack_id)
+                            .iter()
+                            .map(|location| location.file)
+                            .collect::<Vec<_>>()
+                    })
                 });
             brillig_location_maps
         })

compiler/noirc_errors/Cargo.toml

@@ -16,6 +16,7 @@ acvm.workspace = true
 codespan-reporting.workspace = true
 codespan.workspace = true
 fm.workspace = true
+fxhash.workspace = true
 noirc_printable_type.workspace = true
 serde.workspace = true
 serde_with = "3.2.0"

compiler/noirc_evaluator/src/ssa/ir/call_stack.rs → compiler/noirc_errors/src/call_stack.rs

@@ -1,48 +1,71 @@
 use fxhash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
-use noirc_errors::Location;
-
-pub(crate) type CallStack = Vec<Location>;
+use crate::Location;
 
+pub type CallStack = Vec<Location>;
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
-pub(crate) struct CallStackId(u32);
+pub struct CallStackId(u32);
 
 impl CallStackId {
-    pub(crate) fn root() -> Self {
+    pub fn root() -> Self {
         Self::new(0)
     }
 
-    fn new(id: usize) -> Self {
+    pub fn new(id: usize) -> Self {
         Self(id as u32)
     }
 
-    pub(crate) fn index(&self) -> usize {
+    pub fn index(&self) -> usize {
         self.0 as usize
     }
 
-    pub(crate) fn is_root(&self) -> bool {
+    pub fn is_root(&self) -> bool {
         self.0 == 0
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Hash)]
+pub struct LocationNodeDebugInfo {
+    pub parent: Option<CallStackId>,
+    pub value: Location,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default, Hash)]
+pub struct LocationTree {
+    pub locations: Vec<LocationNodeDebugInfo>,
+}
+
+impl LocationTree {
+    /// Construct a CallStack from a CallStackId
+    pub fn get_call_stack(&self, mut call_stack: CallStackId) -> CallStack {
+        let mut result = Vec::new();
+        while let Some(parent) = self.locations[call_stack.index()].parent {
+            result.push(self.locations[call_stack.index()].value);
+            call_stack = parent;
+        }
+        result.reverse();
+        result
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub(crate) struct LocationNode {
-    pub(crate) parent: Option<CallStackId>,
-    pub(crate) children: Vec<CallStackId>,
-    pub(crate) children_hash: FxHashMap<u64, CallStackId>,
-    pub(crate) value: Location,
+pub struct LocationNode {
+    pub parent: Option<CallStackId>,
+    pub children: Vec<CallStackId>,
+    pub children_hash: FxHashMap<u64, CallStackId>,
+    pub value: Location,
 }
 
 impl LocationNode {
-    pub(crate) fn new(parent: Option<CallStackId>, value: Location) -> Self {
+    pub fn new(parent: Option<CallStackId>, value: Location) -> Self {
         LocationNode { parent, children: Vec::new(), children_hash: FxHashMap::default(), value }
     }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub(crate) struct CallStackHelper {
-    locations: Vec<LocationNode>,
+pub struct CallStackHelper {
+    pub locations: Vec<LocationNode>,
 }
 
 impl Default for CallStackHelper {
@@ -56,17 +79,18 @@ impl Default for CallStackHelper {
 
 impl CallStackHelper {
     /// Construct a CallStack from a CallStackId
-    pub(crate) fn get_call_stack(&self, mut call_stack: CallStackId) -> CallStack {
+    pub fn get_call_stack(&self, mut call_stack: CallStackId) -> CallStack {
         let mut result = Vec::new();
         while let Some(parent) = self.locations[call_stack.index()].parent {
             result.push(self.locations[call_stack.index()].value);
             call_stack = parent;
         }
+        result.reverse();
         result
     }
 
     /// Returns a new CallStackId which extends the call_stack with the provided call_stack.
-    pub(crate) fn extend_call_stack(
+    pub fn extend_call_stack(
         &mut self,
         mut call_stack: CallStackId,
         locations: &CallStack,
@@ -78,7 +102,7 @@ impl CallStackHelper {
     }
 
     /// Adds a location to the call stack
-    pub(crate) fn add_child(&mut self, call_stack: CallStackId, location: Location) -> CallStackId {
+    pub fn add_child(&mut self, call_stack: CallStackId, location: Location) -> CallStackId {
         let key = fxhash::hash64(&location);
         if let Some(result) = self.locations[call_stack.index()].children_hash.get(&key) {
             if self.locations[result.index()].value == location {
@@ -96,11 +120,7 @@ impl CallStackHelper {
     }
 
     /// Retrieve the CallStackId corresponding to call_stack with the last 'len' locations removed.
-    pub(crate) fn unwind_call_stack(
-        &self,
-        mut call_stack: CallStackId,
-        mut len: usize,
-    ) -> CallStackId {
+    pub fn unwind_call_stack(&self, mut call_stack: CallStackId, mut len: usize) -> CallStackId {
         while len > 0 {
             if let Some(parent) = self.locations[call_stack.index()].parent {
                 len -= 1;
@@ -112,7 +132,7 @@ impl CallStackHelper {
         call_stack
     }
 
-    pub(crate) fn add_location_to_root(&mut self, location: Location) -> CallStackId {
+    pub fn add_location_to_root(&mut self, location: Location) -> CallStackId {
         if self.locations.is_empty() {
             self.locations.push(LocationNode::new(None, location));
             CallStackId::root()
@@ -122,7 +142,18 @@ impl CallStackHelper {
     }
 
     /// Get (or create) a CallStackId corresponding to the given locations
-    pub(crate) fn get_or_insert_locations(&mut self, locations: CallStack) -> CallStackId {
-        self.extend_call_stack(CallStackId::root(), &locations)
+    pub fn get_or_insert_locations(&mut self, locations: &CallStack) -> CallStackId {
+        self.extend_call_stack(CallStackId::root(), locations)
+    }
+
+    // Clone the locations into a LocationTree
+    pub fn to_location_tree(&self) -> LocationTree {
+        LocationTree {
+            locations: self
+                .locations
+                .iter()
+                .map(|node| LocationNodeDebugInfo { value: node.value, parent: node.parent })
+                .collect(),
+        }
     }
 }

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;
@@ -16,6 +17,8 @@ use std::io::Read;
 use std::io::Write;
 use std::mem;
 
+use crate::call_stack::CallStackId;
+use crate::call_stack::LocationTree;
 use crate::Location;
 use noirc_printable_type::PrintableType;
 use serde::{
@@ -96,13 +99,14 @@ impl ProgramDebugInfo {
 #[serde_as]
 #[derive(Default, Debug, Clone, Deserialize, Serialize, Hash)]
 pub struct DebugInfo {
+    pub brillig_locations:
+        BTreeMap<BrilligFunctionId, BTreeMap<BrilligOpcodeLocation, CallStackId>>,
+    pub location_tree: LocationTree,
     /// Map opcode index of an ACIR circuit into the source code location
     /// 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 locations: BTreeMap<OpcodeLocation, Vec<Location>>,
-    pub brillig_locations:
-        BTreeMap<BrilligFunctionId, BTreeMap<BrilligOpcodeLocation, Vec<Location>>>,
+    pub location_map: BTreeMap<AcirOpcodeLocation, CallStackId>,
     pub variables: DebugVariables,
     pub functions: DebugFunctions,
     pub types: DebugTypes,
@@ -113,11 +117,12 @@ pub struct DebugInfo {
 
 impl DebugInfo {
     pub fn new(
-        locations: BTreeMap<OpcodeLocation, Vec<Location>>,
         brillig_locations: BTreeMap<
             BrilligFunctionId,
-            BTreeMap<BrilligOpcodeLocation, Vec<Location>>,
+            BTreeMap<BrilligOpcodeLocation, CallStackId>,
         >,
+        location_map: BTreeMap<AcirOpcodeLocation, CallStackId>,
+        location_tree: LocationTree,
         variables: DebugVariables,
         functions: DebugFunctions,
         types: DebugTypes,
@@ -126,7 +131,15 @@ impl DebugInfo {
             BTreeMap<ProcedureDebugId, (usize, usize)>,
         >,
     ) -> Self {
-        Self { locations, brillig_locations, variables, functions, types, brillig_procedure_locs }
+        Self {
+            brillig_locations,
+            location_map,
+            location_tree,
+            variables,
+            functions,
+            types,
+            brillig_procedure_locs,
+        }
     }
 
     /// Updates the locations map when the [`Circuit`][acvm::acir::circuit::Circuit] is modified.
@@ -136,16 +149,25 @@ impl DebugInfo {
     /// Note: One old `OpcodeLocation` might have transformed into more than one new `OpcodeLocation`.
     #[tracing::instrument(level = "trace", skip(self, update_map))]
     pub fn update_acir(&mut self, update_map: AcirTransformationMap) {
-        let old_locations = mem::take(&mut self.locations);
+        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| {
-                self.locations.insert(new_opcode_location, source_locations.clone());
+            update_map.new_acir_locations(old_opcode_location).for_each(|new_opcode_location| {
+                self.location_map.insert(new_opcode_location, source_locations);
             });
         }
     }
 
+    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>> {
-        self.locations.get(loc).cloned()
+        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)),
+        }
     }
 }

compiler/noirc_errors/src/lib.rs

@@ -3,6 +3,7 @@
 #![warn(unreachable_pub)]
 #![warn(clippy::semicolon_if_nothing_returned)]
 
+pub mod call_stack;
 pub mod debug_info;
 mod position;
 pub mod reporter;