fix(mem2reg): Handle aliases in function last store cleanup and additional alias unit test

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
@@ -111,7 +111,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: BTreeSet<InstructionId>,
 
     /// Track a value's last load across all blocks.
@@ -181,7 +181,12 @@
                     self.last_loads.get(store_address).is_none()
                 };
 
-                if remove_load && !is_reference_param {
+                let is_reference_alias = block
+                    .expressions
+                    .get(store_address)
+                    .map_or(false, |expression| matches!(expression, Expression::Dereference(_)));
+
+                if remove_load && !is_reference_param && !is_reference_alias {
                     self.instructions_to_remove.insert(*store_instruction);
                 }
             }
@@ -789,4 +794,98 @@
         // We expect the last eq to be optimized out
         assert_eq!(b1_instructions.len(), 0);
     }
+
+    #[test]
+    fn keep_store_to_alias_in_loop_block() {
+        // This test makes sure the instruction `store Field 2 at v5` in b2 remains after mem2reg.
+        // Although the only instruction on v5 is a lone store without any loads,
+        // v5 is an alias of the reference v0 which is stored in v2.
+        // This test makes sure that we are not inadvertently removing stores to aliases across blocks.
+        //
+        // acir(inline) fn main f0 {
+        //     b0():
+        //       v0 = allocate
+        //       store Field 0 at v0
+        //       v2 = allocate
+        //       store v0 at v2
+        //       jmp b1(Field 0)
+        //     b1(v3: Field):
+        //       v4 = eq v3, Field 0
+        //       jmpif v4 then: b2, else: b3
+        //     b2():
+        //       v5 = load v2
+        //       store Field 2 at v5
+        //       v8 = add v3, Field 1
+        //       jmp b1(v8)
+        //     b3():
+        //       v9 = load v0
+        //       v10 = eq v9, Field 2
+        //       constrain v9 == Field 2
+        //       v11 = load v2
+        //       v12 = load v10
+        //       v13 = eq v12, Field 2
+        //       constrain v11 == Field 2
+        //       return
+        //   }
+        let main_id = Id::test_new(0);
+        let mut builder = FunctionBuilder::new("main".into(), main_id);
+
+        let v0 = builder.insert_allocate(Type::field());
+        let zero = builder.numeric_constant(0u128, Type::field());
+        builder.insert_store(v0, zero);
+
+        let v2 = builder.insert_allocate(Type::field());
+        // Construct alias
+        builder.insert_store(v2, v0);
+        let v2_type = builder.current_function.dfg.type_of_value(v2);
+        assert!(builder.current_function.dfg.value_is_reference(v2));
+
+        let b1 = builder.insert_block();
+        builder.terminate_with_jmp(b1, vec![zero]);
+
+        // Loop header
+        builder.switch_to_block(b1);
+        let v3 = builder.add_block_parameter(b1, Type::field());
+        let is_zero = builder.insert_binary(v3, BinaryOp::Eq, zero);
+
+        let b2 = builder.insert_block();
+        let b3 = builder.insert_block();
+        builder.terminate_with_jmpif(is_zero, b2, b3);
+
+        // Loop body
+        builder.switch_to_block(b2);
+        let v5 = builder.insert_load(v2, v2_type.clone());
+        let two = builder.numeric_constant(2u128, Type::field());
+        builder.insert_store(v5, two);
+        let one = builder.numeric_constant(1u128, Type::field());
+        let v3_plus_one = builder.insert_binary(v3, BinaryOp::Add, one);
+        builder.terminate_with_jmp(b1, vec![v3_plus_one]);
+
+        builder.switch_to_block(b3);
+        let v9 = builder.insert_load(v0, Type::field());
+        let _ = builder.insert_binary(v9, BinaryOp::Eq, two);
+
+        builder.insert_constrain(v9, two, None);
+        let v11 = builder.insert_load(v2, v2_type);
+        let v12 = builder.insert_load(v11, Type::field());
+        let _ = builder.insert_binary(v12, BinaryOp::Eq, two);
+
+        builder.insert_constrain(v11, two, None);
+        builder.terminate_with_return(vec![]);
+
+        let ssa = builder.finish();
+
+        // We expect the same result as above.
+        let ssa = ssa.mem2reg();
+
+        let main = ssa.main();
+        assert_eq!(main.reachable_blocks().len(), 4);
+
+        // The store from the original SSA should remain
+        assert_eq!(count_stores(main.entry_block(), &main.dfg), 2);
+        assert_eq!(count_stores(b2, &main.dfg), 1);
+
+        assert_eq!(count_loads(b2, &main.dfg), 1);
+        assert_eq!(count_loads(b3, &main.dfg), 3);
+    }
 }

test_programs/compile_success_empty/references_aliasing/src/main.nr

@@ -5,6 +5,7 @@ fn main() {
     assert(*xref == 101);
 
     regression_2445();
+    single_alias_inside_loop();
 }
 
 fn increment(mut r: &mut Field) {
@@ -26,3 +27,15 @@ fn regression_2445() {
     assert(**array[0] == 2);
     assert(**array[1] == 2);
 }
+
+fn single_alias_inside_loop() {
+    let mut var = 0;
+    let ref = &mut &mut var;
+
+    for _ in 0..1 {
+        **ref = 2;
+    }
+
+    assert(var == 2);
+    assert(**ref == 2);
+}