fix: Fix remaining issues in Brillig's copy on write arrays (#3593)

* Remove todo by going with new approach to Rc

* Very important commit: fix typo

* fix: make toBits accept arrays and add test

* fix: always codegen inc_rc

* IncrementRC has no side effects?

* Special case inc_rc instead

* IncRc is actually an even more special case of DIE. Also mark Value::Params as used

* Remove abbreviation

---------

Co-authored-by: sirasistant <sirasistant@gmail.com>

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -489,15 +489,18 @@ impl<'block> BrilligBlock<'block> {
                     );
                     let target_len = target_len_variable.extract_register();
 
-                    let target_vector = self
-                        .variables
-                        .define_variable(
-                            self.function_context,
-                            self.brillig_context,
-                            results[1],
-                            dfg,
-                        )
-                        .extract_vector();
+                    let target_vector = match self.variables.define_variable(
+                        self.function_context,
+                        self.brillig_context,
+                        results[1],
+                        dfg,
+                    ) {
+                        BrilligVariable::BrilligArray(array) => {
+                            self.brillig_context.array_to_vector(&array)
+                        }
+                        BrilligVariable::BrilligVector(vector) => vector,
+                        BrilligVariable::Simple(..) => unreachable!("ICE: ToBits on non-array"),
+                    };
 
                     let radix = self.brillig_context.make_constant(2_usize.into());
 

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

@@ -471,26 +471,10 @@ impl FunctionBuilder {
                     self.increment_array_reference_count(value);
                 }
             }
-            Type::Array(element_types, length) => {
+            Type::Array(..) | Type::Slice(..) => {
                 self.insert_instruction(Instruction::IncrementRc { value }, None);
-
-                if element_types.iter().any(|element| element.contains_an_array()) {
-                    for i in 0..length {
-                        let index = self.field_constant(i as u128);
-                        let element_type = element_types[i % element_types.len()].clone();
-                        let element = self.insert_array_get(value, index, element_type);
-                        self.increment_array_reference_count(element);
-                    }
-                }
-            }
-            Type::Slice(element_types) => {
-                self.insert_instruction(Instruction::IncrementRc { value }, None);
-
-                for element_type in element_types.iter() {
-                    if element_type.contains_an_array() {
-                        todo!("inc_rc for nested slices")
-                    }
-                }
+                // If there are nested arrays or slices, we wait until ArrayGet
+                // is issued to increment the count of that array.
             }
         }
     }

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

@@ -7,7 +7,7 @@ use crate::ssa::{
         basic_block::{BasicBlock, BasicBlockId},
         dfg::DataFlowGraph,
         function::Function,
-        instruction::InstructionId,
+        instruction::{Instruction, InstructionId},
         post_order::PostOrder,
         value::{Value, ValueId},
     },
@@ -38,13 +38,20 @@ fn dead_instruction_elimination(function: &mut Function) {
     for block in blocks.as_slice() {
         context.remove_unused_instructions_in_block(function, *block);
     }
+
+    context.remove_increment_rc_instructions(&mut function.dfg);
 }
 
 /// Per function context for tracking unused values and which instructions to remove.
 #[derive(Default)]
 struct Context {
     used_values: HashSet<ValueId>,
     instructions_to_remove: HashSet<InstructionId>,
+
+    /// IncrementRc instructions must be revisited after the main DIE pass since
+    /// they are technically side-effectful but we stil want to remove them if their
+    /// `value` parameter is not used elsewhere.
+    increment_rc_instructions: Vec<(InstructionId, BasicBlockId)>,
 }
 
 impl Context {
@@ -67,14 +74,19 @@ impl Context {
         let block = &function.dfg[block_id];
         self.mark_terminator_values_as_used(function, block);
 
-        for instruction in block.instructions().iter().rev() {
-            if self.is_unused(*instruction, function) {
-                self.instructions_to_remove.insert(*instruction);
+        for instruction_id in block.instructions().iter().rev() {
+            if self.is_unused(*instruction_id, function) {
+                self.instructions_to_remove.insert(*instruction_id);
             } else {
-                let instruction = &function.dfg[*instruction];
-                instruction.for_each_value(|value| {
-                    self.mark_used_instruction_results(&function.dfg, value);
-                });
+                let instruction = &function.dfg[*instruction_id];
+
+                if let Instruction::IncrementRc { .. } = instruction {
+                    self.increment_rc_instructions.push((*instruction_id, block_id));
+                } else {
+                    instruction.for_each_value(|value| {
+                        self.mark_used_instruction_results(&function.dfg, value);
+                    });
+                }
             }
         }
 
@@ -119,11 +131,28 @@ impl Context {
                     self.mark_used_instruction_results(dfg, *elem);
                 }
             }
+            Value::Param { .. } => {
+                self.used_values.insert(value_id);
+            }
             _ => {
                 // Does not comprise of any instruction results
             }
         }
     }
+
+    fn remove_increment_rc_instructions(self, dfg: &mut DataFlowGraph) {
+        for (increment_rc, block) in self.increment_rc_instructions {
+            let value = match &dfg[increment_rc] {
+                Instruction::IncrementRc { value } => *value,
+                other => unreachable!("Expected IncrementRc instruction, found {other:?}"),
+            };
+
+            // This could be more efficient if we have to remove multiple IncrementRcs in a single block
+            if !self.used_values.contains(&value) {
+                dfg[block].instructions_mut().retain(|instruction| *instruction != increment_rc);
+            }
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -210,6 +210,13 @@ impl<'a> FunctionContext<'a> {
         for element in elements {
             element.for_each(|element| {
                 let element = element.eval(self);
+
+                // If we're referencing a sub-array in a larger nested array we need to
+                // increase the reference count of the sub array. This maintains a
+                // pessimistic reference count (since some are likely moved rather than shared)
+                // which is important for Brillig's copy on write optimization. This has no
+                // effect in ACIR code.
+                self.builder.increment_array_reference_count(element);
                 array.push_back(element);
             });
         }
@@ -301,6 +308,7 @@ impl<'a> FunctionContext<'a> {
         } else {
             (array_or_slice[0], None)
         };
+
         self.codegen_array_index(
             array,
             index_value,
@@ -345,7 +353,13 @@ impl<'a> FunctionContext<'a> {
                 }
                 _ => unreachable!("must have array or slice but got {array_type}"),
             }
-            self.builder.insert_array_get(array, offset, typ).into()
+
+            // Reference counting in brillig relies on us incrementing reference
+            // counts when nested arrays/slices are constructed or indexed. This
+            // has no effect in ACIR code.
+            let result = self.builder.insert_array_get(array, offset, typ);
+            self.builder.increment_array_reference_count(result);
+            result.into()
         }))
     }
 

tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "brillig_cow"
+type = "bin"
+authors = [""]
+
+[dependencies]

tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml

@@ -0,0 +1,7 @@
+original = [0, 1, 2, 3, 4]
+index = 2
+
+[expected_result]
+original = [0, 1, 2, 3, 4]
+modified_once = [0, 1, 27, 3, 4]
+modified_twice = [0, 1, 27, 27, 4]

tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr

@@ -0,0 +1,54 @@
+// Tests the copy on write optimization for arrays. We look for cases where we are modifying an array in place when we shouldn't.
+
+global ARRAY_SIZE = 5;
+
+struct ExecutionResult {
+    original: [Field; ARRAY_SIZE],
+    modified_once: [Field; ARRAY_SIZE],
+    modified_twice: [Field; ARRAY_SIZE],
+}
+
+impl ExecutionResult {
+    fn is_equal(self, other: ExecutionResult) -> bool {
+        (self.original == other.original) &
+        (self.modified_once == other.modified_once) &
+        (self.modified_twice == other.modified_twice)
+    }
+}
+
+fn modify_in_inlined_constrained(original: [Field; ARRAY_SIZE], index: u64) -> ExecutionResult {
+    let mut modified = original;
+
+    modified[index] = 27;
+
+    let modified_once = modified;
+
+    modified[index+1] = 27;
+
+    ExecutionResult {
+        original,
+        modified_once,
+        modified_twice: modified
+    }
+}
+
+unconstrained fn modify_in_unconstrained(original: [Field; ARRAY_SIZE], index: u64) -> ExecutionResult {
+    let mut modified = original;
+
+    modified[index] = 27;
+
+    let modified_once = modified;
+
+    modified[index+1] = 27;
+
+    ExecutionResult {
+        original,
+        modified_once,
+        modified_twice: modified
+    }
+}
+
+unconstrained fn main(original: [Field; ARRAY_SIZE], index: u64, expected_result: ExecutionResult) {
+    assert(expected_result.is_equal(modify_in_unconstrained(original, index)));
+    assert(expected_result.is_equal(modify_in_inlined_constrained(original, index)));
+}