From 24cc19ed9f29792c7b056124b2adf87fc6c18e42 Mon Sep 17 00:00:00 2001 From: jfecher Date: Thu, 5 Dec 2024 10:24:56 -0600 Subject: [PATCH] fix: Optimize array ref counts to copy arrays much less often (#6685) Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com> --- .../brillig_ir/procedures/array_copy.rs | 2 + .../src/ssa/function_builder/mod.rs | 24 ++++-- .../src/ssa/ir/function_inserter.rs | 2 +- .../noirc_evaluator/src/ssa/ir/instruction.rs | 13 +-- compiler/noirc_evaluator/src/ssa/ir/types.rs | 9 ++ .../src/ssa/opt/constant_folding.rs | 82 +++++++++++++------ .../src/ssa/opt/loop_invariant.rs | 3 +- .../src/ssa/ssa_gen/context.rs | 56 +++++++++---- .../noirc_evaluator/src/ssa/ssa_gen/mod.rs | 62 ++++++++++---- .../src/monomorphization/ast.rs | 6 ++ .../array_dedup_regression/Nargo.toml | 6 ++ .../array_dedup_regression/Prover.toml | 1 + .../array_dedup_regression/src/main.nr | 21 +++++ .../reference_counts/src/main.nr | 56 +++++++++++-- 14 files changed, 264 insertions(+), 79 deletions(-) create mode 100644 test_programs/execution_success/array_dedup_regression/Nargo.toml create mode 100644 test_programs/execution_success/array_dedup_regression/Prover.toml create mode 100644 test_programs/execution_success/array_dedup_regression/src/main.nr diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/array_copy.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/array_copy.rs index 67f7cf2dc34..0a6e8824223 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/array_copy.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/procedures/array_copy.rs @@ -69,6 +69,8 @@ pub(super) fn compile_array_copy_procedure( 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); } }); } diff --git a/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs b/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs index 0479f8da0b7..0ae61404442 100644 --- a/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs +++ b/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 } } } diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs index a0c23ad70aa..6ebd2aa1105 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs +++ b/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); diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs index e34e0070d4b..76409f6a20a 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs +++ b/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) } } } diff --git a/compiler/noirc_evaluator/src/ssa/ir/types.rs b/compiler/noirc_evaluator/src/ssa/ir/types.rs index 16f4b8d2431..4e4f7e8aa62 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/types.rs +++ b/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> { match self { Type::Array(element_types, _) | Type::Slice(element_types) => element_types, diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index d1d84c305d2..93ca428c6d0 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/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,24 +332,17 @@ 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); @@ -346,7 +350,7 @@ impl<'brillig> Context<'brillig> { 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, - 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] diff --git a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs b/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs index 2e5c44cce76..290d2a33846 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs +++ b/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 diff --git a/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs b/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs index 8a09dba64c4..116e0de4ecd 100644 --- a/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs +++ b/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. @@ -159,7 +159,8 @@ impl<'a> FunctionContext<'a> { let parameter_value = Self::map_type(parameter_type, |typ| { let value = self.builder.add_parameter(typ); if mutable { - self.new_mutable_variable(value) + // This will wrap any `mut var: T` in a reference and increase the rc of an array if needed + self.new_mutable_variable(value, true) } else { value.into() } @@ -170,9 +171,17 @@ impl<'a> FunctionContext<'a> { /// Allocate a single slot of memory and store into it the given initial value of the variable. /// Always returns a Value::Mutable wrapping the allocate instruction. - pub(super) fn new_mutable_variable(&mut self, value_to_store: ValueId) -> Value { + pub(super) fn new_mutable_variable( + &mut self, + value_to_store: ValueId, + increment_array_rc: bool, + ) -> Value { let element_type = self.builder.current_function.dfg.type_of_value(value_to_store); - self.builder.increment_array_reference_count(value_to_store); + + if increment_array_rc { + self.builder.increment_array_reference_count(value_to_store); + } + let alloc = self.builder.insert_allocate(element_type); self.builder.insert_store(alloc, value_to_store); let typ = self.builder.type_of_value(value_to_store); @@ -904,35 +913,52 @@ impl<'a> FunctionContext<'a> { } } - /// Increments the reference count of all parameters. Returns the entry block of the function. + /// Increments the reference count of mutable reference array parameters. + /// Any mutable-value (`mut a: [T; N]` versus `a: &mut [T; N]`) are already incremented + /// by `FunctionBuilder::add_parameter_to_scope`. + /// 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 { let entry = self.builder.current_function.entry_block(); let parameters = self.builder.current_function.dfg.block_parameters(entry).to_vec(); + 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.get_contained_array().clone()) + && 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).to_vec(); - - dropped_parameters.retain(|parameter| !terminator_args.contains(parameter)); + pub(crate) fn end_scope( + &mut self, + mut incremented_params: HashSet, + 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); } diff --git a/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs b/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs index d28236bd360..2fe0a38af00 100644 --- a/compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs +++ b/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(()) @@ -195,8 +195,7 @@ impl<'a> FunctionContext<'a> { fn codegen_literal(&mut self, literal: &ast::Literal) -> Result { match literal { ast::Literal::Array(array) => { - let elements = - try_vecmap(&array.contents, |element| self.codegen_expression(element))?; + let elements = self.codegen_array_elements(&array.contents)?; let typ = Self::convert_type(&array.typ).flatten(); Ok(match array.typ { @@ -207,8 +206,7 @@ impl<'a> FunctionContext<'a> { }) } ast::Literal::Slice(array) => { - let elements = - try_vecmap(&array.contents, |element| self.codegen_expression(element))?; + let elements = self.codegen_array_elements(&array.contents)?; let typ = Self::convert_type(&array.typ).flatten(); Ok(match array.typ { @@ -245,18 +243,33 @@ impl<'a> FunctionContext<'a> { } } + fn codegen_array_elements( + &mut self, + elements: &[Expression], + ) -> Result, RuntimeError> { + try_vecmap(elements, |element| { + let value = self.codegen_expression(element)?; + Ok((value, element.is_array_or_slice_literal())) + }) + } + fn codegen_string(&mut self, string: &str) -> Values { let elements = vecmap(string.as_bytes(), |byte| { - self.builder.numeric_constant(*byte as u128, Type::unsigned(8)).into() + let char = self.builder.numeric_constant(*byte as u128, Type::unsigned(8)); + (char.into(), false) }); let typ = Self::convert_non_tuple_type(&ast::Type::String(elements.len() as u32)); self.codegen_array(elements, typ) } // Codegen an array but make sure that we do not have a nested slice + /// + /// The bool aspect of each array element indicates whether the element is an array constant + /// or not. If it is, we avoid incrementing the reference count because we consider the + /// constant to be moved into this larger array constant. fn codegen_array_checked( &mut self, - elements: Vec, + elements: Vec<(Values, bool)>, typ: Type, ) -> Result { if typ.is_nested_slice() { @@ -273,11 +286,15 @@ impl<'a> FunctionContext<'a> { /// stored next to the other fields in memory. So an array such as [(1, 2), (3, 4)] is /// stored the same as the array [1, 2, 3, 4]. /// + /// The bool aspect of each array element indicates whether the element is an array constant + /// or not. If it is, we avoid incrementing the reference count because we consider the + /// constant to be moved into this larger array constant. + /// /// The value returned from this function is always that of the allocate instruction. - fn codegen_array(&mut self, elements: Vec, typ: Type) -> Values { + fn codegen_array(&mut self, elements: Vec<(Values, bool)>, typ: Type) -> Values { let mut array = im::Vector::new(); - for element in elements { + for (element, is_array_constant) in elements { element.for_each(|element| { let element = element.eval(self); @@ -286,7 +303,10 @@ impl<'a> FunctionContext<'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); + if !is_array_constant { + self.builder.increment_array_reference_count(element); + } + array.push_back(element); }); } @@ -662,14 +682,22 @@ impl<'a> FunctionContext<'a> { fn codegen_let(&mut self, let_expr: &ast::Let) -> Result { let mut values = self.codegen_expression(&let_expr.expression)?; + // Don't mutate the reference count if we're assigning an array literal to a Let: + // `let mut foo = [1, 2, 3];` + // we consider the array to be moved, so we should have an initial rc of just 1. + let should_inc_rc = !let_expr.expression.is_array_or_slice_literal(); + values = values.map(|value| { let value = value.eval(self); Tree::Leaf(if let_expr.mutable { - self.new_mutable_variable(value) + self.new_mutable_variable(value, should_inc_rc) } else { - // `new_mutable_variable` already increments rcs internally - self.builder.increment_array_reference_count(value); + // `new_mutable_variable` increments rcs internally so we have to + // handle it separately for the immutable case + if should_inc_rc { + self.builder.increment_array_reference_count(value); + } value::Value::Normal(value) }) }); @@ -728,10 +756,14 @@ impl<'a> FunctionContext<'a> { fn codegen_assign(&mut self, assign: &ast::Assign) -> Result { let lhs = self.extract_current_value(&assign.lvalue)?; let rhs = self.codegen_expression(&assign.expression)?; + let should_inc_rc = !assign.expression.is_array_or_slice_literal(); rhs.clone().for_each(|value| { let value = value.eval(self); - self.builder.increment_array_reference_count(value); + + if should_inc_rc { + self.builder.increment_array_reference_count(value); + } }); self.assign_new_value(lhs, rhs); diff --git a/compiler/noirc_frontend/src/monomorphization/ast.rs b/compiler/noirc_frontend/src/monomorphization/ast.rs index 8f6817dc15d..5d9b66f4f96 100644 --- a/compiler/noirc_frontend/src/monomorphization/ast.rs +++ b/compiler/noirc_frontend/src/monomorphization/ast.rs @@ -48,6 +48,12 @@ pub enum Expression { Continue, } +impl Expression { + pub fn is_array_or_slice_literal(&self) -> bool { + matches!(self, Expression::Literal(Literal::Array(_) | Literal::Slice(_))) + } +} + /// A definition is either a local (variable), function, or is a built-in /// function that will be generated or referenced by the compiler later. #[derive(Debug, Clone, PartialEq, Eq, Hash)] diff --git a/test_programs/execution_success/array_dedup_regression/Nargo.toml b/test_programs/execution_success/array_dedup_regression/Nargo.toml new file mode 100644 index 00000000000..16a708743ed --- /dev/null +++ b/test_programs/execution_success/array_dedup_regression/Nargo.toml @@ -0,0 +1,6 @@ +[package] +name = "array_dedup_regression" +type = "bin" +authors = [""] + +[dependencies] diff --git a/test_programs/execution_success/array_dedup_regression/Prover.toml b/test_programs/execution_success/array_dedup_regression/Prover.toml new file mode 100644 index 00000000000..3aea0c58ce5 --- /dev/null +++ b/test_programs/execution_success/array_dedup_regression/Prover.toml @@ -0,0 +1 @@ +x = 0 diff --git a/test_programs/execution_success/array_dedup_regression/src/main.nr b/test_programs/execution_success/array_dedup_regression/src/main.nr new file mode 100644 index 00000000000..5506d55b9e7 --- /dev/null +++ b/test_programs/execution_success/array_dedup_regression/src/main.nr @@ -0,0 +1,21 @@ +unconstrained fn main(x: u32) { + let a1 = [1, 2, 3, 4, 5]; + + for i in 0..5 { + let mut a2 = [1, 2, 3, 4, 5]; + a2[x + i] = 128; + println(a2); + + if i != 0 { + assert(a2[x + i - 1] != 128); + } + } + + // Can't use `== [1, 2, 3, 4, 5]` here, that make_array may get + // deduplicated to equal a1 in the bugged version + assert_eq(a1[0], 1); + assert_eq(a1[1], 2); + assert_eq(a1[2], 3); + assert_eq(a1[3], 4); + assert_eq(a1[4], 5); +} diff --git a/test_programs/execution_success/reference_counts/src/main.nr b/test_programs/execution_success/reference_counts/src/main.nr index 86b5c812472..8de4d0f2508 100644 --- a/test_programs/execution_success/reference_counts/src/main.nr +++ b/test_programs/execution_success/reference_counts/src/main.nr @@ -1,10 +1,19 @@ +use std::mem::array_refcount; + fn main() { let mut array = [0, 1, 2]; - assert_refcount(array, 2); + assert_refcount(array, 1); + + borrow(array, array_refcount(array)); + borrow_mut(&mut array, array_refcount(array)); + copy_mut(array, array_refcount(array)); - borrow(array, std::mem::array_refcount(array)); - borrow_mut(&mut array, std::mem::array_refcount(array)); - copy_mut(array, std::mem::array_refcount(array)); + borrow_mut_two(&mut array, &mut array, array_refcount(array)); + + let mut u32_array = [0, 1, 2]; + let rc1 = array_refcount(array); + let rc2 = array_refcount(u32_array); + borrow_mut_two_separate(&mut array, &mut u32_array, rc1, rc2); } fn borrow(array: [Field; 3], rc_before_call: u32) { @@ -13,19 +22,48 @@ fn borrow(array: [Field; 3], rc_before_call: u32) { } fn borrow_mut(array: &mut [Field; 3], rc_before_call: u32) { - assert_refcount(*array, rc_before_call + 1); - array[0] = 5; + // Optimization: inc_rc isn't needed since there is only one array (`array`) + // of the same type that `array` can be modified through + assert_refcount(*array, rc_before_call + 0); + array[0] = 3; println(array[0]); } fn copy_mut(mut array: [Field; 3], rc_before_call: u32) { assert_refcount(array, rc_before_call + 1); - array[0] = 6; + array[0] = 4; println(array[0]); } -fn assert_refcount(array: [Field; 3], expected: u32) { - let count = std::mem::array_refcount(array); +/// Borrow the same array mutably through both parameters, inc_rc is necessary here, although +/// only one is needed to bring the rc from 1 to 2. +fn borrow_mut_two(array1: &mut [Field; 3], array2: &mut [Field; 3], rc_before_call: u32) { + assert_refcount(*array1, rc_before_call + 1); + assert_refcount(*array2, rc_before_call + 1); + array1[0] = 5; + array2[0] = 6; + println(array1[0]); // array1 & 2 alias, so this should also print 6 + println(array2[0]); +} + +/// Borrow a different array: we should be able to reason that these types cannot be mutably +/// aliased since they're different types so we don't need any inc_rc instructions. +fn borrow_mut_two_separate( + array1: &mut [Field; 3], + array2: &mut [u32; 3], + rc_before_call1: u32, + rc_before_call2: u32, +) { + assert_refcount(*array1, rc_before_call1 + 0); + assert_refcount(*array2, rc_before_call2 + 0); + array1[0] = 7; + array2[0] = 8; + println(array1[0]); + println(array2[0]); +} + +fn assert_refcount(array: [T; 3], expected: u32) { + let count = array_refcount(array); // All refcounts are zero when running this as a constrained program if std::runtime::is_unconstrained() {