feat(mem2reg): Remove trivial stores

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

@@ -77,7 +77,7 @@ use crate::ssa::{
         instruction::{Instruction, InstructionId, TerminatorInstruction},
         post_order::PostOrder,
         types::Type,
-        value::ValueId,
+        value::{Value, ValueId},
     },
     ssa_gen::Ssa,
 };
@@ -117,6 +117,25 @@ struct PerFunctionContext<'f> {
     /// Track a value's last load across all blocks.
     /// If a value is not used in anymore loads we can remove the last store to that value.
     last_loads: HashMap<ValueId, (InstructionId, BasicBlockId)>,
+
+    /// Flag for tracking whether we had to perform a re-load as part of the Brillig CoW optimization.
+    /// Stores made as part of this optimization should not be removed.
+    /// We want to catch stores of this nature:
+    /// ```text
+    /// v3 = load v1
+    //  inc_rc v3
+    //  v4 = load v1
+    //  inc_rc v4
+    //  store v4 at v1
+    //  store v3 at v2
+    /// ```
+    ///
+    /// We keep track of an optional boolean flag as we go through instructions.
+    /// If the flag exists it means we have hit a load instruction.
+    /// If the flag is false it means we have processed a single load, while if the flag is true
+    /// it means we have performed a re-load.
+    /// The field is reset to `None` on every instruction that is not a load, inc_rc, dec_rc, or function call.
+    inside_rc_reload: Option<bool>,
 }
 
 impl<'f> PerFunctionContext<'f> {
@@ -131,6 +150,7 @@ impl<'f> PerFunctionContext<'f> {
             blocks: BTreeMap::new(),
             instructions_to_remove: BTreeSet::new(),
             last_loads: HashMap::default(),
+            inside_rc_reload: None,
         }
     }
 
@@ -275,12 +295,15 @@ impl<'f> PerFunctionContext<'f> {
 
                 // If the load is known, replace it with the known value and remove the load
                 if let Some(value) = references.get_known_value(address) {
-                    let result = self.inserter.function.dfg.instruction_results(instruction)[0];
                     self.inserter.map_value(result, value);
                     self.instructions_to_remove.insert(instruction);
                 } else {
                     references.mark_value_used(address, self.inserter.function);
 
+                    references.expressions.insert(result, Expression::Other(result));
+                    references.aliases.insert(Expression::Other(result), AliasSet::known(result));
+                    references.set_known_value(result, address);
+
                     self.last_loads.insert(address, (instruction, block_id));
                 }
             }
@@ -296,6 +319,18 @@ impl<'f> PerFunctionContext<'f> {
                     self.instructions_to_remove.insert(*last_store);
                 }
 
+                let known_value = references.get_known_value(value);
+                if let Some(known_value) = known_value {
+                    let known_value_is_address = known_value == address;
+                    if let Some(from_rc) = self.inside_rc_reload {
+                        if known_value_is_address && !from_rc {
+                            self.instructions_to_remove.insert(instruction);
+                        }
+                    } else if known_value_is_address {
+                        self.instructions_to_remove.insert(instruction);
+                    }
+                }
+
                 references.set_known_value(address, value);
                 references.last_stores.insert(address, instruction);
             }
@@ -350,6 +385,44 @@ impl<'f> PerFunctionContext<'f> {
             Instruction::Call { arguments, .. } => self.mark_all_unknown(arguments, references),
             _ => (),
         }
+
+        self.track_rc_reload_state(instruction);
+    }
+
+    /// Update the `inside_rc_reload` context variable.
+    /// To maintain the same value ids, we must run this method inside `analyze_instruction` so that
+    /// we operate on the newly pushed instruction id.
+    /// This method should also always come after running analysis on the new instruction.
+    fn track_rc_reload_state(&mut self, instruction: InstructionId) {
+        match &self.inserter.function.dfg[instruction] {
+            Instruction::Load { .. } => {
+                if self.inside_rc_reload.is_some() {
+                    self.inside_rc_reload = Some(true);
+                } else {
+                    self.inside_rc_reload = Some(false);
+                }
+            }
+            Instruction::Call { arguments, .. } => {
+                for arg in arguments {
+                    if let Value::Instruction { instruction, .. } =
+                        &self.inserter.function.dfg[*arg]
+                    {
+                        let instruction = &self.inserter.function.dfg[*instruction];
+                        if let Instruction::Load { .. } = instruction {
+                            if self.inside_rc_reload.is_some() {
+                                self.inside_rc_reload = Some(true);
+                            } else {
+                                self.inside_rc_reload = Some(false);
+                            }
+                        }
+                    }
+                }
+            }
+            Instruction::IncrementRc { .. } | Instruction::DecrementRc { .. } => {
+                // Do nothing. We want the reload state to remain the same.
+            }
+            _ => self.inside_rc_reload = None,
+        }
     }
 
     fn check_array_aliasing(&self, references: &mut Block, array: ValueId) {

test_programs/execution_success/brillig_loop_size_regression/src/main.nr

@@ -7,8 +7,10 @@ struct EnumEmulation {
 unconstrained fn main() -> pub Field {
     let mut emulated_enum = EnumEmulation { a: Option::some(1), b: Option::none(), c: Option::none() };
 
+    // Do a copy to optimize out loads in the loop
+    let copy_enum = emulated_enum;
     for _ in 0..1 {
-        assert_eq(emulated_enum.a.unwrap(), 1);
+        assert_eq(copy_enum.a.unwrap(), 1);
     }
 
     emulated_enum.a = Option::some(2);