From b84009ca428a5790acf53a6c027146b706170574 Mon Sep 17 00:00:00 2001
From: Maxim Vezenov <mvezenov@gmail.com>
Date: Thu, 5 Sep 2024 11:46:23 -0400
Subject: [PATCH] feat(perf): Remove known store values that equal the store
 address in mem2reg (#5935)

# Description

## Problem\*

Partially resolves #4535

Replaces https://github.com/noir-lang/noir/pull/5865

## Summary\*

When we see a load we mark the address of that load as being a known
value of the load result. When we reach a store instuction, if that
store value has a known value which is equal to the address of the store
we can remove that store.

We also check whether the last instruction was an `inc_rc` or a
`dec_rc`. If it was we do not remove the store.

## Additional Context


## Documentation\*

Check one:
- [X] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [X] I have tested the changes locally.
- [X] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
---
 .../noirc_evaluator/src/ssa/opt/mem2reg.rs    | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
index 3d98f4126cf..4d3666bea1a 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
@@ -117,6 +117,10 @@ 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)>,
+
+    /// Track whether the last instruction is an inc_rc/dec_rc instruction.
+    /// If it is we should not remove any known store values.
+    inside_rc_reload: bool,
 }
 
 impl<'f> PerFunctionContext<'f> {
@@ -131,6 +135,7 @@ impl<'f> PerFunctionContext<'f> {
             blocks: BTreeMap::new(),
             instructions_to_remove: BTreeSet::new(),
             last_loads: HashMap::default(),
+            inside_rc_reload: false,
         }
     }
 
@@ -281,6 +286,10 @@ impl<'f> PerFunctionContext<'f> {
                 } 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 +305,14 @@ 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 known_value_is_address && !self.inside_rc_reload {
+                        self.instructions_to_remove.insert(instruction);
+                    }
+                }
+
                 references.set_known_value(address, value);
                 references.last_stores.insert(address, instruction);
             }
@@ -350,6 +367,18 @@ impl<'f> PerFunctionContext<'f> {
             Instruction::Call { arguments, .. } => self.mark_all_unknown(arguments, references),
             _ => (),
         }
+
+        self.track_rc_reload_state(instruction);
+    }
+
+    fn track_rc_reload_state(&mut self, instruction: InstructionId) {
+        match &self.inserter.function.dfg[instruction] {
+            // We just had an increment or decrement to an array's reference counter
+            Instruction::IncrementRc { .. } | Instruction::DecrementRc { .. } => {
+                self.inside_rc_reload = true;
+            }
+            _ => self.inside_rc_reload = false,
+        }
     }
 
     fn check_array_aliasing(&self, references: &mut Block, array: ValueId) {