chore: Add optimization to #4716

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -622,7 +622,7 @@ impl<'block> BrilligBlock<'block> {
                     destination_variable,
                 );
             }
-            Instruction::ArraySet { array, index, value, .. } => {
+            Instruction::ArraySet { array, index, value, mutable: _ } => {
                 let source_variable = self.convert_ssa_value(*array, dfg);
                 let index_register = self.convert_ssa_single_addr_value(*index, dfg);
                 let value_variable = self.convert_ssa_value(*value, dfg);

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

@@ -530,7 +530,10 @@ impl<'dfg> InsertInstructionResult<'dfg> {
         match self {
             InsertInstructionResult::SimplifiedTo(value) => *value,
             InsertInstructionResult::SimplifiedToMultiple(values) => values[0],
-            InsertInstructionResult::Results(_, results) => results[0],
+            InsertInstructionResult::Results(_, results) => {
+                assert_eq!(results.len(), 1);
+                results[0]
+            }
             InsertInstructionResult::InstructionRemoved => {
                 panic!("Instruction was removed, no results")
             }

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

@@ -354,7 +354,9 @@ fn simplify_slice_push_back(
     slice_sizes.insert(set_last_slice_value, slice_size / element_size);
     slice_sizes.insert(new_slice, slice_size / element_size);
 
-    let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes);
+    let unknown = HashMap::default();
+    let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes, &unknown, None);
+
     let new_slice = value_merger.merge_values(
         len_not_equals_capacity,
         len_equals_capacity,

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

@@ -181,7 +181,7 @@ fn display_instruction_inner(
             let index = show(*index);
             let value = show(*value);
             let mutable = if *mutable { " mut" } else { "" };
-            writeln!(f, "array_set{mutable} {array}, index {index}, value {value}",)
+            writeln!(f, "array_set{mutable} {array}, index {index}, value {value}")
         }
         Instruction::IncrementRc { value } => {
             writeln!(f, "inc_rc {}", show(*value))

compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/value_merger.rs

@@ -1,29 +1,36 @@
 use acvm::FieldElement;
-use fxhash::FxHashMap as HashMap;
+use fxhash::{FxHashMap as HashMap, FxHashSet};
 
 use crate::ssa::ir::{
     basic_block::BasicBlockId,
-    dfg::{CallStack, DataFlowGraph},
+    dfg::{CallStack, DataFlowGraph, InsertInstructionResult},
     instruction::{BinaryOp, Instruction},
     types::Type,
-    value::ValueId,
+    value::{Value, ValueId},
 };
 
 pub(crate) struct ValueMerger<'a> {
     dfg: &'a mut DataFlowGraph,
     block: BasicBlockId,
+
+    current_condition: Option<ValueId>,
+
     // Maps SSA array values with a slice type to their size.
     // This must be computed before merging values.
     slice_sizes: &'a mut HashMap<ValueId, usize>,
+
+    array_set_conditionals: &'a HashMap<ValueId, ValueId>,
 }
 
 impl<'a> ValueMerger<'a> {
     pub(crate) fn new(
         dfg: &'a mut DataFlowGraph,
         block: BasicBlockId,
         slice_sizes: &'a mut HashMap<ValueId, usize>,
+        array_set_conditionals: &'a HashMap<ValueId, ValueId>,
+        current_condition: Option<ValueId>,
     ) -> Self {
-        ValueMerger { dfg, block, slice_sizes }
+        ValueMerger { dfg, block, slice_sizes, array_set_conditionals, current_condition }
     }
 
     /// Merge two values a and b from separate basic blocks to a single value.
@@ -131,6 +138,18 @@ impl<'a> ValueMerger<'a> {
             _ => panic!("Expected array type"),
         };
 
+        let actual_length = len * element_types.len();
+
+        if let Some(result) = self.try_merge_only_changed_indices(
+            then_condition,
+            else_condition,
+            then_value,
+            else_value,
+            actual_length,
+        ) {
+            return result;
+        }
+
         for i in 0..len {
             for (element_index, element_type) in element_types.iter().enumerate() {
                 let index = ((i * element_types.len() + element_index) as u128).into();
@@ -266,4 +285,94 @@ impl<'a> ValueMerger<'a> {
             }
         }
     }
+
+    fn try_merge_only_changed_indices(
+        &mut self,
+        then_condition: ValueId,
+        else_condition: ValueId,
+        then_value: ValueId,
+        else_value: ValueId,
+        array_length: usize,
+    ) -> Option<ValueId> {
+        let mut seen_then = FxHashSet::default();
+        let mut seen_else = FxHashSet::default();
+        let mut found = false;
+
+        let mut current_then = then_value;
+        let mut current_else = else_value;
+
+        let mut changed_indices = FxHashSet::default();
+
+        // Arbitrarily limit this to looking at at most 10 past ArraySet operations.
+        // If there are more than that, we assume 2 completely separate arrays are being merged.
+        for _ in 0..10 {
+            seen_then.insert(current_then);
+            seen_else.insert(current_else);
+
+            if seen_then.contains(&current_else) || seen_else.contains(&current_then) {
+                found = true;
+                break;
+            }
+
+            current_then = self.find_previous_array_set(current_then, &mut changed_indices)?;
+            current_else = self.find_previous_array_set(current_else, &mut changed_indices)?;
+        }
+
+        if !found || changed_indices.len() >= array_length {
+            return None;
+        }
+
+        let mut array = then_value;
+
+        for (index, element_type, condition) in changed_indices {
+            let typevars = Some(vec![element_type.clone()]);
+
+            let instruction = Instruction::EnableSideEffects { condition };
+            self.insert_instruction(instruction);
+
+            let mut get_element = |array, typevars| {
+                let get = Instruction::ArrayGet { array, index };
+                self.dfg
+                    .insert_instruction_and_results(get, self.block, typevars, CallStack::new())
+                    .first()
+            };
+
+            let then_element = get_element(then_value, typevars.clone());
+            let else_element = get_element(else_value, typevars);
+
+            let value =
+                self.merge_values(then_condition, else_condition, then_element, else_element);
+
+            let instruction = Instruction::ArraySet { array, index, value, mutable: false };
+            array = self.insert_instruction(instruction).first();
+        }
+
+        let instruction = Instruction::EnableSideEffects { condition: self.current_condition? };
+        self.insert_instruction(instruction);
+
+        Some(array)
+    }
+
+    fn insert_instruction(&mut self, instruction: Instruction) -> InsertInstructionResult {
+        self.dfg.insert_instruction_and_results(instruction, self.block, None, CallStack::new())
+    }
+
+    fn find_previous_array_set(
+        &self,
+        value: ValueId,
+        changed_indices: &mut FxHashSet<(ValueId, Type, ValueId)>,
+    ) -> Option<ValueId> {
+        match &self.dfg[value] {
+            Value::Instruction { instruction, .. } => match &self.dfg[*instruction] {
+                Instruction::ArraySet { array, index, value, .. } => {
+                    let condition = *self.array_set_conditionals.get(index)?;
+                    let element_type = self.dfg.type_of_value(*value);
+                    changed_indices.insert((*index, element_type, condition));
+                    Some(*array)
+                }
+                _ => Some(value),
+            },
+            _ => Some(value),
+        }
+    }
 }

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

@@ -42,12 +42,20 @@ impl Ssa {
 #[derive(Default)]
 struct Context {
     slice_sizes: HashMap<ValueId, usize>,
+
+    // Maps array_set result -> element that was overwritten by that instruction.
+    // Used to undo array_sets while merging values
+    prev_array_set_elem_values: HashMap<ValueId, ValueId>,
+
+    // Maps array_set result -> enable_side_effects_if value which was active during it.
+    array_set_conditionals: HashMap<ValueId, ValueId>,
 }
 
 impl Context {
     fn remove_if_else(&mut self, function: &mut Function) {
         let block = function.entry_block();
         let instructions = function.dfg[block].take_instructions();
+        let mut current_conditional = function.dfg.make_constant(FieldElement::one(), Type::bool());
 
         for instruction in instructions {
             match &function.dfg[instruction] {
@@ -60,8 +68,14 @@ impl Context {
                     let typ = function.dfg.type_of_value(then_value);
                     assert!(!matches!(typ, Type::Numeric(_)));
 
-                    let mut value_merger =
-                        ValueMerger::new(&mut function.dfg, block, &mut self.slice_sizes);
+                    let mut value_merger = ValueMerger::new(
+                        &mut function.dfg,
+                        block,
+                        &mut self.slice_sizes,
+                        &self.array_set_conditionals,
+                        Some(current_conditional),
+                    );
+
                     let value = value_merger.merge_values(
                         then_condition,
                         else_condition,
@@ -98,10 +112,15 @@ impl Context {
                     let results = function.dfg.instruction_results(instruction);
                     let result = if results.len() == 2 { results[1] } else { results[0] };
 
+                    self.array_set_conditionals.insert(*array, current_conditional);
                     let old_capacity = self.get_or_find_capacity(&function.dfg, *array);
                     self.slice_sizes.insert(result, old_capacity);
                     function.dfg[block].instructions_mut().push(instruction);
                 }
+                Instruction::EnableSideEffects { condition } => {
+                    current_conditional = *condition;
+                    function.dfg[block].instructions_mut().push(instruction);
+                }
                 _ => {
                     function.dfg[block].instructions_mut().push(instruction);
                 }