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

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);