feat(optimization): Deduplicate more instructions (#5457)

# Description

## Problem\*

## Summary\*

@TomAFrench's recent SSA changes got me thinking more about the passes
and I was looking at how we deduplicated instructions. We never
deduplicated instructions which can require a predicate so I've added
the ability to do so here.

We only deduplicate these instructions after flattening and to do so we
cache the SideEffectsEnabled var in addition to the instruction - and
only if the instruction actually relies on the variable.

## 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.

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

@@ -288,25 +288,33 @@ impl Instruction {
     }
 
     /// Indicates if the instruction can be safely replaced with the results of another instruction with the same inputs.
-    pub(crate) fn can_be_deduplicated(&self, dfg: &DataFlowGraph) -> bool {
+    /// If `deduplicate_with_predicate` is set, we assume we're deduplicating with the instruction
+    /// and its predicate, rather than just the instruction. Setting this means instructions that
+    /// rely on predicates can be deduplicated as well.
+    pub(crate) fn can_be_deduplicated(
+        &self,
+        dfg: &DataFlowGraph,
+        deduplicate_with_predicate: bool,
+    ) -> bool {
         use Instruction::*;
 
         match self {
             // These either have side-effects or interact with memory
-            Constrain(..)
-            | EnableSideEffects { .. }
+            EnableSideEffects { .. }
             | Allocate
             | Load { .. }
             | Store { .. }
             | IncrementRc { .. }
-            | DecrementRc { .. }
-            | RangeCheck { .. } => false,
+            | DecrementRc { .. } => false,
 
             Call { func, .. } => match dfg[*func] {
                 Value::Intrinsic(intrinsic) => !intrinsic.has_side_effects(),
                 _ => false,
             },
 
+            // We can deduplicate these instructions if we know the predicate is also the same.
+            Constrain(..) | RangeCheck { .. } => deduplicate_with_predicate,
+
             // These can have different behavior depending on the EnableSideEffectsIf context.
             // Replacing them with a similar instruction potentially enables replacing an instruction
             // with one that was disabled. See
@@ -317,7 +325,9 @@ impl Instruction {
             | Truncate { .. }
             | IfElse { .. }
             | ArrayGet { .. }
-            | ArraySet { .. } => !self.requires_acir_gen_predicate(dfg),
+            | ArraySet { .. } => {
+                deduplicate_with_predicate || !self.requires_acir_gen_predicate(dfg)
+            }
         }
     }
 
@@ -372,7 +382,7 @@ impl Instruction {
     }
 
     /// If true the instruction will depends on enable_side_effects context during acir-gen
-    fn requires_acir_gen_predicate(&self, dfg: &DataFlowGraph) -> bool {
+    pub(crate) fn requires_acir_gen_predicate(&self, dfg: &DataFlowGraph) -> bool {
         match self {
             Instruction::Binary(binary)
                 if matches!(binary.operator, BinaryOp::Div | BinaryOp::Mod) =>

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

@@ -87,12 +87,16 @@ struct Context {
     block_queue: Vec<BasicBlockId>,
 }
 
+/// HashMap from (Instruction, side_effects_enabled_var) to the results of the instruction.
+/// Stored as a two-level map to avoid cloning Instructions during the `.get` call.
+type InstructionResultCache = HashMap<Instruction, HashMap<Option<ValueId>, Vec<ValueId>>>;
+
 impl Context {
     fn fold_constants_in_block(&mut self, function: &mut Function, block: BasicBlockId) {
         let instructions = function.dfg[block].take_instructions();
 
         // Cache of instructions without any side-effects along with their outputs.
-        let mut cached_instruction_results: HashMap<Instruction, Vec<ValueId>> = HashMap::default();
+        let mut cached_instruction_results = HashMap::default();
 
         // Contains sets of values which are constrained to be equivalent to each other.
         //
@@ -124,7 +128,7 @@ impl Context {
         dfg: &mut DataFlowGraph,
         block: BasicBlockId,
         id: InstructionId,
-        instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        instruction_result_cache: &mut InstructionResultCache,
         constraint_simplification_mappings: &mut HashMap<ValueId, HashMap<ValueId, ValueId>>,
         side_effects_enabled_var: &mut ValueId,
     ) {
@@ -134,7 +138,9 @@ impl Context {
         let old_results = dfg.instruction_results(id).to_vec();
 
         // If a copy of this instruction exists earlier in the block, then reuse the previous results.
-        if let Some(cached_results) = instruction_result_cache.get(&instruction) {
+        if let Some(cached_results) =
+            Self::get_cached(dfg, instruction_result_cache, &instruction, *side_effects_enabled_var)
+        {
             Self::replace_result_ids(dfg, &old_results, cached_results);
             return;
         }
@@ -150,6 +156,7 @@ impl Context {
             dfg,
             instruction_result_cache,
             constraint_simplification_mapping,
+            *side_effects_enabled_var,
         );
 
         // If we just inserted an `Instruction::EnableSideEffects`, we need to update `side_effects_enabled_var`
@@ -224,8 +231,9 @@ impl Context {
         instruction: Instruction,
         instruction_results: Vec<ValueId>,
         dfg: &DataFlowGraph,
-        instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        instruction_result_cache: &mut InstructionResultCache,
         constraint_simplification_mapping: &mut HashMap<ValueId, ValueId>,
+        side_effects_enabled_var: ValueId,
     ) {
         if self.use_constraint_info {
             // If the instruction was a constraint, then create a link between the two `ValueId`s
@@ -258,8 +266,15 @@ impl Context {
 
         // If the instruction doesn't have side-effects and if it won't interact with enable_side_effects during acir_gen,
         // we cache the results so we can reuse them if the same instruction appears again later in the block.
-        if instruction.can_be_deduplicated(dfg) {
-            instruction_result_cache.insert(instruction, instruction_results);
+        if instruction.can_be_deduplicated(dfg, self.use_constraint_info) {
+            let use_predicate =
+                self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg);
+            let predicate = use_predicate.then_some(side_effects_enabled_var);
+
+            instruction_result_cache
+                .entry(instruction)
+                .or_default()
+                .insert(predicate, instruction_results);
         }
     }
 
@@ -273,6 +288,25 @@ impl Context {
             dfg.set_value_from_id(*old_result, *new_result);
         }
     }
+
+    fn get_cached<'a>(
+        dfg: &DataFlowGraph,
+        instruction_result_cache: &'a mut InstructionResultCache,
+        instruction: &Instruction,
+        side_effects_enabled_var: ValueId,
+    ) -> Option<&'a Vec<ValueId>> {
+        let results_for_instruction = instruction_result_cache.get(instruction);
+
+        // See if there's a cached version with no predicate first
+        if let Some(results) = results_for_instruction.and_then(|map| map.get(&None)) {
+            return Some(results);
+        }
+
+        let predicate =
+            instruction.requires_acir_gen_predicate(dfg).then_some(side_effects_enabled_var);
+
+        results_for_instruction.and_then(|map| map.get(&predicate))
+    }
 }
 
 #[cfg(test)]
@@ -725,4 +759,86 @@ mod test {
         let ending_instruction_count = instructions.len();
         assert_eq!(starting_instruction_count, ending_instruction_count);
     }
+
+    #[test]
+    fn deduplicate_instructions_with_predicates() {
+        // fn main f0 {
+        //   b0(v0: bool, v1: bool, v2: [u32; 2]):
+        //     enable_side_effects v0
+        //     v3 = array_get v2, index u32 0
+        //     v4 = array_set v2, index u32 1, value: u32 2
+        //     v5 = array_get v4, index u32 0
+        //     constrain_eq v3, v5
+        //     enable_side_effects v1
+        //     v6 = array_get v2, index u32 0
+        //     v7 = array_set v2, index u32 1, value: u32 2
+        //     v8 = array_get v7, index u32 0
+        //     constrain_eq v6, v8
+        //     enable_side_effects v0
+        //     v9 = array_get v2, index u32 0
+        //     v10 = array_set v2, index u32 1, value: u32 2
+        //     v11 = array_get v10, index u32 0
+        //     constrain_eq v9, v11
+        // }
+        let main_id = Id::test_new(0);
+
+        // Compiling main
+        let mut builder = FunctionBuilder::new("main".into(), main_id);
+
+        let v0 = builder.add_parameter(Type::bool());
+        let v1 = builder.add_parameter(Type::bool());
+        let v2 = builder.add_parameter(Type::Array(Rc::new(vec![Type::field()]), 2));
+
+        let zero = builder.numeric_constant(0u128, Type::length_type());
+        let one = builder.numeric_constant(1u128, Type::length_type());
+        let two = builder.numeric_constant(2u128, Type::length_type());
+
+        builder.insert_enable_side_effects_if(v0);
+        let v3 = builder.insert_array_get(v2, zero, Type::length_type());
+        let v4 = builder.insert_array_set(v2, one, two);
+        let v5 = builder.insert_array_get(v4, zero, Type::length_type());
+        builder.insert_constrain(v3, v5, None);
+
+        builder.insert_enable_side_effects_if(v1);
+        let v6 = builder.insert_array_get(v2, zero, Type::length_type());
+        let v7 = builder.insert_array_set(v2, one, two);
+        let v8 = builder.insert_array_get(v7, zero, Type::length_type());
+        builder.insert_constrain(v6, v8, None);
+
+        // We expect all these instructions after the 'enable side effects' instruction to be removed.
+        builder.insert_enable_side_effects_if(v0);
+        let v9 = builder.insert_array_get(v2, zero, Type::length_type());
+        let v10 = builder.insert_array_set(v2, one, two);
+        let v11 = builder.insert_array_get(v10, zero, Type::length_type());
+        builder.insert_constrain(v9, v11, None);
+
+        let ssa = builder.finish();
+        println!("{ssa}");
+
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+        assert_eq!(instructions.len(), 15);
+
+        // Expected output:
+        //
+        // fn main f0 {
+        //   b0(v0: bool, v1: bool, v2: [Field; 2]):
+        //     enable_side_effects v0
+        //     v3 = array_get v2, index Field 0
+        //     v4 = array_set v2, index Field 1, value: Field 2
+        //     v5 = array_get v4, index Field 0
+        //     constrain_eq v3, v5
+        //     enable_side_effects v1
+        //     v7 = array_set v2, index Field 1, value: Field 2
+        //     v8 = array_get v7, index Field 0
+        //     constrain_eq v3, v8
+        //     enable_side_effects v0
+        // }
+        let ssa = ssa.fold_constants_using_constraints();
+        println!("{ssa}");
+
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+        assert_eq!(instructions.len(), 10);
+    }
 }