fix(ssa): Perform mem2reg before DIE

compiler/noirc_evaluator/src/ssa.rs

@@ -218,9 +218,17 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
         ),
         SsaPass::new(Ssa::make_constrain_not_equal_instructions, "Adding constrain not equal"),
         SsaPass::new(Ssa::check_u128_mul_overflow, "Check u128 mul overflow"),
-        SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
+        // Simplifying the CFG can have a positive effect on mem2reg: every time we unify with a
+        // yet-to-be-visited predecessor we forget known values; less blocks mean less unification.
         SsaPass::new(Ssa::simplify_cfg, "Simplifying"),
+        // We cannot run mem2reg after DIE, because it removes Store instructions.
+        // We have to run it before, to give it a chance to turn Store+Load into known values.
         SsaPass::new(Ssa::mem2reg, "Mem2Reg"),
+        // Removing unreachable instructions before DIE, so it gets rid of loads that mem2reg couldn't,
+        // if they are unreachable and would cause the DIE post-checks to fail.
+        SsaPass::new(Ssa::remove_unreachable_instructions, "Remove Unreachable Instructions")
+            .and_then(Ssa::remove_unreachable_functions),
+        SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
         SsaPass::new(Ssa::array_set_optimization, "Array Set Optimizations"),
         // The Brillig globals pass expected that we have the used globals map set for each function.
         // The used globals map is determined during DIE, so we should duplicate entry points before a DIE pass run.
@@ -238,10 +246,8 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
         // because the creation of shifted index constants can reuse their IDs.
         SsaPass::new(Ssa::brillig_array_get_and_set, "Brillig Array Get and Set Optimizations"),
         // Perform another DIE pass to update the used globals after offsetting Brillig indexes.
-        SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
-        SsaPass::new(Ssa::remove_unreachable_instructions, "Remove Unreachable Instructions")
-            // A function can be potentially unreachable post-DIE if all calls to that function were removed,
-            // or after the removal of unreachable instructions.
+        SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination")
+            // A function can be potentially unreachable post-DIE if all calls to that function were removed.
             .and_then(Ssa::remove_unreachable_functions),
         SsaPass::new(Ssa::checked_to_unchecked, "Checked to unchecked"),
         SsaPass::new_try(

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

@@ -54,6 +54,14 @@ impl Ssa {
         flattened: bool,
         skip_brillig: bool,
     ) -> Ssa {
+        // Perform post-checks on the SSA.
+        let check = |ssa: Ssa| {
+            // Check that we have established the properties expected from this pass.
+            #[cfg(debug_assertions)]
+            ssa.functions.values().for_each(|f| die_post_check(f, flattened));
+            ssa
+        };
+
         let mut previous_unused_params = None;
         loop {
             let (new_ssa, result) =
@@ -64,15 +72,16 @@ impl Ssa {
                 .unused_parameters
                 .values()
                 .any(|block_map| block_map.values().any(|params| !params.is_empty()));
+
             // If there are no unused parameters, return early
             if !has_unused {
-                return new_ssa;
+                return check(new_ssa);
             }
 
             if let Some(previous) = &previous_unused_params {
                 // If no changes to dead parameters occurred, return early
                 if previous == &result.unused_parameters {
-                    return new_ssa;
+                    return check(new_ssa);
                 }
             }
 
@@ -431,16 +440,7 @@ fn can_be_eliminated_if_unused(
         | Noop
         | MakeArray { .. } => true,
 
-        // Store instructions must be removed by DIE in acir code, any load
-        // instructions should already be unused by that point.
-        //
-        // Note that this check assumes that it is being performed after the flattening
-        // pass and after the last mem2reg pass. This is currently the case for the DIE
-        // pass where this check is done, but does mean that we cannot perform mem2reg
-        // after the DIE pass.
-        Store { .. } => {
-            flattened && function.runtime().is_acir() && function.reachable_blocks().len() == 1
-        }
+        Store { .. } => should_remove_store(function, flattened),
 
         Constrain(..)
         | ConstrainNotEqual(..)
@@ -610,6 +610,38 @@ impl RcTracker {
     }
 }
 
+/// Store instructions must be removed by DIE in acir code, any load
+/// instructions should already be unused by that point.
+///
+/// Note that this check assumes that it is being performed after the flattening
+/// pass and after the last mem2reg pass. This is currently the case for the DIE
+/// pass where this check is done, but does mean that we cannot perform mem2reg
+/// after the DIE pass.
+fn should_remove_store(func: &Function, flattened: bool) -> bool {
+    flattened && func.runtime().is_acir() && func.reachable_blocks().len() == 1
+}
+
+/// Check post-execution properties:
+/// * Store and Load instructions should be removed from ACIR after flattening.
+#[cfg(debug_assertions)]
+fn die_post_check(func: &Function, flattened: bool) {
+    if should_remove_store(func, flattened) {
+        for block_id in func.reachable_blocks() {
+            for (i, instruction_id) in func.dfg[block_id].instructions().iter().enumerate() {
+                let instruction = &func.dfg[*instruction_id];
+                if matches!(instruction, Instruction::Load { .. } | Instruction::Store { .. }) {
+                    panic!(
+                        "not expected to have Load or Store instruction after DIE in an ACIR function: {} {} / {block_id} / {i}: {:?}",
+                        func.name(),
+                        func.id(),
+                        instruction
+                    );
+                }
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
@@ -1046,7 +1078,7 @@ mod test {
             v0 = allocate -> &mut Field
             store Field 0 at v0
             return
-        }  
+        }
         "#);
     }
 

test_programs/execution_success/regression_8779/Nargo.toml

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

test_programs/execution_success/regression_8779/Prover.toml

@@ -0,0 +1,2 @@
+x = 0
+return = true

test_programs/execution_success/regression_8779/src/main.nr

@@ -0,0 +1,4 @@
+fn main(x: u32) -> pub bool {
+    let d: [&mut bool; 1] = [&mut true];
+    *d[x]
+}