fix(mem2reg): Keep last store for reference in array used only in an array get

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

@@ -225,6 +225,7 @@ impl<'f> PerFunctionContext<'f> {
     ) -> bool {
         let reference_parameters = self.reference_parameters();
 
+        // Check whether the store address has an alias that crosses an entry point boundary (e.g. a Call or Return)
         if let Some(expression) = block.expressions.get(store_address) {
             if let Some(aliases) = block.aliases.get(expression) {
                 let allocation_aliases_parameter =
@@ -476,9 +477,13 @@ impl<'f> PerFunctionContext<'f> {
             }
             Instruction::ArrayGet { array, .. } => {
                 let result = self.inserter.function.dfg.instruction_results(instruction)[0];
-                references.mark_value_used(*array, self.inserter.function);
 
                 if self.inserter.function.dfg.value_is_reference(result) {
+                    references.for_each_alias_of(*array, |_, alias| {
+                        self.instruction_input_references.insert(alias);
+                    });
+                    references.mark_value_used(*array, self.inserter.function);
+
                     let array = *array;
                     let expression = Expression::ArrayElement(Box::new(Expression::Other(array)));
 
@@ -518,7 +523,7 @@ impl<'f> PerFunctionContext<'f> {
             }
             Instruction::Call { arguments, .. } => {
                 // 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.
+                // This prevents us potentially removing a last store from a preceding block or is altered within another function.
                 for arg in arguments {
                     references.for_each_alias_of(*arg, |_, alias| {
                         self.instruction_input_references.insert(alias);
@@ -643,6 +648,13 @@ impl<'f> PerFunctionContext<'f> {
                 }
             }
             TerminatorInstruction::Return { return_values, .. } => {
+                // We need to appropriately mark each alias of a reference as being used as a return terminator argument.
+                // This prevents us potentially removing a last store from a preceding block or is altered within another function.
+                for return_value in return_values {
+                    references.for_each_alias_of(*return_value, |_, alias| {
+                        self.instruction_input_references.insert(alias);
+                    });
+                }
                 // Removing all `last_stores` for each returned reference is more important here
                 // than setting them all to ReferenceValue::Unknown since no other block should
                 // have a block with a Return terminator as a predecessor anyway.
@@ -1289,7 +1301,7 @@ mod tests {
     }
 
     #[test]
-    fn keep_last_store_used_in_make_array_returned_from_function() {
+    fn keep_last_store_in_make_array_returned_from_function() {
         let src = "
         brillig(inline) fn main f0 {
           b0():
@@ -1312,6 +1324,129 @@ mod tests {
         assert_normalized_ssa_equals(ssa, src);
     }
 
+    #[test]
+    fn keep_last_store_in_make_array_used_in_array_get_that_returns_result() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0():
+            v1 = allocate -> &mut u1
+            store u1 0 at v1
+            v3 = make_array [v1] : [&mut u1]
+            jmpif u1 1 then: b1, else: b2
+          b1():
+            jmp b3(u32 0)
+          b2():
+            jmp b3(u32 0)
+          b3(v0: u32):
+            constrain v0 == u32 0
+            v6 = array_get v3, index v0 -> &mut u1
+            v7 = load v6 -> u1
+            return v7
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+        assert_normalized_ssa_equals(ssa, src);
+    }
+
+    #[test]
+    fn keep_last_store_in_diff_block_from_make_array_used_in_array_get_that_returns_result() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0():
+            v1 = allocate -> &mut u1
+            store u1 0 at v1
+            jmpif u1 1 then: b1, else: b2
+          b1():
+            v3 = make_array [v1] : [&mut u1]
+            jmp b3(u32 0)
+          b2():
+            jmp b3(u32 0)
+          b3(v0: u32):
+            constrain v0 == u32 0
+            v6 = array_get v3, index v0 -> &mut u1
+            v7 = load v6 -> u1
+            return v7
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+        assert_normalized_ssa_equals(ssa, src);
+    }
+
+    #[test]
+    fn remove_last_store_in_make_array_that_is_never_used() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0():
+            v0 = allocate -> &mut u1
+            store u1 1 at v0
+            jmp b1()
+          b1():
+            v2 = make_array [v0] : [&mut u1]
+            return
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+
+        assert_ssa_snapshot!(ssa, @r"
+        brillig(inline) fn main f0 {
+          b0():
+            v0 = allocate -> &mut u1
+            jmp b1()
+          b1():
+            v1 = make_array [v0] : [&mut u1]
+            return
+        }
+        ");
+    }
+
+    #[test]
+    fn keep_last_store_in_make_array_returned_from_function_separate_blocks() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: u1):
+            v1 = call f1(v0) -> [&mut u1; 1]
+            v3 = array_get v1, index u32 0 -> &mut u1
+            store u32 1 at v3
+            v5 = load v3 -> u1
+            return v5
+        }
+        brillig(inline_always) fn foo f1 {
+          b0(v0: u1):
+            v1 = allocate -> &mut u1
+            store v0 at v1
+            v2 = allocate -> &mut u32
+            store u32 0 at v2
+            jmp b1()
+          b1():
+            v4 = load v2 -> u32
+            v6 = eq v4, u32 1
+            jmpif v6 then: b2, else: b3
+          b2():
+            jmp b4()
+          b3():
+            v7 = load v2 -> u32
+            v8 = add v7, u32 1
+            store v8 at v2
+            jmp b5()
+          b4():
+            v9 = make_array [v1] : [&mut u1; 1]
+            return v9
+          b5():
+            jmp b1()
+        }
+        ";
+
+        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 = "

test_programs/execution_success/array_with_refs_from_param/Nargo.toml

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

test_programs/execution_success/array_with_refs_from_param/Prover.toml

@@ -0,0 +1 @@
+return = 1

test_programs/execution_success/array_with_refs_from_param/src/main.nr

@@ -0,0 +1,7 @@
+// Regression for issue #8803 (https://github.com/noir-lang/noir/issues/8803)
+unconstrained fn main() -> pub bool {
+    false >= *func_3(&mut false)[if *func_3(&mut true)[0] { 0 } else { 0 }]
+}
+fn func_3(b: &mut bool) -> [&mut bool] {
+    [b]
+}