diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index d5e958f032a..0e06a36fd94 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/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()); diff --git a/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs b/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs index 0c3ff08557e..144e37a3f31 100644 --- a/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs +++ b/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. } } } diff --git a/compiler/noirc_evaluator/src/ssa/opt/die.rs b/compiler/noirc_evaluator/src/ssa/opt/die.rs index 9bc2f15abbb..53cdf72bbbf 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/die.rs +++ b/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,6 +38,8 @@ 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. @@ -45,6 +47,11 @@ fn dead_instruction_elimination(function: &mut Function) { struct Context { used_values: HashSet, instructions_to_remove: HashSet, + + /// 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)] diff --git a/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs b/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs index 3a4edf47c81..53f1bc863be 100644 --- a/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs +++ b/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() })) } diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml b/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml new file mode 100644 index 00000000000..d191eb53ddf --- /dev/null +++ b/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml @@ -0,0 +1,6 @@ +[package] +name = "brillig_cow" +type = "bin" +authors = [""] + +[dependencies] diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml b/tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml new file mode 100644 index 00000000000..6533d218b15 --- /dev/null +++ b/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] diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr b/tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr new file mode 100644 index 00000000000..7d847e085fe --- /dev/null +++ b/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))); +}