feat(perf): Handle array set optimization across blocks for Brillig f…

…unctions (#6153)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

Part of general effort to reduce Brillig bytecode sizes

## Summary\*

We now go through all blocks of Brillig functions to determine whether
we can mutate an array directly rather than copy it.

We are still pretty conservative with our checks to make sure that we do
not inadvertently mark the wrong array set mutable.
- If we have seen a last array set whose array comes from a Load result
we do not make that array set mutable unless we are in a return block.
We can expand upon the cases in which we mark an array loaded from a
reference as mutable in follow-ups.
- If the array is used in the terminator of a block we do not mark its
array set mutable.

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

---------

Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>

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

@@ -2,12 +2,13 @@ use crate::ssa::{
     ir::{
         basic_block::BasicBlockId,
         dfg::DataFlowGraph,
-        instruction::{Instruction, InstructionId},
+        instruction::{Instruction, InstructionId, TerminatorInstruction},
         types::Type::{Array, Slice},
+        value::ValueId,
     },
     ssa_gen::Ssa,
 };
-use fxhash::{FxHashMap as HashMap, FxHashSet};
+use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
 
 impl Ssa {
     /// Map arrays with the last instruction that uses it
@@ -16,28 +17,42 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn array_set_optimization(mut self) -> Self {
         for func in self.functions.values_mut() {
-            let mut reachable_blocks = func.reachable_blocks();
-            let block = if !func.runtime().is_entry_point() {
+            let reachable_blocks = func.reachable_blocks();
+
+            if !func.runtime().is_entry_point() {
                 assert_eq!(reachable_blocks.len(), 1, "Expected there to be 1 block remaining in Acir function for array_set optimization");
-                reachable_blocks.pop_first().unwrap()
-            } else {
-                // We only apply the array set optimization in the return block of Brillig functions
-                func.find_last_block()
-            };
+            }
+            let mut array_to_last_use = HashMap::default();
+            let mut instructions_to_update = HashSet::default();
+            let mut arrays_from_load = HashSet::default();
 
-            let instructions_to_update = analyze_last_uses(&func.dfg, block);
-            make_mutable(&mut func.dfg, block, instructions_to_update);
+            for block in reachable_blocks.iter() {
+                analyze_last_uses(
+                    &func.dfg,
+                    *block,
+                    &mut array_to_last_use,
+                    &mut instructions_to_update,
+                    &mut arrays_from_load,
+                );
+            }
+            for block in reachable_blocks {
+                make_mutable(&mut func.dfg, block, &instructions_to_update);
+            }
         }
         self
     }
 }
 
-/// Returns the set of ArraySet instructions that can be made mutable
+/// Builds the set of ArraySet instructions that can be made mutable
 /// because their input value is unused elsewhere afterward.
-fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet<InstructionId> {
+fn analyze_last_uses(
+    dfg: &DataFlowGraph,
+    block_id: BasicBlockId,
+    array_to_last_use: &mut HashMap<ValueId, InstructionId>,
+    instructions_that_can_be_made_mutable: &mut HashSet<InstructionId>,
+    arrays_from_load: &mut HashSet<ValueId>,
+) {
     let block = &dfg[block_id];
-    let mut array_to_last_use = HashMap::default();
-    let mut instructions_that_can_be_made_mutable = FxHashSet::default();
 
     for instruction_id in block.instructions() {
         match &dfg[*instruction_id] {
@@ -54,7 +69,22 @@ fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet<I
                 if let Some(existing) = array_to_last_use.insert(array, *instruction_id) {
                     instructions_that_can_be_made_mutable.remove(&existing);
                 }
-                instructions_that_can_be_made_mutable.insert(*instruction_id);
+                // If the array we are setting does not come from a load we can safely mark it mutable.
+                // If the array comes from a load we may potentially being mutating an array at a reference
+                // that is loaded from by other values.
+                let terminator = dfg[block_id].unwrap_terminator();
+                // If we are in a return block we are not concerned about the array potentially being mutated again.
+                let is_return_block = matches!(terminator, TerminatorInstruction::Return { .. });
+                // We also want to check that the array is not part of the terminator arguments, as this means it is used again.
+                let mut array_in_terminator = false;
+                terminator.for_each_value(|value| {
+                    if value == array {
+                        array_in_terminator = true;
+                    }
+                });
+                if (!arrays_from_load.contains(&array) || is_return_block) && !array_in_terminator {
+                    instructions_that_can_be_made_mutable.insert(*instruction_id);
+                }
             }
             Instruction::Call { arguments, .. } => {
                 for argument in arguments {
@@ -68,18 +98,22 @@ fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet<I
                     }
                 }
             }
+            Instruction::Load { .. } => {
+                let result = dfg.instruction_results(*instruction_id)[0];
+                if matches!(dfg.type_of_value(result), Array { .. } | Slice { .. }) {
+                    arrays_from_load.insert(result);
+                }
+            }
             _ => (),
         }
     }
-
-    instructions_that_can_be_made_mutable
 }
 
 /// Make each ArraySet instruction in `instructions_to_update` mutable.
 fn make_mutable(
     dfg: &mut DataFlowGraph,
     block_id: BasicBlockId,
-    instructions_to_update: FxHashSet<InstructionId>,
+    instructions_to_update: &HashSet<InstructionId>,
 ) {
     if instructions_to_update.is_empty() {
         return;
@@ -105,3 +139,129 @@ fn make_mutable(
 
     *dfg[block_id].instructions_mut() = instructions;
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use im::vector;
+
+    use crate::ssa::{
+        function_builder::FunctionBuilder,
+        ir::{
+            function::RuntimeType,
+            instruction::{BinaryOp, Instruction},
+            map::Id,
+            types::Type,
+        },
+    };
+
+    #[test]
+    fn array_set_in_loop_with_conditional_clone() {
+        // We want to make sure that we do not mark a single array set mutable which is loaded
+        // from and cloned in a loop. If the array is inadvertently marked mutable, and is cloned in a previous iteration
+        // of the loop, its clone will also be altered.
+        //
+        // acir(inline) fn main f0 {
+        //     b0():
+        //       v2 = allocate
+        //       store [Field 0, Field 0, Field 0, Field 0, Field 0] at v2
+        //       v3 = allocate
+        //       store [Field 0, Field 0, Field 0, Field 0, Field 0] at v3
+        //       jmp b1(u32 0)
+        //     b1(v5: u32):
+        //       v7 = lt v5, u32 5
+        //       jmpif v7 then: b3, else: b2
+        //     b3():
+        //       v8 = eq v5, u32 5
+        //       jmpif v8 then: b4, else: b5
+        //     b4():
+        //       v9 = load v2
+        //       store v9 at v3
+        //       jmp b5()
+        //     b5():
+        //       v10 = load v2
+        //       v12 = array_set v10, index v5, value Field 20
+        //       store v12 at v2
+        //       v14 = add v5, u32 1
+        //       jmp b1(v14)
+        //     b2():
+        //       return
+        //   }
+        let main_id = Id::test_new(0);
+        let mut builder = FunctionBuilder::new("main".into(), main_id);
+        builder.set_runtime(RuntimeType::Brillig);
+
+        let array_type = Type::Array(Arc::new(vec![Type::field()]), 5);
+        let zero = builder.field_constant(0u128);
+        let array_constant =
+            builder.array_constant(vector![zero, zero, zero, zero, zero], array_type.clone());
+
+        let v2 = builder.insert_allocate(array_type.clone());
+
+        builder.insert_store(v2, array_constant);
+
+        let v3 = builder.insert_allocate(array_type.clone());
+        builder.insert_store(v3, array_constant);
+
+        let b1 = builder.insert_block();
+        let zero_u32 = builder.numeric_constant(0u128, Type::unsigned(32));
+        builder.terminate_with_jmp(b1, vec![zero_u32]);
+
+        // Loop header
+        builder.switch_to_block(b1);
+        let v5 = builder.add_block_parameter(b1, Type::unsigned(32));
+        let five = builder.numeric_constant(5u128, Type::unsigned(32));
+        let v7 = builder.insert_binary(v5, BinaryOp::Lt, five);
+
+        let b2 = builder.insert_block();
+        let b3 = builder.insert_block();
+        let b4 = builder.insert_block();
+        let b5 = builder.insert_block();
+        builder.terminate_with_jmpif(v7, b3, b2);
+
+        // Loop body
+        // b3 is the if statement conditional
+        builder.switch_to_block(b3);
+        let two = builder.numeric_constant(5u128, Type::unsigned(32));
+        let v8 = builder.insert_binary(v5, BinaryOp::Eq, two);
+        builder.terminate_with_jmpif(v8, b4, b5);
+
+        // b4 is the rest of the loop after the if statement
+        builder.switch_to_block(b4);
+        let v9 = builder.insert_load(v2, array_type.clone());
+        builder.insert_store(v3, v9);
+        builder.terminate_with_jmp(b5, vec![]);
+
+        builder.switch_to_block(b5);
+        let v10 = builder.insert_load(v2, array_type.clone());
+        let twenty = builder.field_constant(20u128);
+        let v12 = builder.insert_array_set(v10, v5, twenty);
+        builder.insert_store(v2, v12);
+        let one = builder.numeric_constant(1u128, Type::unsigned(32));
+        let v14 = builder.insert_binary(v5, BinaryOp::Add, one);
+        builder.terminate_with_jmp(b1, vec![v14]);
+
+        builder.switch_to_block(b2);
+        builder.terminate_with_return(vec![]);
+
+        let ssa = builder.finish();
+        // We expect the same result as above
+        let ssa = ssa.array_set_optimization();
+
+        let main = ssa.main();
+        assert_eq!(main.reachable_blocks().len(), 6);
+
+        let array_set_instructions = main.dfg[b5]
+            .instructions()
+            .iter()
+            .filter(|instruction| matches!(&main.dfg[**instruction], Instruction::ArraySet { .. }))
+            .collect::<Vec<_>>();
+
+        assert_eq!(array_set_instructions.len(), 1);
+        if let Instruction::ArraySet { mutable, .. } = &main.dfg[*array_set_instructions[0]] {
+            // The single array set should not be marked mutable
+            assert!(!mutable);
+        }
+    }
+}