fix(mem2reg): Remove possibility of underflow

compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs

@@ -56,7 +56,7 @@
 //!
 //! Repeating this algorithm for each block in the function in program order should result in
 //! optimizing out most known loads. However, identifying all aliases correctly has been proven
 //! undecidable in general (Landi, 1992). So this pass will not always optimize out all loads
 //! that could theoretically be optimized out. This pass can be performed at any time in the
 //! SSA optimization pipeline, although it will be more successful the simpler the program's CFG is.
 //! This pass is currently performed several times to enable other passes - most notably being
@@ -144,7 +144,7 @@
     /// Load and Store instructions that should be removed at the end of the pass.
     ///
     /// We avoid removing individual instructions as we go since removing elements
     /// from the middle of Vecs many times will be slower than a single call to `retain`.
     instructions_to_remove: HashSet<InstructionId>,
 
     /// Track a value's last load across all blocks.
@@ -288,6 +288,8 @@
             .filter(|param| self.inserter.function.dfg.value_is_reference(**param))
             .collect::<BTreeSet<_>>();
 
+        // Must collect here as we are immutably borrowing `self` to fetch the reference parameters
+        let mut values_to_reduce_counts = Vec::new();
         for (allocation, instruction) in &references.last_stores {
             if let Some(expression) = references.expressions.get(allocation) {
                 if let Some(aliases) = references.aliases.get(expression) {
@@ -297,13 +299,15 @@
                     // If `allocation_aliases_parameter` is known to be false
                     if allocation_aliases_parameter == Some(false) {
                         self.instructions_to_remove.insert(*instruction);
-                        if let Some(context) = self.load_results.get_mut(allocation) {
-                            context.uses -= 1;
-                        }
+                        values_to_reduce_counts.push(*allocation);
                     }
                 }
             }
         }
+
+        for value in values_to_reduce_counts {
+            self.reduce_load_result_count(value);
+        }
     }
 
     fn increase_load_ref_counts(&mut self, value: ValueId) {
@@ -406,25 +410,17 @@
                     else {
                         panic!("Should have a store instruction here");
                     };
-                    if let Some(context) = self.load_results.get_mut(&address) {
-                        context.uses -= 1;
-                    }
-                    if let Some(context) = self.load_results.get_mut(&value) {
-                        context.uses -= 1;
-                    }
+                    self.reduce_load_result_count(address);
+                    self.reduce_load_result_count(value);
                 }
 
                 let known_value = references.get_known_value(value);
                 if let Some(known_value) = known_value {
                     let known_value_is_address = known_value == address;
                     if known_value_is_address {
                         self.instructions_to_remove.insert(instruction);
-                        if let Some(context) = self.load_results.get_mut(&address) {
-                            context.uses -= 1;
-                        }
-                        if let Some(context) = self.load_results.get_mut(&value) {
-                            context.uses -= 1;
-                        }
+                        self.reduce_load_result_count(address);
+                        self.reduce_load_result_count(value);
                     } else {
                         references.last_stores.insert(address, instruction);
                     }
@@ -617,6 +613,14 @@
         }
     }
 
+    fn reduce_load_result_count(&mut self, value: ValueId) {
+        if let Some(context) = self.load_results.get_mut(&value) {
+            if context.uses != 0 {
+                context.uses -= 1;
+            }
+        }
+    }
+
     fn recursively_add_values(&self, value: ValueId, set: &mut HashSet<ValueId>) {
         set.insert(value);
         if let Some((elements, _)) = self.inserter.function.dfg.get_array_constant(value) {
@@ -741,7 +745,9 @@
                 if all_loads_removed && !store_alias_used {
                     self.instructions_to_remove.insert(*store_instruction);
                     if let Some((_, counter)) = remaining_last_stores.get_mut(store_address) {
-                        *counter -= 1;
+                        if *counter != 0 {
+                            *counter -= 1;
+                        }
                     }
                 } else if let Some((_, counter)) = remaining_last_stores.get_mut(store_address) {
                     *counter += 1;