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

# Description

## Problem\*

Partially
#5925 (comment) and
benefits other post mem2reg function cleanup work
#5925. Want to add some more
aliasing cases.

Works towards making the alias testing surrounding mem2reg more robust.

## Summary\*

I have added a mem2reg unit test where the code is storing to an alias
inside of a loop. This triggered a failure that has been fixed in this
PR. I also confirmed that this unit test would have caught the error
triggered by #5935 in
AztecProtocol/aztec-packages#8378.

We have a simple post-mem2reg process that operates over the final
mem2reg state to determine if there are any stores that were missed that
can be removed. The check currently does not look at whether the store
we are removing is an alias.

For now I simply block removing this store if it is an alias.

## Additional Context

In follow-ups we can work on handling of aliases in this per function
state to remove aliases we know we can simplify.

## Documentation\*

Check one:
- [X] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [X] I have tested the changes locally.
- [X] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

Author: vezenovm

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

@@ -181,7 +181,12 @@ impl<'f> PerFunctionContext<'f> {
                     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 @@ mod tests {
         // 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);
+}