chore: Rework mem2reg load & store removal

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

@@ -21,7 +21,7 @@ use super::{
     dfg::{CallStack, DataFlowGraph},
     function::Function,
     map::Id,
-    types::{NumericType, Type},
+    types::Type,
     value::{Value, ValueId},
 };
 
@@ -413,6 +413,7 @@ impl Instruction {
 
             Constrain(..)
             | EnableSideEffectsIf { .. }
+            | Store { .. }
             | IncrementRc { .. }
             | DecrementRc { .. }
             | RangeCheck { .. } => false,

compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs

@@ -1,9 +1,9 @@
 use acvm::{acir::AcirField, FieldElement};
 use serde::{Deserialize, Serialize};
 
-use super::{
-    DataFlowGraph, Instruction, InstructionResultType, NumericType, SimplifyResult, Type, ValueId,
-};
+use crate::ssa::ir::types::NumericType;
+
+use super::{DataFlowGraph, Instruction, InstructionResultType, SimplifyResult, Type, ValueId};
 
 /// Binary Operations allowed in the IR.
 /// Aside from the comparison operators (Eq and Lt), all operators

compiler/noirc_evaluator/src/ssa/ir/instruction/cast.rs

@@ -1,7 +1,9 @@
 use acvm::{acir::AcirField, FieldElement};
 use num_bigint::BigUint;
 
-use super::{DataFlowGraph, Instruction, NumericType, SimplifyResult, Type, Value, ValueId};
+use crate::ssa::ir::types::NumericType;
+
+use super::{DataFlowGraph, Instruction, SimplifyResult, Type, Value, ValueId};
 
 /// Try to simplify this cast instruction. If the instruction can be simplified to a known value,
 /// that value is returned. Otherwise None is returned.

compiler/noirc_evaluator/src/ssa/ir/post_order.rs

@@ -49,7 +49,7 @@ impl PostOrder {
                         stack.push((Visit::Last, block_id));
                         // Stack successors for visiting. Because items are taken from the top of the
                         // stack, we push the item that's due for a visit first to the top.
-                        for successor_id in func.dfg[block_id].successors().rev() {
+                        for successor_id in func.dfg[block_id].successors() {
                             if !visited.contains(&successor_id) {
                                 // This not visited check would also be cover by the next
                                 // iteration, but checking here two saves an iteration per successor.
@@ -147,4 +147,6 @@ mod tests {
         let block_a_id = func.entry_block();
         assert_eq!(post_order.0, [block_d_id, block_f_id, block_e_id, block_b_id, block_a_id]);
     }
+
+    // TODO: Add `loop` test
 }

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

@@ -7,6 +7,7 @@ use noirc_errors::Location;
 use crate::ssa::{
     ir::{
         basic_block::{BasicBlock, BasicBlockId},
+        cfg::ControlFlowGraph,
         dfg::DataFlowGraph,
         function::Function,
         instruction::{BinaryOp, Instruction, InstructionId, Intrinsic},
@@ -17,69 +18,106 @@ use crate::ssa::{
     ssa_gen::{Ssa, SSA_WORD_SIZE},
 };
 
-use super::rc::{pop_rc_for, RcInstruction};
+use super::{
+    mem2reg::{AliasSet, StoreInstructionAliases},
+    rc::{pop_rc_for, RcInstruction},
+};
 
 impl Ssa {
     /// Performs Dead Instruction Elimination (DIE) to remove any instructions with
     /// unused results.
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn dead_instruction_elimination(mut self) -> Ssa {
         for function in self.functions.values_mut() {
-            function.dead_instruction_elimination(true);
+            Context::new(function).dead_instruction_elimination(function, true);
         }
         self
     }
 }
 
-impl Function {
+pub(crate) fn dce_with_store_aliases(function: &mut Function, aliases: StoreInstructionAliases) {
+    let mut context = Context::new(function);
+    context.store_aliases = Some(aliases);
+    context.dead_instruction_elimination(function, true);
+}
+
+/// Per function context for tracking unused values and which instructions to remove.
+struct Context {
+    used_values: HashSet<ValueId>,
+    instructions_to_remove: HashSet<InstructionId>,
+
+    /// IncrementRc & DecrementRc instructions must be revisited after the main DIE pass since
+    /// they technically contain side-effects but we still want to remove them if their
+    /// `value` parameter is not used elsewhere.
+    rc_instructions: Vec<(InstructionId, BasicBlockId)>,
+
+    /// If we have information on which store instructions may alias other references we
+    /// can remove unused store instructions as well. This information is given by mem2reg
+    /// and is on each store instruction since which ValueIds are aliased would change over time
+    /// (so we cannot index by ValueId).
+    store_aliases: Option<StoreInstructionAliases>,
+
+    cfg: ControlFlowGraph,
+
+    visited_blocks: HashSet<BasicBlockId>,
+}
+
+impl Context {
+    fn new(function: &Function) -> Self {
+        Self {
+            used_values: HashSet::default(),
+            instructions_to_remove: HashSet::default(),
+            rc_instructions: Vec::new(),
+            store_aliases: None,
+            visited_blocks: HashSet::default(),
+            cfg: ControlFlowGraph::with_function(function),
+        }
+    }
+
     /// Removes any unused instructions in the reachable blocks of the given function.
     ///
     /// The blocks of the function are iterated in post order, such that any blocks containing
     /// instructions that reference results from an instruction in another block are evaluated first.
     /// If we did not iterate blocks in this order we could not safely say whether or not the results
     /// of its instructions are needed elsewhere.
-    pub(crate) fn dead_instruction_elimination(&mut self, insert_out_of_bounds_checks: bool) {
-        let mut context = Context::default();
-        for call_data in &self.dfg.data_bus.call_data {
-            context.mark_used_instruction_results(&self.dfg, call_data.array_id);
+    fn dead_instruction_elimination(
+        mut self,
+        function: &mut Function,
+        insert_out_of_bounds_checks: bool,
+    ) {
+        for call_data in &function.dfg.data_bus.call_data {
+            Self::mark_used_instruction_results(
+                &function.dfg,
+                call_data.array_id,
+                &mut self.used_values,
+            );
         }
 
         let mut inserted_out_of_bounds_checks = false;
 
-        let blocks = PostOrder::with_function(self);
+        let blocks = PostOrder::with_function(function);
         for block in blocks.as_slice() {
-            inserted_out_of_bounds_checks |= context.remove_unused_instructions_in_block(
-                self,
+            inserted_out_of_bounds_checks |= self.remove_unused_instructions_in_block(
+                function,
                 *block,
                 insert_out_of_bounds_checks,
             );
+            self.visited_blocks.insert(*block);
         }
 
         // If we inserted out of bounds check, let's run the pass again with those new
         // instructions (we don't want to remove those checks, or instructions that are
         // dependencies of those checks)
         if inserted_out_of_bounds_checks {
-            self.dead_instruction_elimination(false);
-            return;
+            let store_aliases = self.store_aliases;
+            self = Context::new(function);
+            self.store_aliases = store_aliases;
+            return self.dead_instruction_elimination(function, false);
         }
 
-        context.remove_rc_instructions(&mut self.dfg);
+        self.remove_rc_instructions(&mut function.dfg);
     }
-}
-
-/// Per function context for tracking unused values and which instructions to remove.
-#[derive(Default)]
-struct Context {
-    used_values: HashSet<ValueId>,
-    instructions_to_remove: HashSet<InstructionId>,
 
-    /// IncrementRc & DecrementRc instructions must be revisited after the main DIE pass since
-    /// they technically contain side-effects but we still want to remove them if their
-    /// `value` parameter is not used elsewhere.
-    rc_instructions: Vec<(InstructionId, BasicBlockId)>,
-}
-
-impl Context {
     /// Steps backwards through the instruction of the given block, amassing a set of used values
     /// as it goes, and at the same time marking instructions for removal if they haven't appeared
     /// in the set thus far.
@@ -106,34 +144,40 @@ impl Context {
         let block = &function.dfg[block_id];
         self.mark_terminator_values_as_used(function, block);
 
-        let instructions_len = block.instructions().len();
-
         let mut rc_tracker = RcTracker::default();
 
         // Indexes of instructions that might be out of bounds.
         // We'll remove those, but before that we'll insert bounds checks for them.
         let mut possible_index_out_of_bounds_indexes = Vec::new();
 
-        for (instruction_index, instruction_id) in block.instructions().iter().rev().enumerate() {
+        for (instruction_index, instruction_id) in block.instructions().iter().enumerate().rev() {
             let instruction = &function.dfg[*instruction_id];
 
-            if self.is_unused(*instruction_id, function) {
+            if self.is_unused(*instruction_id, function, block_id) {
                 self.instructions_to_remove.insert(*instruction_id);
 
                 if insert_out_of_bounds_checks
                     && instruction_might_result_in_out_of_bounds(function, instruction)
                 {
                     possible_index_out_of_bounds_indexes
-                        .push(instructions_len - instruction_index - 1);
+                        .push(instruction_index);
                 }
             } else {
                 use Instruction::*;
                 if matches!(instruction, IncrementRc { .. } | DecrementRc { .. }) {
                     self.rc_instructions.push((*instruction_id, block_id));
                 } else {
                     instruction.for_each_value(|value| {
-                        self.mark_used_instruction_results(&function.dfg, value);
+                        Self::mark_used_instruction_results(
+                            &function.dfg,
+                            value,
+                            &mut self.used_values,
+                        );
                     });
+
+                    if matches!(instruction, Store { .. }) {
+                        self.mark_used_store_aliases(&function.dfg, *instruction_id);
+                    }
                 }
             }
 
@@ -169,12 +213,21 @@ impl Context {
     ///
     /// An instruction can be removed as long as it has no side-effects, and none of its result
     /// values have been referenced.
-    fn is_unused(&self, instruction_id: InstructionId, function: &Function) -> bool {
+    fn is_unused(
+        &self,
+        instruction_id: InstructionId,
+        function: &Function,
+        block: BasicBlockId,
+    ) -> bool {
         let instruction = &function.dfg[instruction_id];
 
         if instruction.can_eliminate_if_unused(function) {
             let results = function.dfg.instruction_results(instruction_id);
-            results.iter().all(|result| !self.used_values.contains(result))
+            results.iter().all(|result| {
+                !self.used_values.contains(result)
+            })
+        } else if let Instruction::Store { address, .. } = instruction {
+            self.store_is_unused(instruction_id, *address, block)
         } else if let Instruction::Call { func, arguments } = instruction {
             // TODO: make this more general for instructions which don't have results but have side effects "sometimes" like `Intrinsic::AsWitness`
             let as_witness_id = function.dfg.get_intrinsic(Intrinsic::AsWitness);
@@ -185,21 +238,56 @@ impl Context {
         }
     }
 
+    fn store_is_unused(
+        &self,
+        instruction_id: InstructionId,
+        address: ValueId,
+        block: BasicBlockId,
+    ) -> bool {
+        if !self.cfg.successors(block).all(|block| self.visited_blocks.contains(&block)) {
+            return false;
+        }
+
+        if let Some(aliases) = Self::get_store_aliases(&self.store_aliases, instruction_id) {
+            aliases.any(|alias| self.used_values.contains(&alias)) == Some(false)
+                && !self.used_values.contains(&address)
+        } else {
+            false
+        }
+    }
+
     /// Adds values referenced by the terminator to the set of used values.
     fn mark_terminator_values_as_used(&mut self, function: &Function, block: &BasicBlock) {
         block.unwrap_terminator().for_each_value(|value| {
-            self.mark_used_instruction_results(&function.dfg, value);
+            Self::mark_used_instruction_results(&function.dfg, value, &mut self.used_values);
         });
     }
 
     /// Inspects a value and marks all instruction results as used.
-    fn mark_used_instruction_results(&mut self, dfg: &DataFlowGraph, value_id: ValueId) {
+    fn mark_used_instruction_results(dfg: &DataFlowGraph, value_id: ValueId, used_values: &mut HashSet<ValueId>) {
         let value_id = dfg.resolve(value_id);
         if matches!(&dfg[value_id], Value::Instruction { .. } | Value::Param { .. }) {
-            self.used_values.insert(value_id);
+            used_values.insert(value_id);
         }
     }
 
+    /// Marks each alias of `address` for the given store instruction as used
+    fn mark_used_store_aliases(&mut self, dfg: &DataFlowGraph, instruction: InstructionId) {
+        if let Some(aliases) = Self::get_store_aliases(&self.store_aliases, instruction) {
+            aliases.for_each(|alias| {
+                Self::mark_used_instruction_results(dfg, alias, &mut self.used_values)
+            })
+        }
+    }
+
+    // unknown stores mess this up...
+    fn get_store_aliases(
+        store_aliases: &Option<StoreInstructionAliases>,
+        instruction: InstructionId,
+    ) -> Option<&AliasSet> {
+        store_aliases.as_ref().and_then(|aliases| aliases.get(&instruction))
+    }
+
     fn remove_rc_instructions(self, dfg: &mut DataFlowGraph) {
         for (rc, block) in self.rc_instructions {
             let value = match &dfg[rc] {
@@ -282,11 +370,9 @@ impl Context {
 
             let call_stack = function.dfg.get_call_stack(instruction_id);
 
-            let (lhs, rhs) = if function.dfg.get_numeric_constant(*index).is_some() {
+            let lhs = if function.dfg.get_numeric_constant(*index).is_some() {
                 // If we are here it means the index is known but out of bounds. That's always an error!
-                let false_const = function.dfg.make_constant(false.into(), Type::bool());
-                let true_const = function.dfg.make_constant(true.into(), Type::bool());
-                (false_const, true_const)
+                function.dfg.make_constant(false.into(), Type::bool())
             } else {
                 // `index` will be relative to the flattened array length, so we need to take that into account
                 let array_length = function.dfg.type_of_value(*array).flattened_size();
@@ -309,11 +395,10 @@ impl Context {
                     None,
                     call_stack.clone(),
                 );
-                let is_index_out_of_bounds = is_index_out_of_bounds.first();
-                let true_const = function.dfg.make_constant(true.into(), Type::bool());
-                (is_index_out_of_bounds, true_const)
+                is_index_out_of_bounds.first()
             };
 
+            let rhs = function.dfg.make_constant(true.into(), Type::bool());
             let (lhs, rhs) = apply_side_effects(
                 side_effects_condition,
                 lhs,

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

@@ -392,6 +392,7 @@ impl InlineContext {
 
         context.blocks.insert(context.source_function.entry_block(), entry_block);
         context.inline_blocks(ssa);
+
         // translate databus values
         let databus = entry_point.dfg.data_bus.map_values(|t| context.translate_value(t));