feat(perf): Allow array set last uses optimization in return block of…

… Brillig functions (#6119)

# Description

## Problem\*

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

Part of general effort to reduce Brillig bytecode sizes.

## Summary\*

We currently only perform the array set last uses optimization on ACIR
functions. This makes sense as the pass only looks only last uses in a
single block. We can start to expand the last uses opt to encompass more
than just ACIR functions and make it aware of the last array set of an
array across blocks. This is quite complex though so this PR currently
just attempts to look at the return block of Brillig functions and see
the effects.

## 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/brillig/brillig_gen/brillig_block.rs

@@ -666,7 +666,7 @@ impl<'block> BrilligBlock<'block> {
 
                 self.brillig_context.deallocate_register(items_pointer);
             }
-            Instruction::ArraySet { array, index, value, mutable: _ } => {
+            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);
@@ -688,6 +688,7 @@ impl<'block> BrilligBlock<'block> {
                     destination_variable,
                     index_register,
                     value_variable,
+                    *mutable,
                 );
             }
             Instruction::RangeCheck { value, max_bit_size, assert_message } => {
@@ -868,34 +869,49 @@ impl<'block> BrilligBlock<'block> {
         destination_variable: BrilligVariable,
         index_register: SingleAddrVariable,
         value_variable: BrilligVariable,
+        mutable: bool,
     ) {
         assert!(index_register.bit_size == BRILLIG_MEMORY_ADDRESSING_BIT_SIZE);
         match (source_variable, destination_variable) {
             (
                 BrilligVariable::BrilligArray(source_array),
                 BrilligVariable::BrilligArray(destination_array),
             ) => {
-                self.brillig_context.call_array_copy_procedure(source_array, destination_array);
+                if !mutable {
+                    self.brillig_context.call_array_copy_procedure(source_array, destination_array);
+                }
             }
             (
                 BrilligVariable::BrilligVector(source_vector),
                 BrilligVariable::BrilligVector(destination_vector),
             ) => {
-                self.brillig_context.call_vector_copy_procedure(source_vector, destination_vector);
+                if !mutable {
+                    self.brillig_context
+                        .call_vector_copy_procedure(source_vector, destination_vector);
+                }
             }
             _ => unreachable!("ICE: array set on non-array"),
         }
 
+        let destination_for_store = if mutable { source_variable } else { destination_variable };
         // Then set the value in the newly created array
         let items_pointer =
-            self.brillig_context.codegen_make_array_or_vector_items_pointer(destination_variable);
+            self.brillig_context.codegen_make_array_or_vector_items_pointer(destination_for_store);
 
         self.brillig_context.codegen_store_with_offset(
             items_pointer,
             index_register,
             value_variable.extract_register(),
         );
 
+        // If we mutated the source array we want instructions that use the destination array to point to the source array
+        if mutable {
+            self.brillig_context.mov_instruction(
+                destination_variable.extract_register(),
+                source_variable.extract_register(),
+            );
+        }
+
         self.brillig_context.deallocate_register(items_pointer);
     }
 

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

@@ -160,6 +160,17 @@ impl Function {
         let returns = vecmap(self.returns(), |ret| self.dfg.type_of_value(*ret));
         Signature { params, returns }
     }
+
+    /// Finds the block of the function with the Return instruction
+    pub(crate) fn find_last_block(&self) -> BasicBlockId {
+        for block in self.reachable_blocks() {
+            if matches!(self.dfg[block].terminator(), Some(TerminatorInstruction::Return { .. })) {
+                return block;
+            }
+        }
+
+        unreachable!("SSA Function {} has no reachable return instruction!", self.id())
+    }
 }
 
 impl std::fmt::Display for RuntimeType {

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

@@ -16,14 +16,17 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn array_set_optimization(mut self) -> Self {
         for func in self.functions.values_mut() {
-            if !func.runtime().is_entry_point() {
-                let mut reachable_blocks = func.reachable_blocks();
+            let mut reachable_blocks = func.reachable_blocks();
+            let block = 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 block = reachable_blocks.pop_first().unwrap();
-                let instructions_to_update = analyze_last_uses(&func.dfg, block);
-                make_mutable(&mut func.dfg, block, instructions_to_update);
-            }
+            let instructions_to_update = analyze_last_uses(&func.dfg, block);
+            make_mutable(&mut func.dfg, block, instructions_to_update);
         }
         self
     }

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

@@ -2,9 +2,8 @@ use std::collections::{HashMap, HashSet};
 
 use crate::ssa::{
     ir::{
-        basic_block::BasicBlockId,
         function::Function,
-        instruction::{Instruction, InstructionId, TerminatorInstruction},
+        instruction::{Instruction, InstructionId},
         types::Type,
         value::ValueId,
     },
@@ -107,7 +106,7 @@ impl Context {
     /// Find each dec_rc instruction and if the most recent inc_rc instruction for the same value
     /// is not possibly mutated, then we can remove them both. Returns each such pair.
     fn find_rcs_to_remove(&mut self, function: &Function) -> HashSet<InstructionId> {
-        let last_block = Self::find_last_block(function);
+        let last_block = function.find_last_block();
         let mut to_remove = HashSet::new();
 
         for instruction in function.dfg[last_block].instructions() {
@@ -124,20 +123,6 @@ impl Context {
         to_remove
     }
 
-    /// Finds the block of the function with the Return instruction
-    fn find_last_block(function: &Function) -> BasicBlockId {
-        for block in function.reachable_blocks() {
-            if matches!(
-                function.dfg[block].terminator(),
-                Some(TerminatorInstruction::Return { .. })
-            ) {
-                return block;
-            }
-        }
-
-        unreachable!("SSA Function {} has no reachable return instruction!", function.id())
-    }
-
     /// Finds and pops the IncRc for the given array value if possible.
     fn pop_rc_for(&mut self, value: ValueId, function: &Function) -> Option<IncRc> {
         let typ = function.dfg.type_of_value(value);