Generalize inc_rc elision for parameters

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/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 into_contained_array(&self) -> &Type {
+        match self {
+            Type::Numeric(_) | Type::Function => panic!("Expected an array type"),
+            Type::Array(_, _) | Type::Slice(_) => self,
+            Type::Reference(element) => element.into_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

@@ -1179,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();
@@ -1188,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/ssa_gen/context.rs

@@ -20,7 +20,7 @@ use crate::ssa::ir::value::ValueId;
 
 use super::value::{Tree, Value, Values};
 use super::SSA_WORD_SIZE;
-use fxhash::FxHashMap as HashMap;
+use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
 
 /// The FunctionContext is the main context object for translating a
 /// function into SSA form during the SSA-gen pass.
@@ -912,42 +912,50 @@ impl<'a> FunctionContext<'a> {
         }
     }
 
-    /// Increments the reference count of all parameters. Returns the entry block of the function.
+    /// Increments the reference count of mutable array parameters.
+    /// Returns each array id that was incremented.
     ///
     /// This is done on parameters rather than call arguments so that we can optimize out
     /// paired inc/dec instructions within brillig functions more easily.
-    pub(crate) fn increment_parameter_rcs(&mut self) -> BasicBlockId {
+    pub(crate) fn increment_parameter_rcs(&mut self) -> HashSet<ValueId> {
         let entry = self.builder.current_function.entry_block();
         let parameters = self.builder.current_function.dfg.block_parameters(entry).to_vec();
 
-        // Performance hack: always consider the first parameter to be a move
-        for parameter in parameters.into_iter().skip(1) {
+        let mut incremented = HashSet::default();
+        let mut seen_array_types = HashSet::default();
+
+        for parameter in parameters {
             // Avoid reference counts for immutable arrays that aren't behind references.
-            if self.builder.current_function.dfg.value_is_reference(parameter) {
-                self.builder.increment_array_reference_count(parameter);
+            let typ = self.builder.current_function.dfg.type_of_value(parameter);
+
+            if let Type::Reference(element) = typ {
+                if element.contains_an_array() {
+                    // If we haven't already seen this array type, the value may be possibly
+                    // aliased, so issue an inc_rc for it.
+                    if !seen_array_types.insert(element.into_contained_array().clone()) {
+                        if self.builder.increment_array_reference_count(parameter) {
+                            incremented.insert(parameter);
+                        }
+                    }
+                }
             }
         }
 
-        entry
+        incremented
     }
 
     /// Ends a local scope of a function.
     /// This will issue DecrementRc instructions for any arrays in the given starting scope
     /// block's parameters. Arrays that are also used in terminator instructions for the scope are
     /// ignored.
-    pub(crate) fn end_scope(&mut self, scope: BasicBlockId, terminator_args: &[ValueId]) {
-        let mut dropped_parameters = self.builder.current_function.dfg.block_parameters(scope);
-
-        // Skip the first parameter to match the check in `increment_parameter_rcs`
-        if !dropped_parameters.is_empty() {
-            dropped_parameters = &dropped_parameters[1..];
-        }
-
-        let mut dropped_parameters = dropped_parameters.to_vec();
-
-        dropped_parameters.retain(|parameter| !terminator_args.contains(parameter));
+    pub(crate) fn end_scope(
+        &mut self,
+        mut incremented_params: HashSet<ValueId>,
+        terminator_args: &[ValueId],
+    ) {
+        incremented_params.retain(|parameter| !terminator_args.contains(parameter));
 
-        for parameter in dropped_parameters {
+        for parameter in incremented_params {
             if self.builder.current_function.dfg.value_is_reference(parameter) {
                 self.builder.decrement_array_reference_count(parameter);
             }

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

@@ -125,10 +125,10 @@ impl<'a> FunctionContext<'a> {
     /// Codegen a function's body and set its return value to that of its last parameter.
     /// For functions returning nothing, this will be an empty list.
     fn codegen_function_body(&mut self, body: &Expression) -> Result<(), RuntimeError> {
-        let entry_block = self.increment_parameter_rcs();
+        let incremented_params = self.increment_parameter_rcs();
         let return_value = self.codegen_expression(body)?;
         let results = return_value.into_value_list(self);
-        self.end_scope(entry_block, &results);
+        self.end_scope(incremented_params, &results);
 
         self.builder.terminate_with_return(results);
         Ok(())