fix(mem2reg): Keep last stores used in array returned from a function

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

@@ -68,7 +68,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
@@ -459,13 +459,9 @@
                     self.instructions_to_remove.insert(*last_store);
                 }
 
-                if let Some(expression) = references.expressions.get(&value) {
-                    if let Some(aliases) = references.aliases.get(expression) {
-                        aliases.for_each(|alias| {
-                            self.aliased_references.entry(alias).or_default().insert(instruction);
-                        });
-                    }
-                }
+                references.for_each_alias_of(value, |_, alias| {
+                    self.aliased_references.entry(alias).or_default().insert(instruction);
+                });
 
                 references.set_known_value(address, value);
                 // If we see a store to an address, the last load to that address needs to remain.
@@ -521,20 +517,14 @@
                 }
             }
             Instruction::Call { arguments, .. } => {
-                // We want to fetch all aliases of each argument to be marked unknown as an array
-                // containing references internally can potentially be aliased by those references.
-                let mut all_aliases = Vec::new();
+                // We need to appropriately mark each alias of a reference as being used as a call argument.
+                // This prevents us potentially removing a last store that is altered within another function.
                 for arg in arguments {
-                    if let Some(expression) = references.expressions.get(arg) {
-                        if let Some(aliases) = references.aliases.get(expression) {
-                            aliases.for_each(|alias| {
-                                self.instruction_input_references.insert(alias);
-                                all_aliases.push(alias);
-                            });
-                        }
-                    }
+                    references.for_each_alias_of(*arg, |_, alias| {
+                        self.instruction_input_references.insert(alias);
+                    });
                 }
-                self.mark_all_unknown(&all_aliases, references);
+                self.mark_all_unknown(arguments, references);
             }
             Instruction::MakeArray { elements, typ } => {
                 // If `array` is an array constant that contains reference types, then insert each element
@@ -575,7 +565,8 @@
 
     fn mark_all_unknown(&self, values: &[ValueId], references: &mut Block) {
         for value in values {
-            if self.inserter.function.dfg.value_is_reference(*value) {
+            let typ = self.inserter.function.dfg.type_of_value(*value);
+            if Self::contains_references(&typ) {
                 let value = *value;
                 references.set_unknown(value);
                 references.mark_value_used(value, self.inserter.function);
@@ -1296,4 +1287,50 @@
 
         assert_normalized_ssa_equals(ssa, src);
     }
+
+    #[test]
+    fn keep_last_store_used_in_make_array_returned_from_function() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0():
+            v0 = call f1() -> [&mut u1; 2]
+            return
+        }
+        brillig(inline) fn foo f1 {
+          b0():
+            v0 = allocate -> &mut u1
+            store u1 1 at v0
+            v2 = allocate -> &mut u1
+            store u1 0 at v2
+            v4 = make_array [v0, v2] : [&mut u1; 2]
+            return v4
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+        assert_normalized_ssa_equals(ssa, src);
+    }
+
+    #[test]
+    fn keep_last_store_in_make_array_where_aliases_are_none() {
+        let src = "
+        brillig(inline) fn foo f1 {
+          b0(v0: &mut u1):
+            v1 = call f2() -> &mut u1
+            store u1 1 at v1
+            v3 = make_array [v1] : [&mut u1; 1]
+            return v3
+        }
+        brillig(inline) fn get_ref f2 {
+          b0():
+            v0 = allocate -> &mut u1
+            store u1 1 at v0
+            return v0
+        }
+        ";
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+        assert_normalized_ssa_equals(ssa, src);
+    }
 }

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

@@ -172,7 +172,7 @@ impl Block {
     }
 
     /// Iterate through each known alias of the given address and apply the function `f` to each.
-    fn for_each_alias_of<T>(
+    pub(super) fn for_each_alias_of<T>(
         &mut self,
         address: ValueId,
         mut f: impl FnMut(&mut Self, ValueId) -> T,

test_programs/execution_success/array_with_refs_return/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "array_with_refs_return"
+type = "bin"
+authors = [""]
+
+[dependencies]

test_programs/execution_success/array_with_refs_return/src/main.nr

@@ -0,0 +1,11 @@
+// Regression for issue #8786 (https://github.com/noir-lang/noir/issues/8786)
+unconstrained fn main() -> pub [i8; 2] {
+    if (*func_5()[1].0) {
+        [0, 1]
+    } else {
+        [2, 3]
+    }
+}
+unconstrained fn func_5() -> [(&mut bool, str<2>, str<4>, str<4>); 2] {
+    [((&mut true), "AB", "WTGS", "AELL"), ((&mut false), "CD", "PKHQ", "LDWE")]
+}