fix(mem2reg): Update array set value alias set and propagate array get result as alias

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

@@ -533,7 +533,6 @@ impl<'f> PerFunctionContext<'f> {
                 let value = *value;
 
                 let address_aliases = references.get_aliases_for_value(address);
-
                 // If there was another store to this instruction without any (unremoved) loads or
                 // function calls in-between, we can remove the previous store.
                 if !self.aliased_references.contains_key(&address) && !address_aliases.is_unknown()
@@ -569,14 +568,23 @@ impl<'f> PerFunctionContext<'f> {
                         .extend(references.get_aliases_for_value(array).iter());
                     references.mark_value_used(array, self.inserter.function);
 
-                    // An expression for the array might already exist, so try to fetch it first
+                    // An expression for the value might already exist, so try to fetch it first
                     let expression = references.expressions.get(&array).copied();
                     let expression = expression.unwrap_or(Expression::Other(array));
-
                     if let Some(aliases) = references.aliases.get_mut(&expression) {
                         aliases.insert(result);
                     }
 
+                    // Any aliases of the array need to be updated to also include the result of the array get in their alias sets.
+                    for alias in (*references.get_aliases_for_value(array)).clone().iter() {
+                        // An expression for the alias might already exist, so try to fetch it first
+                        let expression = references.expressions.get(&alias).copied();
+                        let expression = expression.unwrap_or(Expression::Other(alias));
+                        if let Some(aliases) = references.aliases.get_mut(&expression) {
+                            aliases.insert(result);
+                        }
+                    }
+
                     // In this SSA:
                     //
                     // v2 = array_get v0, index v1 -> Field
@@ -608,11 +616,15 @@ impl<'f> PerFunctionContext<'f> {
                     } else {
                         AliasSet::unknown()
                     };
-
                     aliases.unify(&references.get_aliases_for_value(*value));
 
                     references.expressions.insert(result, expression);
-                    references.aliases.insert(expression, aliases);
+                    references.aliases.insert(expression, aliases.clone());
+
+                    // The value being stored in the array also needs its aliases updated to match the array itself
+                    let value_expression = references.expressions.get(value).copied();
+                    let value_expression = value_expression.unwrap_or(Expression::Other(*value));
+                    references.aliases.insert(value_expression, aliases);
 
                     // Similar to how we remember that we used a value in a `Store` instruction,
                     // take note that it was used in the `ArraySet`. If this instruction is not

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

@@ -73,6 +73,7 @@ impl Block {
             // Now we have to invalidate every reference we know of
             self.invalidate_all_references();
         } else if let Some(alias) = aliases.single_alias() {
+            // We always know address points to value
             self.set_reference_value(alias, value);
         } else {
             // More than one alias. We're not sure which it refers to so we have to
@@ -82,40 +83,6 @@ impl Block {
             }
         }
 
-        // Issue https://github.com/noir-lang/noir/issues/10070
-        // When setting the value at `address`, it may impact any expression having unknown aliases
-        // because the `address` could reference such expressions.
-        // This is what happens in the issue linked above. An array of references is modified in two branches
-        // and as a result the parameter of the block joining the two branches is set to unknown aliases in `handle_terminator()`.
-        // Any value set here may be done on an address to the unknown aliases. Since we do no know them, we must conservatively
-        // invalidate their values.
-        // Furthermore, if the aliases are known and include the modified `address`, then we must invalidate the aliases values as well.
-        let addresses_to_invalidate: Vec<ValueId> = self
-            .expressions
-            .iter()
-            .filter_map(|(other_address, other_expression)| {
-                // Skip the address we're storing to, since its value is known
-                if *other_address == address {
-                    return None;
-                }
-                if let Some(other_aliases) = self.aliases.get(other_expression) {
-                    if other_aliases.is_unknown() {
-                        // other_address with unknown alias set could alias anything and must be conservatively invalidated
-                        return Some(*other_address);
-                    }
-                    if other_aliases.iter().any(|alias| alias == address) {
-                        //other_address's alias set contains the address we're storing to, so it must be invalidated
-                        return Some(*other_address);
-                    }
-                }
-                None
-            })
-            .collect();
-        // Invalidate addresses with possible aliasing to `address`
-        for addr in addresses_to_invalidate {
-            self.references.remove(&addr);
-        }
-
         // We always know address points to value
         self.set_reference_value(address, value);
     }