feat: speed up transformation of debug messages

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

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use acir::{
     circuit::{opcodes::UnsupportedMemoryOpcode, Circuit, Opcode, OpcodeLocation},
     BlackBoxFunc,
@@ -33,7 +35,49 @@ pub struct AcirTransformationMap {
 }
 
 impl AcirTransformationMap {
-    pub fn new_locations(
+    /// Returns a `BTreeMap` which maps an ACIR index in the untransformed [`Circuit`] to the set of ACIR indices
+    /// in the new [`Circuit`] corresponding to the opcodes.
+    pub fn get_opcode_mapping(&self) -> BTreeMap<usize, Vec<usize>> {
+        if self.acir_opcode_positions.is_empty() {
+            return BTreeMap::new();
+        };
+
+        let mut new_opcode_to_old_map: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
+        let mut index = 0;
+        let mut old_index = 0;
+
+        let mut temp: Vec<usize> = Vec::new();
+        while old_index <= *self.acir_opcode_positions.last().unwrap() {
+            let val = self.acir_opcode_positions[index];
+            match old_index.cmp(&val) {
+                std::cmp::Ordering::Less => {
+                    new_opcode_to_old_map.insert(old_index, temp);
+                    temp = Vec::new();
+                    old_index = val;
+                }
+                std::cmp::Ordering::Equal => {
+                    temp.push(index);
+                    index += 1;
+                    if index == self.acir_opcode_positions.len() {
+                        new_opcode_to_old_map.insert(old_index, temp);
+                        break;
+                    }
+                }
+                std::cmp::Ordering::Greater => {
+                    // We assume that `self.acir_opcodes_positions` is sorted. For this assumption to be broken
+                    // we would require the situation where a circuit `[opcode_1, opcode_2]` is optimized such that
+                    // the new opcodes generated from `opcode_1` would be executed after the opcodes from `opcode_2`.
+                    //
+                    // There is no reason for this to occur so we can discount it.
+                    unreachable!("`old_index` cannot exceed `val`")
+                }
+            }
+        }
+
+        new_opcode_to_old_map
+    }
+
+    fn new_locations(
         &self,
         old_location: OpcodeLocation,
     ) -> impl Iterator<Item = OpcodeLocation> + '_ {

compiler/noirc_errors/src/debug_info.rs

@@ -37,15 +37,36 @@ impl DebugInfo {
     ///
     /// The [`OpcodeLocation`]s are generated with the ACIR, but passing the ACIR through a transformation step
     /// renders the old `OpcodeLocation`s invalid. The AcirTransformationMap is able to map the old `OpcodeLocation` to the new ones.
-    /// Note: One old `OpcodeLocation` might have transformed into more than one new `OpcodeLocation`.
+    /// Note: One old `OpcodeLocation` might have transformed into more than one new `OpcodeLocation`, or removed entirely.
     pub fn update_acir(&mut self, update_map: AcirTransformationMap) {
         log::trace!("Start debug info update");
         let old_locations = mem::take(&mut self.locations);
 
+        // `update_map` maps as `update_map.get(new_acir_index) = old_acir_index` which is many to one.
+        // We want a mapping which is `reverse_mapping_tree.get(old_acir_index) = new_acir_indices`, which is one to many.
+        let reverse_mapping_tree = update_map.get_opcode_mapping();
+
         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());
-            });
+            let old_acir_index = match old_opcode_location {
+                OpcodeLocation::Acir(index) => index,
+                OpcodeLocation::Brillig { acir_index, .. } => acir_index,
+            };
+
+            if let Some(new_acir_indices) = reverse_mapping_tree.get(&old_acir_index) {
+                // Update the `OpcodeLocation`s with the new acir index and write them back into the `locations` map.
+                new_acir_indices
+                    .iter()
+                    .map(|&new_acir_index| match old_opcode_location {
+                        OpcodeLocation::Acir(_) => OpcodeLocation::Acir(new_acir_index),
+                        OpcodeLocation::Brillig { brillig_index, .. } => {
+                            // brillig indices are unaffected by optimizations.
+                            OpcodeLocation::Brillig { acir_index: new_acir_index, brillig_index }
+                        }
+                    })
+                    .for_each(|new_opcode_location| {
+                        self.locations.insert(new_opcode_location, source_locations.clone());
+                    });
+            }
         }
         log::trace!("Finish debug info update");
     }