chore: put RcTracker as part of the DIE context

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

@@ -135,9 +135,8 @@ struct Context {
     /// them just yet.
     flattened: bool,
 
-    // When tracking mutations we consider arrays with the same type as all being possibly mutated.
-    // This we consider to span all blocks of the functions.
-    mutated_array_types: HashSet<Type>,
+    /// Track IncrementRc instructions per block to determine whether they are useless.
+    rc_tracker: RcTracker,
 }
 
 impl Context {
@@ -167,10 +166,8 @@ impl Context {
         let block = &function.dfg[block_id];
         self.mark_terminator_values_as_used(function, block);
 
-        // Lend the shared array type to the tracker.
-        let mut mutated_array_types = std::mem::take(&mut self.mutated_array_types);
-        let mut rc_tracker = RcTracker::new(&mut mutated_array_types);
-        rc_tracker.mark_terminator_arrays_as_used(function, block);
+        self.rc_tracker.new_block();
+        self.rc_tracker.mark_terminator_arrays_as_used(function, block);
 
         let instructions_len = block.instructions().len();
 
@@ -203,12 +200,11 @@ impl Context {
                 }
             }
 
-            rc_tracker.track_inc_rcs_to_remove(*instruction_id, function);
+            self.rc_tracker.track_inc_rcs_to_remove(*instruction_id, function);
         }
 
-        self.instructions_to_remove.extend(rc_tracker.get_non_mutated_arrays(&function.dfg));
-        self.instructions_to_remove.extend(rc_tracker.rc_pairs_to_remove);
-
+        self.instructions_to_remove.extend(self.rc_tracker.get_non_mutated_arrays(&function.dfg));
+        self.instructions_to_remove.extend(self.rc_tracker.rc_pairs_to_remove.drain());
         // If there are some instructions that might trigger an out of bounds error,
         // first add constrain checks. Then run the DIE pass again, which will remove those
         // but leave the constrains (any any value needed by those constrains)
@@ -228,9 +224,6 @@ impl Context {
             .instructions_mut()
             .retain(|instruction| !self.instructions_to_remove.contains(instruction));
 
-        // Take the mutated array back.
-        self.mutated_array_types = mutated_array_types;
-
         false
     }
 
@@ -279,11 +272,15 @@ impl Context {
                 let typ = typ.get_contained_array();
                 // Want to store the array type which is being referenced,
                 // because it's the underlying array that the `inc_rc` is associated with.
-                self.mutated_array_types.insert(typ.clone());
+                self.add_mutated_array_type(typ.clone());
             }
         }
     }
 
+    fn add_mutated_array_type(&mut self, typ: Type) {
+        self.rc_tracker.mutated_array_types.insert(typ.get_contained_array().clone());
+    }
+
     /// Go through the RC instructions collected when we figured out which values were unused;
     /// for each RC that refers to an unused value, remove the RC as well.
     fn remove_rc_instructions(&self, dfg: &mut DataFlowGraph) {
@@ -615,8 +612,9 @@ fn apply_side_effects(
     (lhs, rhs)
 }
 
+#[derive(Default)]
 /// Per block RC tracker.
-struct RcTracker<'a> {
+struct RcTracker {
     // We can track IncrementRc instructions per block to determine whether they are useless.
     // IncrementRc and DecrementRc instructions are normally side effectual instructions, but we remove
     // them if their value is not used anywhere in the function. However, even when their value is used, their existence
@@ -631,23 +629,21 @@ struct RcTracker<'a> {
     // If an array is the same type as one of those non-mutated array types, we can safely remove all IncrementRc instructions on that array.
     inc_rcs: HashMap<ValueId, HashSet<InstructionId>>,
     // Mutated arrays shared across the blocks of the function.
-    mutated_array_types: &'a mut HashSet<Type>,
+    // When tracking mutations we consider arrays with the same type as all being possibly mutated.
+    mutated_array_types: HashSet<Type>,
     // The SSA often creates patterns where after simplifications we end up with repeat
     // IncrementRc instructions on the same value. We track whether the previous instruction was an IncrementRc,
     // and if the current instruction is also an IncrementRc on the same value we remove the current instruction.
     // `None` if the previous instruction was anything other than an IncrementRc
     previous_inc_rc: Option<ValueId>,
 }
 
-impl<'a> RcTracker<'a> {
-    fn new(mutated_array_types: &'a mut HashSet<Type>) -> Self {
-        Self {
-            rcs_with_possible_pairs: Default::default(),
-            rc_pairs_to_remove: Default::default(),
-            inc_rcs: Default::default(),
-            previous_inc_rc: Default::default(),
-            mutated_array_types,
-        }
+impl RcTracker {
+    fn new_block(&mut self) {
+        self.rcs_with_possible_pairs.clear();
+        self.rc_pairs_to_remove.clear();
+        self.inc_rcs.clear();
+        self.previous_inc_rc = Default::default();
     }
 
     fn mark_terminator_arrays_as_used(&mut self, function: &Function, block: &BasicBlock) {
@@ -1128,4 +1124,38 @@ mod test {
         ";
         assert_normalized_ssa_equals(ssa, expected);
     }
+
+    #[test]
+    fn do_not_remove_inc_rc_if_mutated_in_other_block() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: &mut [Field; 3]):
+            v1 = load v0 -> [Field; 3]
+            inc_rc v1
+            jmp b1()
+          b1():
+            v2 = load v0 -> [Field; 3]
+            v3 = array_set v2, index u32 0, value u32 0
+            store v3 at v0
+            return
+        }
+        ";
+        let ssa = Ssa::from_str(src).unwrap();
+
+        let expected = "
+        brillig(inline) fn main f0 {
+          b0(v0: &mut [Field; 3]):
+            v1 = load v0 -> [Field; 3]
+            inc_rc v1
+            jmp b1()
+          b1():
+            v2 = load v0 -> [Field; 3]
+            v4 = array_set v2, index u32 0, value u32 0
+            store v4 at v0
+            return
+        }
+        ";
+        let ssa = ssa.dead_instruction_elimination();
+        assert_normalized_ssa_equals(ssa, expected);
+    }
 }