fix: Optimize array ref counts to copy arrays much less often (#6685)

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

compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/array_copy.rs

@@ -69,6 +69,8 @@ pub(super) fn compile_array_copy_procedure<F: AcirField + DebugToString>(
                 BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
                 1_usize.into(),
             );
+            // Decrease the original ref count now that this copy is no longer pointing to it
+            ctx.codegen_usize_op(rc.address, rc.address, BrilligBinaryOp::Sub, 1);
         }
     });
 }

compiler/noirc_evaluator/src/ssa/function_builder/mod.rs

@@ -441,29 +441,38 @@ impl FunctionBuilder {
     /// Insert instructions to increment the reference count of any array(s) stored
     /// within the given value. If the given value is not an array and does not contain
     /// any arrays, this does nothing.
-    pub(crate) fn increment_array_reference_count(&mut self, value: ValueId) {
-        self.update_array_reference_count(value, true);
+    ///
+    /// Returns whether a reference count instruction was issued.
+    pub(crate) fn increment_array_reference_count(&mut self, value: ValueId) -> bool {
+        self.update_array_reference_count(value, true)
     }
 
     /// Insert instructions to decrement the reference count of any array(s) stored
     /// within the given value. If the given value is not an array and does not contain
     /// any arrays, this does nothing.
-    pub(crate) fn decrement_array_reference_count(&mut self, value: ValueId) {
-        self.update_array_reference_count(value, false);
+    ///
+    /// Returns whether a reference count instruction was issued.
+    pub(crate) fn decrement_array_reference_count(&mut self, value: ValueId) -> bool {
+        self.update_array_reference_count(value, false)
     }
 
     /// Increment or decrement the given value's reference count if it is an array.
     /// If it is not an array, this does nothing. Note that inc_rc and dec_rc instructions
     /// are ignored outside of unconstrained code.
-    fn update_array_reference_count(&mut self, value: ValueId, increment: bool) {
+    ///
+    /// Returns whether a reference count instruction was issued.
+    fn update_array_reference_count(&mut self, value: ValueId, increment: bool) -> bool {
         match self.type_of_value(value) {
-            Type::Numeric(_) => (),
-            Type::Function => (),
+            Type::Numeric(_) => false,
+            Type::Function => false,
             Type::Reference(element) => {
                 if element.contains_an_array() {
                     let reference = value;
                     let value = self.insert_load(reference, element.as_ref().clone());
                     self.update_array_reference_count(value, increment);
+                    true
+                } else {
+                    false
                 }
             }
             Type::Array(..) | Type::Slice(..) => {
@@ -474,6 +483,7 @@ impl FunctionBuilder {
                 } else {
                     self.insert_dec_rc(value);
                 }
+                true
             }
         }
     }

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

@@ -129,7 +129,7 @@ impl<'f> FunctionInserter<'f> {
         // another MakeArray instruction. Note that this assumes the function inserter is inserting
         // in control-flow order. Otherwise we could refer to ValueIds defined later in the program.
         let make_array = if let Instruction::MakeArray { elements, typ } = &instruction {
-            if self.array_is_constant(elements) {
+            if self.array_is_constant(elements) && self.function.runtime().is_acir() {
                 if let Some(fetched_value) = self.get_cached_array(elements, typ) {
                     assert_eq!(results.len(), 1);
                     self.values.insert(results[0], fetched_value);

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

@@ -429,7 +429,7 @@ impl Instruction {
     /// conditional on whether the caller wants the predicate to be taken into account or not.
     pub(crate) fn can_be_deduplicated(
         &self,
-        dfg: &DataFlowGraph,
+        function: &Function,
         deduplicate_with_predicate: bool,
     ) -> bool {
         use Instruction::*;
@@ -443,7 +443,7 @@ impl Instruction {
             | IncrementRc { .. }
             | DecrementRc { .. } => false,
 
-            Call { func, .. } => match dfg[*func] {
+            Call { func, .. } => match function.dfg[*func] {
                 Value::Intrinsic(intrinsic) => {
                     intrinsic.can_be_deduplicated(deduplicate_with_predicate)
                 }
@@ -453,8 +453,11 @@ impl Instruction {
             // We can deduplicate these instructions if we know the predicate is also the same.
             Constrain(..) | RangeCheck { .. } => deduplicate_with_predicate,
 
-            // This should never be side-effectful
-            MakeArray { .. } => true,
+            // Arrays can be mutated in unconstrained code so code that handles this case must
+            // take care to track whether the array was possibly mutated or not before
+            // deduplicating. Since we don't know if the containing pass checks for this, we
+            // can only assume these are safe to deduplicate in constrained code.
+            MakeArray { .. } => function.runtime().is_acir(),
 
             // These can have different behavior depending on the EnableSideEffectsIf context.
             // Replacing them with a similar instruction potentially enables replacing an instruction
@@ -467,7 +470,7 @@ impl Instruction {
             | IfElse { .. }
             | ArrayGet { .. }
             | ArraySet { .. } => {
-                deduplicate_with_predicate || !self.requires_acir_gen_predicate(dfg)
+                deduplicate_with_predicate || !self.requires_acir_gen_predicate(&function.dfg)
             }
         }
     }

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

@@ -190,6 +190,15 @@ impl Type {
         }
     }
 
+    /// Retrieves the array or slice type within this type, or panics if there is none.
+    pub(crate) fn get_contained_array(&self) -> &Type {
+        match self {
+            Type::Numeric(_) | Type::Function => panic!("Expected an array type"),
+            Type::Array(_, _) | Type::Slice(_) => self,
+            Type::Reference(element) => element.get_contained_array(),
+        }
+    }
+
     pub(crate) fn element_types(self) -> Arc<Vec<Type>> {
         match self {
             Type::Array(element_types, _) | Type::Slice(element_types) => element_types,

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

@@ -159,7 +159,7 @@ impl Function {
             }
 
             context.visited_blocks.insert(block);
-            context.fold_constants_in_block(&mut self.dfg, &mut dom, block);
+            context.fold_constants_in_block(self, &mut dom, block);
         }
     }
 }
@@ -266,36 +266,38 @@ impl<'brillig> Context<'brillig> {
 
     fn fold_constants_in_block(
         &mut self,
-        dfg: &mut DataFlowGraph,
+        function: &mut Function,
         dom: &mut DominatorTree,
         block: BasicBlockId,
     ) {
-        let instructions = dfg[block].take_instructions();
+        let instructions = function.dfg[block].take_instructions();
 
         // Default side effect condition variable with an enabled state.
-        let mut side_effects_enabled_var = dfg.make_constant(FieldElement::one(), Type::bool());
+        let mut side_effects_enabled_var =
+            function.dfg.make_constant(FieldElement::one(), Type::bool());
 
         for instruction_id in instructions {
             self.fold_constants_into_instruction(
-                dfg,
+                function,
                 dom,
                 block,
                 instruction_id,
                 &mut side_effects_enabled_var,
             );
         }
-        self.block_queue.extend(dfg[block].successors());
+        self.block_queue.extend(function.dfg[block].successors());
     }
 
     fn fold_constants_into_instruction(
         &mut self,
-        dfg: &mut DataFlowGraph,
+        function: &mut Function,
         dom: &mut DominatorTree,
         mut block: BasicBlockId,
         id: InstructionId,
         side_effects_enabled_var: &mut ValueId,
     ) {
         let constraint_simplification_mapping = self.get_constraint_map(*side_effects_enabled_var);
+        let dfg = &mut function.dfg;
 
         let instruction =
             Self::resolve_instruction(id, block, dfg, dom, constraint_simplification_mapping);
@@ -308,6 +310,15 @@ impl<'brillig> Context<'brillig> {
         {
             match cache_result {
                 CacheResult::Cached(cached) => {
+                    // We track whether we may mutate MakeArray instructions before we deduplicate
+                    // them but we still need to issue an extra inc_rc in case they're mutated afterward.
+                    if matches!(instruction, Instruction::MakeArray { .. }) {
+                        let value = *cached.last().unwrap();
+                        let inc_rc = Instruction::IncrementRc { value };
+                        let call_stack = dfg.get_call_stack(id);
+                        dfg.insert_instruction_and_results(inc_rc, block, None, call_stack);
+                    }
+
                     Self::replace_result_ids(dfg, &old_results, cached);
                     return;
                 }
@@ -321,32 +332,25 @@ impl<'brillig> Context<'brillig> {
             }
         };
 
-        let new_results =
         // First try to inline a call to a brillig function with all constant arguments.
-        Self::try_inline_brillig_call_with_all_constants(
+        let new_results = Self::try_inline_brillig_call_with_all_constants(
             &instruction,
             &old_results,
             block,
             dfg,
             self.brillig_info,
         )
+        // Otherwise, try inserting the instruction again to apply any optimizations using the newly resolved inputs.
         .unwrap_or_else(|| {
-            // Otherwise, try inserting the instruction again to apply any optimizations using the newly resolved inputs.
-            Self::push_instruction(
-                id,
-                instruction.clone(),
-                &old_results,
-                block,
-                dfg,
-            )
+            Self::push_instruction(id, instruction.clone(), &old_results, block, dfg)
         });
 
         Self::replace_result_ids(dfg, &old_results, &new_results);
 
         self.cache_instruction(
             instruction.clone(),
             new_results,
-            dfg,
+            function,
             *side_effects_enabled_var,
             block,
         );
@@ -433,7 +437,7 @@ impl<'brillig> Context<'brillig> {
         &mut self,
         instruction: Instruction,
         instruction_results: Vec<ValueId>,
-        dfg: &DataFlowGraph,
+        function: &Function,
         side_effects_enabled_var: ValueId,
         block: BasicBlockId,
     ) {
@@ -442,11 +446,11 @@ impl<'brillig> Context<'brillig> {
             // to map from the more complex to the simpler value.
             if let Instruction::Constrain(lhs, rhs, _) = instruction {
                 // These `ValueId`s should be fully resolved now.
-                if let Some((complex, simple)) = simplify(dfg, lhs, rhs) {
+                if let Some((complex, simple)) = simplify(&function.dfg, lhs, rhs) {
                     self.get_constraint_map(side_effects_enabled_var)
                         .entry(complex)
                         .or_default()
-                        .add(dfg, simple, block);
+                        .add(&function.dfg, simple, block);
                 }
             }
         }
@@ -463,7 +467,7 @@ impl<'brillig> Context<'brillig> {
         // we can simplify the operation when we take into account the predicate.
         if let Instruction::ArraySet { index, value, .. } = &instruction {
             let use_predicate =
-                self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg);
+                self.use_constraint_info && instruction.requires_acir_gen_predicate(&function.dfg);
             let predicate = use_predicate.then_some(side_effects_enabled_var);
 
             let array_get = Instruction::ArrayGet { array: instruction_results[0], index: *index };
@@ -476,12 +480,19 @@ impl<'brillig> Context<'brillig> {
                 .cache(block, vec![*value]);
         }
 
+        self.remove_possibly_mutated_cached_make_arrays(&instruction, function);
+
         // 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.
         // Others have side effects representing failure, which are implicit in the ACIR code and can also be deduplicated.
-        if instruction.can_be_deduplicated(dfg, self.use_constraint_info) {
+        let can_be_deduplicated =
+            instruction.can_be_deduplicated(function, self.use_constraint_info);
+
+        // We also allow deduplicating MakeArray instructions that we have tracked which haven't
+        // been mutated.
+        if can_be_deduplicated || matches!(instruction, Instruction::MakeArray { .. }) {
             let use_predicate =
-                self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg);
+                self.use_constraint_info && instruction.requires_acir_gen_predicate(&function.dfg);
             let predicate = use_predicate.then_some(side_effects_enabled_var);
 
             self.cached_instruction_results
@@ -688,6 +699,26 @@ impl<'brillig> Context<'brillig> {
             }
         }
     }
+
+    fn remove_possibly_mutated_cached_make_arrays(
+        &mut self,
+        instruction: &Instruction,
+        function: &Function,
+    ) {
+        use Instruction::{ArraySet, Store};
+
+        // Should we consider calls to slice_push_back and similar to be mutating operations as well?
+        if let Store { value: array, .. } | ArraySet { array, .. } = instruction {
+            let instruction = match &function.dfg[*array] {
+                Value::Instruction { instruction, .. } => &function.dfg[*instruction],
+                _ => return,
+            };
+
+            if matches!(instruction, Instruction::MakeArray { .. }) {
+                self.cached_instruction_results.remove(instruction);
+            }
+        }
+    }
 }
 
 impl ResultCache {
@@ -1148,6 +1179,7 @@ mod test {
         // fn main f0 {
         //   b0(v0: u64):
         //     v1 = make_array [v0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0, u64 0]
+        //     inc_rc v1
         //     v5 = call keccakf1600(v1)
         // }
         let ssa = ssa.fold_constants();
@@ -1157,7 +1189,7 @@ mod test {
         let main = ssa.main();
         let instructions = main.dfg[main.entry_block()].instructions();
         let ending_instruction_count = instructions.len();
-        assert_eq!(ending_instruction_count, 2);
+        assert_eq!(ending_instruction_count, 3);
     }
 
     #[test]

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

@@ -189,8 +189,7 @@ impl<'f> LoopInvariantContext<'f> {
                 !self.defined_in_loop.contains(&value) || self.loop_invariants.contains(&value);
         });
 
-        let can_be_deduplicated = instruction
-            .can_be_deduplicated(&self.inserter.function.dfg, false)
+        let can_be_deduplicated = instruction.can_be_deduplicated(self.inserter.function, false)
             || self.can_be_deduplicated_from_upper_bound(&instruction);
 
         is_loop_invariant && can_be_deduplicated