diff --git a/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr b/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr index 4582444c8f7..02057732f35 100644 --- a/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr +++ b/crates/nargo_cli/tests/execution_success/references_aliasing/src/main.nr @@ -1,10 +1,29 @@ -fn increment(mut r: &mut Field) { - *r = *r + 1; -} - fn main() { let mut x = 100; let mut xref = &mut x; increment(xref); assert(*xref == 101); + + regression_2445(); +} + +fn increment(mut r: &mut Field) { + *r = *r + 1; +} + +// If aliasing within arrays and constant folding within the mem2reg pass aren't +// handled, we'll fail to optimize out all the references in this function. +fn regression_2445() { + let mut var = 0; + let ref = &mut &mut var; + + let mut array = [ref, ref]; + + **array[0] = 1; + **array[1] = 2; + + assert(var == 2); + assert(**ref == 2); + assert(**array[0] == 2); + assert(**array[1] == 2); } diff --git a/crates/noirc_evaluator/src/ssa/ir/dfg.rs b/crates/noirc_evaluator/src/ssa/ir/dfg.rs index a9523afb2a4..ea00ad3bbc1 100644 --- a/crates/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/crates/noirc_evaluator/src/ssa/ir/dfg.rs @@ -171,7 +171,7 @@ impl DataFlowGraph { let id = self.make_instruction(instruction, ctrl_typevars); self.blocks[block].insert_instruction(id); self.locations.insert(id, call_stack); - InsertInstructionResult::Results(self.instruction_results(id)) + InsertInstructionResult::Results(id, self.instruction_results(id)) } } } @@ -478,7 +478,8 @@ impl std::ops::IndexMut for DataFlowGraph { // be a list of results or a single ValueId if the instruction was simplified // to an existing value. pub(crate) enum InsertInstructionResult<'dfg> { - Results(&'dfg [ValueId]), + /// Results is the standard case containing the instruction id and the results of that instruction. + Results(InstructionId, &'dfg [ValueId]), SimplifiedTo(ValueId), SimplifiedToMultiple(Vec), InstructionRemoved, @@ -490,7 +491,7 @@ impl<'dfg> InsertInstructionResult<'dfg> { match self { InsertInstructionResult::SimplifiedTo(value) => *value, InsertInstructionResult::SimplifiedToMultiple(values) => values[0], - InsertInstructionResult::Results(results) => results[0], + InsertInstructionResult::Results(_, results) => results[0], InsertInstructionResult::InstructionRemoved => { panic!("Instruction was removed, no results") } @@ -501,7 +502,7 @@ impl<'dfg> InsertInstructionResult<'dfg> { /// This is used for instructions returning multiple results like function calls. pub(crate) fn results(self) -> Cow<'dfg, [ValueId]> { match self { - InsertInstructionResult::Results(results) => Cow::Borrowed(results), + InsertInstructionResult::Results(_, results) => Cow::Borrowed(results), InsertInstructionResult::SimplifiedTo(result) => Cow::Owned(vec![result]), InsertInstructionResult::SimplifiedToMultiple(results) => Cow::Owned(results), InsertInstructionResult::InstructionRemoved => Cow::Owned(vec![]), @@ -513,7 +514,7 @@ impl<'dfg> InsertInstructionResult<'dfg> { match self { InsertInstructionResult::SimplifiedTo(_) => 1, InsertInstructionResult::SimplifiedToMultiple(results) => results.len(), - InsertInstructionResult::Results(results) => results.len(), + InsertInstructionResult::Results(_, results) => results.len(), InsertInstructionResult::InstructionRemoved => 0, } } diff --git a/crates/noirc_evaluator/src/ssa/ir/function_inserter.rs b/crates/noirc_evaluator/src/ssa/ir/function_inserter.rs index bc0084b6d4e..d9aee785016 100644 --- a/crates/noirc_evaluator/src/ssa/ir/function_inserter.rs +++ b/crates/noirc_evaluator/src/ssa/ir/function_inserter.rs @@ -30,7 +30,7 @@ impl<'f> FunctionInserter<'f> { pub(crate) fn resolve(&mut self, mut value: ValueId) -> ValueId { value = self.function.dfg.resolve(value); match self.values.get(&value) { - Some(value) => *value, + Some(value) => self.resolve(*value), None => match &self.function.dfg[value] { super::value::Value::Array { array, typ } => { let array = array.clone(); @@ -47,12 +47,22 @@ impl<'f> FunctionInserter<'f> { /// Insert a key, value pair if the key isn't already present in the map pub(crate) fn try_map_value(&mut self, key: ValueId, value: ValueId) { - self.values.entry(key).or_insert(value); + if key == value { + // This case is technically not needed since try_map_value isn't meant to change + // existing entries, but we should never have a value in the map referring to itself anyway. + self.values.remove(&key); + } else { + self.values.entry(key).or_insert(value); + } } /// Insert a key, value pair in the map pub(crate) fn map_value(&mut self, key: ValueId, value: ValueId) { - self.values.insert(key, value); + if key == value { + self.values.remove(&key); + } else { + self.values.insert(key, value); + } } pub(crate) fn map_instruction(&mut self, id: InstructionId) -> (Instruction, CallStack) { @@ -62,9 +72,27 @@ impl<'f> FunctionInserter<'f> { ) } - pub(crate) fn push_instruction(&mut self, id: InstructionId, block: BasicBlockId) { + /// Maps a terminator in place, replacing any ValueId in the terminator with the + /// resolved version of that value id from this FunctionInserter's internal value mapping. + pub(crate) fn map_terminator_in_place(&mut self, block: BasicBlockId) { + let mut terminator = self.function.dfg[block].take_terminator(); + terminator.mutate_values(|value| self.resolve(value)); + self.function.dfg[block].set_terminator(terminator); + } + + /// Push a new instruction to the given block and return its new InstructionId. + /// If the instruction was simplified out of the program, None is returned. + pub(crate) fn push_instruction( + &mut self, + id: InstructionId, + block: BasicBlockId, + ) -> Option { let (instruction, location) = self.map_instruction(id); - self.push_instruction_value(instruction, id, block, location); + + match self.push_instruction_value(instruction, id, block, location) { + InsertInstructionResult::Results(new_id, _) => Some(new_id), + _ => None, + } } pub(crate) fn push_instruction_value( @@ -110,7 +138,7 @@ impl<'f> FunctionInserter<'f> { values.insert(*old_result, *new_result); } } - InsertInstructionResult::Results(new_results) => { + InsertInstructionResult::Results(_, new_results) => { for (old_result, new_result) in old_results.iter().zip(*new_results) { values.insert(*old_result, *new_result); } diff --git a/crates/noirc_evaluator/src/ssa/ir/instruction.rs b/crates/noirc_evaluator/src/ssa/ir/instruction.rs index c0d2a6251f2..decf5140ad3 100644 --- a/crates/noirc_evaluator/src/ssa/ir/instruction.rs +++ b/crates/noirc_evaluator/src/ssa/ir/instruction.rs @@ -546,6 +546,26 @@ impl TerminatorInstruction { } } + /// Mutate each ValueId to a new ValueId using the given mapping function + pub(crate) fn mutate_values(&mut self, mut f: impl FnMut(ValueId) -> ValueId) { + use TerminatorInstruction::*; + match self { + JmpIf { condition, .. } => { + *condition = f(*condition); + } + Jmp { arguments, .. } => { + for argument in arguments { + *argument = f(*argument); + } + } + Return { return_values } => { + for return_value in return_values { + *return_value = f(*return_value); + } + } + } + } + /// Apply a function to each value pub(crate) fn for_each_value(&self, mut f: impl FnMut(ValueId) -> T) { use TerminatorInstruction::*; @@ -866,7 +886,7 @@ pub(crate) enum SimplifyResult { /// a function such as a tuple SimplifiedToMultiple(Vec), - /// Replace this function with an simpler but equivalent function. + /// Replace this function with an simpler but equivalent instruction. SimplifiedToInstruction(Instruction), /// Remove the instruction, it is unnecessary diff --git a/crates/noirc_evaluator/src/ssa/opt/constant_folding.rs b/crates/noirc_evaluator/src/ssa/opt/constant_folding.rs index d4ad64da519..4deec522f0d 100644 --- a/crates/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/crates/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -98,7 +98,7 @@ impl Context { ) { InsertInstructionResult::SimplifiedTo(new_result) => vec![new_result], InsertInstructionResult::SimplifiedToMultiple(new_results) => new_results, - InsertInstructionResult::Results(new_results) => new_results.to_vec(), + InsertInstructionResult::Results(_, new_results) => new_results.to_vec(), InsertInstructionResult::InstructionRemoved => vec![], }; assert_eq!(old_results.len(), new_results.len()); diff --git a/crates/noirc_evaluator/src/ssa/opt/inlining.rs b/crates/noirc_evaluator/src/ssa/opt/inlining.rs index 55424f8f32f..ba247631e61 100644 --- a/crates/noirc_evaluator/src/ssa/opt/inlining.rs +++ b/crates/noirc_evaluator/src/ssa/opt/inlining.rs @@ -392,7 +392,7 @@ impl<'function> PerFunctionContext<'function> { self.context.call_stack.pop_back(); } - let new_results = InsertInstructionResult::Results(&new_results); + let new_results = InsertInstructionResult::Results(call_id, &new_results); Self::insert_new_instruction_results(&mut self.values, old_results, new_results); } @@ -435,7 +435,7 @@ impl<'function> PerFunctionContext<'function> { values.insert(*old_result, new_result); } } - InsertInstructionResult::Results(new_results) => { + InsertInstructionResult::Results(_, new_results) => { for (old_result, new_result) in old_results.iter().zip(new_results) { values.insert(*old_result, *new_result); } diff --git a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs index 4d2531fa27c..84af1053ed7 100644 --- a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -69,6 +69,7 @@ use crate::ssa::{ cfg::ControlFlowGraph, dom::DominatorTree, function::Function, + function_inserter::FunctionInserter, instruction::{Instruction, InstructionId, TerminatorInstruction}, post_order::PostOrder, types::Type, @@ -83,20 +84,22 @@ impl Ssa { pub(crate) fn mem2reg(mut self) -> Ssa { for function in self.functions.values_mut() { let mut context = PerFunctionContext::new(function); - context.mem2reg(function); - context.remove_instructions(function); + context.mem2reg(); + context.remove_instructions(); } self } } -struct PerFunctionContext { +struct PerFunctionContext<'f> { cfg: ControlFlowGraph, post_order: PostOrder, dom_tree: DominatorTree, blocks: BTreeMap, + inserter: FunctionInserter<'f>, + /// Load and Store instructions that should be removed at the end of the pass. /// /// We avoid removing individual instructions as we go since removing elements @@ -143,8 +146,8 @@ enum ReferenceValue { Known(ValueId), } -impl PerFunctionContext { - fn new(function: &Function) -> Self { +impl<'f> PerFunctionContext<'f> { + fn new(function: &'f mut Function) -> Self { let cfg = ControlFlowGraph::with_function(function); let post_order = PostOrder::with_function(function); let dom_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); @@ -153,6 +156,7 @@ impl PerFunctionContext { cfg, post_order, dom_tree, + inserter: FunctionInserter::new(function), blocks: BTreeMap::new(), instructions_to_remove: BTreeSet::new(), } @@ -162,14 +166,14 @@ impl PerFunctionContext { /// /// This function is expected to be the same one that the internal cfg, post_order, and /// dom_tree were created from. - fn mem2reg(&mut self, function: &mut Function) { + fn mem2reg(&mut self) { // Iterate each block in reverse post order = forward order - let mut block_order = PostOrder::with_function(function).into_vec(); + let mut block_order = PostOrder::with_function(self.inserter.function).into_vec(); block_order.reverse(); for block in block_order { let references = self.find_starting_references(block); - self.analyze_block(function, block, references); + self.analyze_block(block, references); } } @@ -199,25 +203,20 @@ impl PerFunctionContext { /// This will remove any known loads in the block and track the value of references /// as they are stored to. When this function is finished, the value of each reference /// at the end of this block will be remembered in `self.blocks`. - fn analyze_block( - &mut self, - function: &mut Function, - block: BasicBlockId, - mut references: Block, - ) { - let instructions = function.dfg[block].instructions().to_vec(); + fn analyze_block(&mut self, block: BasicBlockId, mut references: Block) { + let instructions = self.inserter.function.dfg[block].take_instructions(); for instruction in instructions { - self.analyze_instruction(function, &mut references, instruction); + self.analyze_instruction(block, &mut references, instruction); } - self.handle_terminator(function, block, &mut references); + self.handle_terminator(block, &mut references); // If there's only 1 block in the function total, we can remove any remaining last stores // as well. We can't do this if there are multiple blocks since subsequent blocks may // reference these stores. if self.post_order.as_slice().len() == 1 { - self.remove_stores_that_do_not_alias_parameters(function, &references); + self.remove_stores_that_do_not_alias_parameters(&references); } self.blocks.insert(block, references); @@ -225,15 +224,10 @@ impl PerFunctionContext { /// Add all instructions in `last_stores` to `self.instructions_to_remove` which do not /// possibly alias any parameters of the given function. - fn remove_stores_that_do_not_alias_parameters( - &mut self, - function: &Function, - references: &Block, - ) { - let reference_parameters = function - .parameters() - .iter() - .filter(|param| function.dfg.value_is_reference(**param)) + fn remove_stores_that_do_not_alias_parameters(&mut self, references: &Block) { + let parameters = self.inserter.function.parameters().iter(); + let reference_parameters = parameters + .filter(|param| self.inserter.function.dfg.value_is_reference(**param)) .collect::>(); for (allocation, instruction) in &references.last_stores { @@ -252,32 +246,40 @@ impl PerFunctionContext { fn analyze_instruction( &mut self, - function: &mut Function, + block_id: BasicBlockId, references: &mut Block, - instruction: InstructionId, + mut instruction: InstructionId, ) { - match &function.dfg[instruction] { + // If the instruction was simplified and optimized out of the program we shouldn't analyze + // it. Analyzing it could make tracking aliases less accurate if it is e.g. an ArrayGet + // call that used to hold references but has since been optimized out to a known result. + if let Some(new_id) = self.inserter.push_instruction(instruction, block_id) { + instruction = new_id; + } else { + return; + } + + match &self.inserter.function.dfg[instruction] { Instruction::Load { address } => { - let address = function.dfg.resolve(*address); + let address = self.inserter.function.dfg.resolve(*address); - let result = function.dfg.instruction_results(instruction)[0]; - references.remember_dereference(function, address, result); + let result = self.inserter.function.dfg.instruction_results(instruction)[0]; + references.remember_dereference(self.inserter.function, address, result); // If the load is known, replace it with the known value and remove the load if let Some(value) = references.get_known_value(address) { - let result = function.dfg.instruction_results(instruction)[0]; - function.dfg.set_value_from_id(result, value); - + let result = self.inserter.function.dfg.instruction_results(instruction)[0]; + self.inserter.map_value(result, value); self.instructions_to_remove.insert(instruction); } else { - references.mark_value_used(address, function); + references.mark_value_used(address, self.inserter.function); } } Instruction::Store { address, value } => { - let address = function.dfg.resolve(*address); - let value = function.dfg.resolve(*value); + let address = self.inserter.function.dfg.resolve(*address); + let value = self.inserter.function.dfg.resolve(*value); - self.check_array_aliasing(function, references, value); + self.check_array_aliasing(references, value); // If there was another store to this instruction without any (unremoved) loads or // function calls in-between, we can remove the previous store. @@ -289,11 +291,11 @@ impl PerFunctionContext { references.last_stores.insert(address, instruction); } Instruction::Call { arguments, .. } => { - self.mark_all_unknown(arguments, function, references); + self.mark_all_unknown(arguments, references); } Instruction::Allocate => { // Register the new reference - let result = function.dfg.instruction_results(instruction)[0]; + let result = self.inserter.function.dfg.instruction_results(instruction)[0]; references.expressions.insert(result, Expression::Other(result)); let mut aliases = BTreeSet::new(); @@ -301,11 +303,11 @@ impl PerFunctionContext { references.aliases.insert(Expression::Other(result), aliases); } Instruction::ArrayGet { array, .. } => { - let result = function.dfg.instruction_results(instruction)[0]; - references.mark_value_used(*array, function); + let result = self.inserter.function.dfg.instruction_results(instruction)[0]; + references.mark_value_used(*array, self.inserter.function); - if function.dfg.value_is_reference(result) { - let array = function.dfg.resolve(*array); + if self.inserter.function.dfg.value_is_reference(result) { + let array = self.inserter.function.dfg.resolve(*array); let expression = Expression::ArrayElement(Box::new(Expression::Other(array))); if let Some(aliases) = references.aliases.get_mut(&expression) { @@ -314,17 +316,19 @@ impl PerFunctionContext { } } Instruction::ArraySet { array, value, .. } => { - references.mark_value_used(*array, function); + references.mark_value_used(*array, self.inserter.function); - if function.dfg.value_is_reference(*value) { - let result = function.dfg.instruction_results(instruction)[0]; - let array = function.dfg.resolve(*array); + if self.inserter.function.dfg.value_is_reference(*value) { + let result = self.inserter.function.dfg.instruction_results(instruction)[0]; + let array = self.inserter.function.dfg.resolve(*array); let expression = Expression::ArrayElement(Box::new(Expression::Other(array))); if let Some(aliases) = references.aliases.get_mut(&expression) { aliases.insert(result); - } else if let Some((elements, _)) = function.dfg.get_array_constant(array) { + } else if let Some((elements, _)) = + self.inserter.function.dfg.get_array_constant(array) + { // TODO: This should be a unification of each alias set // If any are empty, the whole should be as well. for reference in elements { @@ -339,8 +343,8 @@ impl PerFunctionContext { } } - fn check_array_aliasing(&self, function: &Function, references: &mut Block, array: ValueId) { - if let Some((elements, typ)) = function.dfg.get_array_constant(array) { + fn check_array_aliasing(&self, references: &mut Block, array: ValueId) { + if let Some((elements, typ)) = self.inserter.function.dfg.get_array_constant(array) { if Self::contains_references(&typ) { // TODO: Check if type directly holds references or holds arrays that hold references let expr = Expression::ArrayElement(Box::new(Expression::Other(array))); @@ -373,12 +377,12 @@ impl PerFunctionContext { } } - fn mark_all_unknown(&self, values: &[ValueId], function: &Function, references: &mut Block) { + fn mark_all_unknown(&self, values: &[ValueId], references: &mut Block) { for value in values { - if function.dfg.value_is_reference(*value) { - let value = function.dfg.resolve(*value); + if self.inserter.function.dfg.value_is_reference(*value) { + let value = self.inserter.function.dfg.resolve(*value); references.set_unknown(value); - references.mark_value_used(value, function); + references.mark_value_used(value, self.inserter.function); } } } @@ -386,31 +390,28 @@ impl PerFunctionContext { /// Remove any instructions in `self.instructions_to_remove` from the current function. /// This is expected to contain any loads which were replaced and any stores which are /// no longer needed. - fn remove_instructions(&self, function: &mut Function) { + fn remove_instructions(&mut self) { // The order we iterate blocks in is not important for block in self.post_order.as_slice() { - function.dfg[*block] + self.inserter.function.dfg[*block] .instructions_mut() .retain(|instruction| !self.instructions_to_remove.contains(instruction)); } } - fn handle_terminator( - &self, - function: &mut Function, - block: BasicBlockId, - references: &mut Block, - ) { - match function.dfg[block].unwrap_terminator() { + fn handle_terminator(&mut self, block: BasicBlockId, references: &mut Block) { + self.inserter.map_terminator_in_place(block); + + match self.inserter.function.dfg[block].unwrap_terminator() { TerminatorInstruction::JmpIf { .. } => (), // Nothing to do TerminatorInstruction::Jmp { destination, arguments, .. } => { - let destination_parameters = function.dfg[*destination].parameters(); + let destination_parameters = self.inserter.function.dfg[*destination].parameters(); assert_eq!(destination_parameters.len(), arguments.len()); // Add an alias for each reference parameter for (parameter, argument) in destination_parameters.iter().zip(arguments) { - if function.dfg.value_is_reference(*parameter) { - let argument = function.dfg.resolve(*argument); + if self.inserter.function.dfg.value_is_reference(*parameter) { + let argument = self.inserter.function.dfg.resolve(*argument); if let Some(expression) = references.expressions.get(&argument) { if let Some(aliases) = references.aliases.get_mut(expression) { @@ -425,7 +426,7 @@ impl PerFunctionContext { // Removing all `last_stores` for each returned reference is more important here // than setting them all to ReferenceValue::Unknown since no other block should // have a block with a Return terminator as a predecessor anyway. - self.mark_all_unknown(return_values, function, references); + self.mark_all_unknown(return_values, references); } } } @@ -719,12 +720,18 @@ mod tests { // Store is needed by the return value, and can't be removed assert_eq!(count_stores(block_id, &func.dfg), 1); + let instructions = func.dfg[block_id].instructions(); + assert_eq!(instructions.len(), 2); let ret_val_id = match func.dfg[block_id].terminator().unwrap() { - TerminatorInstruction::Return { return_values } => return_values.first().unwrap(), + TerminatorInstruction::Return { return_values } => *return_values.first().unwrap(), _ => unreachable!(), }; - assert_eq!(func.dfg[*ret_val_id], func.dfg[v0]); + + // Since the mem2reg pass simplifies as it goes, the id of the allocate instruction result + // is most likely no longer v0. We have to retrieve the new id here. + let alloca_id = func.dfg.instruction_results(instructions[0])[0]; + assert_eq!(ret_val_id, alloca_id); } fn count_stores(block: BasicBlockId, dfg: &DataFlowGraph) -> usize { @@ -786,12 +793,12 @@ mod tests { // Expected result: // acir fn main f0 { // b0(): - // v0 = allocate - // store Field 5 at v0 + // v7 = allocate + // store Field 5 at v7 // jmp b1(Field 5) // b1(v3: Field): - // store Field 6 at v0 - // return v3, Field 5, Field 6 // Optimized to constants 5 and 6 + // store Field 6 at v7 + // return v3, Field 5, Field 6 // } let ssa = ssa.mem2reg(); @@ -873,14 +880,13 @@ mod tests { // Expected result: // acir fn main f0 { // b0(): - // v0 = allocate - // store Field 0 at v0 - // v2 = allocate - // store v0 at v2 + // v9 = allocate + // store Field 0 at v9 + // v10 = allocate + // store v9 at v10 // jmp b1() // b1(): - // store Field 2 at v0 - // v8 = eq Field 1, Field 2 + // store Field 2 at v9 // return // } let ssa = ssa.mem2reg(); @@ -899,18 +905,7 @@ mod tests { let b1_instructions = main.dfg[b1].instructions(); - // The last instruction in b1 should be a binary operation - match &main.dfg[*b1_instructions.last().unwrap()] { - Instruction::Binary(binary) => { - let lhs = - main.dfg.get_numeric_constant(binary.lhs).expect("Expected constant value"); - let rhs = - main.dfg.get_numeric_constant(binary.rhs).expect("Expected constant value"); - - assert_eq!(lhs, rhs); - assert_eq!(lhs, FieldElement::from(2u128)); - } - _ => unreachable!(), - } + // We expect the last eq to be optimized out + assert_eq!(b1_instructions.len(), 1); } }