fix: Implement constant folding during the mem2reg pass (#2464)

Co-authored-by: Maxim Vezenov <[email protected]>

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

@@ -1,10 +1,29 @@
-fn increment(mut r: &mut Field) {
-    *r = *r + 1;
-}
-
 fn main() {
     let mut x = 100;
     let mut xref = &mut x;
     increment(xref);
     assert(*xref == 101);
+
+    regression_2445();
+}
+
+fn increment(mut r: &mut Field) {
+    *r = *r + 1;
+}
+
+// If aliasing within arrays and constant folding within the mem2reg pass aren't
+// handled, we'll fail to optimize out all the references in this function.
+fn regression_2445() {
+    let mut var = 0;
+    let ref = &mut &mut var;
+
+    let mut array = [ref, ref];
+
+    **array[0] = 1;
+    **array[1] = 2;
+
+    assert(var == 2);
+    assert(**ref == 2);
+    assert(**array[0] == 2);
+    assert(**array[1] == 2);
 }

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

@@ -171,7 +171,7 @@ impl DataFlowGraph {
                 let id = self.make_instruction(instruction, ctrl_typevars);
                 self.blocks[block].insert_instruction(id);
                 self.locations.insert(id, call_stack);
-                InsertInstructionResult::Results(self.instruction_results(id))
+                InsertInstructionResult::Results(id, self.instruction_results(id))
             }
         }
     }
@@ -478,7 +478,8 @@ impl std::ops::IndexMut<BasicBlockId> for DataFlowGraph {
 // be a list of results or a single ValueId if the instruction was simplified
 // to an existing value.
 pub(crate) enum InsertInstructionResult<'dfg> {
-    Results(&'dfg [ValueId]),
+    /// Results is the standard case containing the instruction id and the results of that instruction.
+    Results(InstructionId, &'dfg [ValueId]),
     SimplifiedTo(ValueId),
     SimplifiedToMultiple(Vec<ValueId>),
     InstructionRemoved,
@@ -490,7 +491,7 @@ impl<'dfg> InsertInstructionResult<'dfg> {
         match self {
             InsertInstructionResult::SimplifiedTo(value) => *value,
             InsertInstructionResult::SimplifiedToMultiple(values) => values[0],
-            InsertInstructionResult::Results(results) => results[0],
+            InsertInstructionResult::Results(_, results) => results[0],
             InsertInstructionResult::InstructionRemoved => {
                 panic!("Instruction was removed, no results")
             }
@@ -501,7 +502,7 @@ impl<'dfg> InsertInstructionResult<'dfg> {
     /// This is used for instructions returning multiple results like function calls.
     pub(crate) fn results(self) -> Cow<'dfg, [ValueId]> {
         match self {
-            InsertInstructionResult::Results(results) => Cow::Borrowed(results),
+            InsertInstructionResult::Results(_, results) => Cow::Borrowed(results),
             InsertInstructionResult::SimplifiedTo(result) => Cow::Owned(vec![result]),
             InsertInstructionResult::SimplifiedToMultiple(results) => Cow::Owned(results),
             InsertInstructionResult::InstructionRemoved => Cow::Owned(vec![]),
@@ -513,7 +514,7 @@ impl<'dfg> InsertInstructionResult<'dfg> {
         match self {
             InsertInstructionResult::SimplifiedTo(_) => 1,
             InsertInstructionResult::SimplifiedToMultiple(results) => results.len(),
-            InsertInstructionResult::Results(results) => results.len(),
+            InsertInstructionResult::Results(_, results) => results.len(),
             InsertInstructionResult::InstructionRemoved => 0,
         }
     }

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

@@ -30,7 +30,7 @@ impl<'f> FunctionInserter<'f> {
     pub(crate) fn resolve(&mut self, mut value: ValueId) -> ValueId {
         value = self.function.dfg.resolve(value);
         match self.values.get(&value) {
-            Some(value) => *value,
+            Some(value) => self.resolve(*value),
             None => match &self.function.dfg[value] {
                 super::value::Value::Array { array, typ } => {
                     let array = array.clone();
@@ -47,12 +47,22 @@ impl<'f> FunctionInserter<'f> {
 
     /// Insert a key, value pair if the key isn't already present in the map
     pub(crate) fn try_map_value(&mut self, key: ValueId, value: ValueId) {
-        self.values.entry(key).or_insert(value);
+        if key == value {
+            // This case is technically not needed since try_map_value isn't meant to change
+            // existing entries, but we should never have a value in the map referring to itself anyway.
+            self.values.remove(&key);
+        } else {
+            self.values.entry(key).or_insert(value);
+        }
     }
 
     /// Insert a key, value pair in the map
     pub(crate) fn map_value(&mut self, key: ValueId, value: ValueId) {
-        self.values.insert(key, value);
+        if key == value {
+            self.values.remove(&key);
+        } else {
+            self.values.insert(key, value);
+        }
     }
 
     pub(crate) fn map_instruction(&mut self, id: InstructionId) -> (Instruction, CallStack) {
@@ -62,9 +72,27 @@ impl<'f> FunctionInserter<'f> {
         )
     }
 
-    pub(crate) fn push_instruction(&mut self, id: InstructionId, block: BasicBlockId) {
+    /// Maps a terminator in place, replacing any ValueId in the terminator with the
+    /// resolved version of that value id from this FunctionInserter's internal value mapping.
+    pub(crate) fn map_terminator_in_place(&mut self, block: BasicBlockId) {
+        let mut terminator = self.function.dfg[block].take_terminator();
+        terminator.mutate_values(|value| self.resolve(value));
+        self.function.dfg[block].set_terminator(terminator);
+    }
+
+    /// Push a new instruction to the given block and return its new InstructionId.
+    /// If the instruction was simplified out of the program, None is returned.
+    pub(crate) fn push_instruction(
+        &mut self,
+        id: InstructionId,
+        block: BasicBlockId,
+    ) -> Option<InstructionId> {
         let (instruction, location) = self.map_instruction(id);
-        self.push_instruction_value(instruction, id, block, location);
+
+        match self.push_instruction_value(instruction, id, block, location) {
+            InsertInstructionResult::Results(new_id, _) => Some(new_id),
+            _ => None,
+        }
     }
 
     pub(crate) fn push_instruction_value(
@@ -110,7 +138,7 @@ impl<'f> FunctionInserter<'f> {
                     values.insert(*old_result, *new_result);
                 }
             }
-            InsertInstructionResult::Results(new_results) => {
+            InsertInstructionResult::Results(_, new_results) => {
                 for (old_result, new_result) in old_results.iter().zip(*new_results) {
                     values.insert(*old_result, *new_result);
                 }

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

@@ -546,6 +546,26 @@ impl TerminatorInstruction {
         }
     }
 
+    /// Mutate each ValueId to a new ValueId using the given mapping function
+    pub(crate) fn mutate_values(&mut self, mut f: impl FnMut(ValueId) -> ValueId) {
+        use TerminatorInstruction::*;
+        match self {
+            JmpIf { condition, .. } => {
+                *condition = f(*condition);
+            }
+            Jmp { arguments, .. } => {
+                for argument in arguments {
+                    *argument = f(*argument);
+                }
+            }
+            Return { return_values } => {
+                for return_value in return_values {
+                    *return_value = f(*return_value);
+                }
+            }
+        }
+    }
+
     /// Apply a function to each value
     pub(crate) fn for_each_value<T>(&self, mut f: impl FnMut(ValueId) -> T) {
         use TerminatorInstruction::*;
@@ -882,7 +902,7 @@ pub(crate) enum SimplifyResult {
     /// a function such as a tuple
     SimplifiedToMultiple(Vec<ValueId>),
 
-    /// Replace this function with an simpler but equivalent function.
+    /// Replace this function with an simpler but equivalent instruction.
     SimplifiedToInstruction(Instruction),
 
     /// Remove the instruction, it is unnecessary

crates/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -98,7 +98,7 @@ impl Context {
         ) {
             InsertInstructionResult::SimplifiedTo(new_result) => vec![new_result],
             InsertInstructionResult::SimplifiedToMultiple(new_results) => new_results,
-            InsertInstructionResult::Results(new_results) => new_results.to_vec(),
+            InsertInstructionResult::Results(_, new_results) => new_results.to_vec(),
             InsertInstructionResult::InstructionRemoved => vec![],
         };
         assert_eq!(old_results.len(), new_results.len());

crates/noirc_evaluator/src/ssa/opt/inlining.rs

@@ -392,7 +392,7 @@ impl<'function> PerFunctionContext<'function> {
             self.context.call_stack.pop_back();
         }
 
-        let new_results = InsertInstructionResult::Results(&new_results);
+        let new_results = InsertInstructionResult::Results(call_id, &new_results);
         Self::insert_new_instruction_results(&mut self.values, old_results, new_results);
     }
 
@@ -435,7 +435,7 @@ impl<'function> PerFunctionContext<'function> {
                     values.insert(*old_result, new_result);
                 }
             }
-            InsertInstructionResult::Results(new_results) => {
+            InsertInstructionResult::Results(_, new_results) => {
                 for (old_result, new_result) in old_results.iter().zip(new_results) {
                     values.insert(*old_result, *new_result);
                 }