fix(mem2reg): reuse existing expression and add missing alias

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

@@ -513,18 +513,26 @@ impl<'f> PerFunctionContext<'f> {
             Instruction::ArrayGet { array, .. } => {
                 let result = self.inserter.function.dfg.instruction_results(instruction)[0];
 
-                let array = *array;
-                let array_typ = self.inserter.function.dfg.type_of_value(array);
-                if array_typ.contains_reference() {
+                if self.inserter.function.dfg.type_of_value(result).contains_reference() {
+                    let array = *array;
                     self.instruction_input_references
                         .extend(references.get_aliases_for_value(array).iter());
                     references.mark_value_used(array, self.inserter.function);
 
-                    let expression = Expression::ArrayElement(array);
+                    // An expression for the array might already exist, so try to fetch it first
+                    let expression = references.expressions.get(&array).copied();
+                    let expression = expression.unwrap_or(Expression::ArrayElement(array));
 
                     if let Some(aliases) = references.aliases.get_mut(&expression) {
                         aliases.insert(result);
                     }
+
+                    // In this SSA:
+                    //
+                    // v2 = array_get v0, index v1 -> Field
+                    //
+                    // make v2 point to v0 so they share the same alias set
+                    references.expressions.insert(result, expression);
                 }
             }
             Instruction::ArraySet { array, value, .. } => {
@@ -535,7 +543,8 @@ impl<'f> PerFunctionContext<'f> {
                     let result = self.inserter.function.dfg.instruction_results(instruction)[0];
                     let array = *array;
 
-                    let expression = Expression::ArrayElement(array);
+                    let expression = references.expressions.get(&array).copied();
+                    let expression = expression.unwrap_or(Expression::ArrayElement(array));
 
                     let mut aliases = if let Some(aliases) = references.aliases.get_mut(&expression)
                     {
@@ -1957,4 +1966,36 @@ mod tests {
         }
         ");
     }
+
+    #[test]
+    fn does_not_reuse_load_from_aliased_array_element() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: &mut Field, v1: &mut Field, v2: u32):
+            v3 = make_array [v0] : [&mut Field; 1]
+            v4 = array_set v3, index v2, value v1
+            v5 = load v0 -> Field
+            v6 = array_get v4, index v2 -> &mut Field
+            store Field 0 at v6
+            v7 = load v0 -> Field
+            return v7
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.mem2reg();
+        assert_ssa_snapshot!(ssa, @r#"
+        brillig(inline) fn main f0 {
+          b0(v0: &mut Field, v1: &mut Field, v2: u32):
+            v3 = make_array [v0] : [&mut Field; 1]
+            v4 = array_set v3, index v2, value v1
+            v5 = load v0 -> Field
+            constrain v2 == u32 0, "Index out of bounds"
+            v7 = array_get v4, index u32 0 -> &mut Field
+            store Field 0 at v7
+            v9 = load v0 -> Field
+            return v9
+        }
+        "#);
+    }
 }

compiler/noirc_evaluator/src/ssa/ssa_gen/program.rs

@@ -165,6 +165,12 @@ impl Ssa {
     }
 }
 
+impl std::fmt::Display for Ssa {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.print_without_locations().fmt(f)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::ssa::ir::map::Id;