feat(ssa): Switch mem2reg pass to be per function rather than per blo…

…ck (#2243)

Co-authored-by: jfecher <[email protected]>

crates/nargo_cli/tests/execution_success/array_dynamic/src/main.nr

@@ -21,7 +21,7 @@ fn dyn_array(mut x: [u32; 5], y: Field, z: Field) {
     assert(x[y] == 111);
     assert(x[z] == 101);
     x[z] = 0;
-   assert(x[y] == 111);
+    assert(x[y] == 111);
     assert(x[1] == 0);
     if y as u32 < 10 {
       x[y] = x[y] - 2;

crates/nargo_cli/tests/execution_success/references/src/main.nr

@@ -35,6 +35,12 @@ fn main(mut x: Field) {
     regression_2054();
     regression_2030();
     regression_2255();
+
+    assert(x == 3);
+    regression_2218_if_inner_if(x, 10);
+    regression_2218_if_inner_else(20, x);
+    regression_2218_else(x, 3);
+    regression_2218_loop(x, 10);
 }
 
 fn add1(x: &mut Field) {
@@ -111,3 +117,116 @@ fn regression_2255() {
 fn regression_2255_helper(mut x: &mut Field) {
     *x = 1;
 }
+
+fn regression_2218(x: Field, y: Field) -> Field {
+    let q = &mut &mut 0;
+    let q1 = *q;
+    let q2 = *q;
+
+    if x != y {
+        *q1 = 1;
+        // Make sure that we correct load reference aliases through multiple blocks
+        if x != 20 {
+            *q1 = 10;
+            *q2 = 2; // now we'd expect q1 == q2 == 2
+
+            assert(*q1 == 2);  
+        } else {
+            *q2 = 15;
+            assert(*q1 == 15);
+        }
+    } else {
+        *q2 = 20;
+        assert(*q1 == 20);  
+    }
+    // Have to assign value to return it
+    let value = *q1;
+    value
+}
+
+fn regression_2218_if_inner_if(x: Field, y: Field) {
+    let value = regression_2218(x, y);
+    assert(value == 2);
+}
+
+fn regression_2218_if_inner_else(x: Field, y: Field) {
+    let value = regression_2218(x, y);
+    assert(value == 15);
+}
+
+fn regression_2218_else(x: Field, y: Field) {
+    let value = regression_2218(x, y);
+    assert(value == 20);  
+}
+
+fn regression_2218_loop(x: Field, y: Field) {
+    let q = &mut &mut 0;
+    let q1 = *q;
+    let q2 = *q;
+
+    for _ in 0..1 {
+        if x != y {
+            *q1 = 10;
+            *q2 = 2; // now we'd expect q1 == q2 == 2
+
+            assert(*q1 == 2);
+        } else {
+            *q2 = 20;
+            assert(*q1 == 20);  
+        } 
+    }
+    assert(*q1 == 2);
+
+    for _ in 0..1 {
+        for _ in 0..5 {
+            if x != y {
+                *q1 = 1;
+                // Make sure that we correct load reference aliases through multiple blocks
+                if x != 20 {
+                    *q1 = 10;
+                    *q2 = 2; // now we'd expect q1 == q2 == 2
+
+                    assert(*q1 == 2);
+                }
+            } else {
+                *q2 = 20;
+                assert(*q1 == 20);
+            }
+        }
+        if x != y {
+            *q1 = 1;
+            for _ in 0..5 {
+                // Make sure that we correct load reference aliases through multiple blocks
+                if x != 20 {
+                    *q1 = 10;
+                    *q2 = 2; // now we'd expect q1 == q2 == 2
+
+                    assert(*q1 == 2);
+                }
+            }
+        } else {
+            *q2 = 20;
+            assert(*q1 == 20);  
+        }
+    }
+    assert(*q1 == 2);  
+
+    if x != y {
+        for _ in 0..5 {
+            if x != y {
+                *q1 = 1;
+                // Make sure that we correct load reference aliases through multiple blocks
+                if x != 20 {
+                    *q1 = 10;
+                    *q2 = 2; // now we'd expect q1 == q2 == 2
+
+                    assert(*q1 == 2);  
+                }
+            }
+        }
+    } else {
+        *q2 = 20;
+        assert(*q1 == 20);  
+    }
+    assert(*q1 == 2);  
+}

crates/noirc_evaluator/src/ssa.rs

@@ -49,13 +49,21 @@ pub(crate) fn optimize_into_acir(
         ssa = ssa
             .inline_functions()
             .print(print_ssa_passes, "After Inlining:")
+            // Run mem2reg with the CFG separated into blocks
+            .mem2reg()
+            .print(print_ssa_passes, "After Mem2Reg:")
             .evaluate_assert_constant()?
             .unroll_loops()?
             .print(print_ssa_passes, "After Unrolling:")
             .simplify_cfg()
             .print(print_ssa_passes, "After Simplifying:")
+            // Run mem2reg before flattening to handle any promotion
+            // of values that can be accessed after loop unrolling
+            .mem2reg()
+            .print(print_ssa_passes, "After Mem2Reg:")
             .flatten_cfg()
             .print(print_ssa_passes, "After Flattening:")
+            // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
             .mem2reg()
             .print(print_ssa_passes, "After Mem2Reg:")
             .fold_constants()

crates/noirc_evaluator/src/ssa/ir/cfg.rs

@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap};
 
 use super::{
     basic_block::{BasicBlock, BasicBlockId},
@@ -10,13 +10,14 @@ use super::{
 struct CfgNode {
     /// Set of blocks that containing jumps that target this block.
     /// The predecessor set has no meaningful order.
-    pub(crate) predecessors: HashSet<BasicBlockId>,
+    pub(crate) predecessors: BTreeSet<BasicBlockId>,
 
     /// Set of blocks that are the targets of jumps in this block.
     /// The successors set has no meaningful order.
-    pub(crate) successors: HashSet<BasicBlockId>,
+    pub(crate) successors: BTreeSet<BasicBlockId>,
 }
 
+#[derive(Clone)]
 /// The Control Flow Graph maintains a mapping of blocks to their predecessors
 /// and successors where predecessors are basic blocks and successors are
 /// basic blocks.
@@ -31,7 +32,7 @@ impl ControlFlowGraph {
         // therefore we must ensure that a node exists for the entry block, regardless of whether
         // it later comes to describe any edges after calling compute.
         let entry_block = func.entry_block();
-        let empty_node = CfgNode { predecessors: HashSet::new(), successors: HashSet::new() };
+        let empty_node = CfgNode { predecessors: BTreeSet::new(), successors: BTreeSet::new() };
         let data = HashMap::from([(entry_block, empty_node)]);
 
         let mut cfg = ControlFlowGraph { data };

crates/noirc_evaluator/src/ssa/ir/function.rs

@@ -1,4 +1,4 @@
-use std::collections::HashSet;
+use std::collections::BTreeSet;
 
 use iter_extended::vecmap;
 
@@ -107,8 +107,8 @@ impl Function {
     /// Note that self.dfg.basic_blocks_iter() iterates over all blocks,
     /// whether reachable or not. This function should be used if you
     /// want to iterate only reachable blocks.
-    pub(crate) fn reachable_blocks(&self) -> HashSet<BasicBlockId> {
-        let mut blocks = HashSet::new();
+    pub(crate) fn reachable_blocks(&self) -> BTreeSet<BasicBlockId> {
+        let mut blocks = BTreeSet::new();
         let mut stack = vec![self.entry_block];
 
         while let Some(block) = stack.pop() {