diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs index 1187ea8cb07..d7b98631d47 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs @@ -354,7 +354,8 @@ fn simplify_slice_push_back( slice_sizes.insert(set_last_slice_value, slice_size / element_size); slice_sizes.insert(new_slice, slice_size / element_size); - let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes); + let modified_indices = HashMap::default(); + let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes, &modified_indices); let new_slice = value_merger.merge_values( len_not_equals_capacity, len_equals_capacity, diff --git a/compiler/noirc_evaluator/src/ssa/ir/types.rs b/compiler/noirc_evaluator/src/ssa/ir/types.rs index 48036580d29..d789225cead 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/types.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/types.rs @@ -130,6 +130,14 @@ impl Type { } } + /// Returns the length of this array type - if it is one + pub(crate) fn array_length(&self) -> Option { + match self { + Type::Array(_, length) => Some(*length), + _ => None, + } + } + /// Returns the flattened size of a Type pub(crate) fn flattened_size(&self) -> usize { match self { diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index e731a7952a6..4c771ed8431 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -131,8 +131,8 @@ //! v11 = mul v4, Field 12 //! v12 = add v10, v11 //! store v12 at v5 (new store) -use fxhash::FxHashMap as HashMap; -use std::collections::{BTreeMap, HashSet}; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; +use std::collections::BTreeMap; use acvm::FieldElement; use iter_extended::vecmap; @@ -142,7 +142,7 @@ use crate::ssa::{ basic_block::BasicBlockId, cfg::ControlFlowGraph, dfg::{CallStack, InsertInstructionResult}, - function::Function, + function::{Function, RuntimeType}, function_inserter::FunctionInserter, instruction::{BinaryOp, Instruction, InstructionId, Intrinsic, TerminatorInstruction}, types::Type, @@ -213,6 +213,8 @@ struct Context<'f> { arguments_stack: Vec>, } +pub(crate) type ModifiedArrayIndices = HashMap)>; + #[derive(Clone)] pub(crate) struct Store { old_value: ValueId, @@ -231,6 +233,9 @@ struct ConditionalBranch { store_values: HashMap, // The allocations accumulated when processing the branch local_allocations: HashSet, + + // Each array index that was stored to an array of a given type within this block. + array_indexes_set: ModifiedArrayIndices, } struct ConditionalContext { @@ -248,7 +253,7 @@ fn flatten_function_cfg(function: &mut Function) { // This pass may run forever on a brillig function. // Analyze will check if the predecessors have been processed and push the block to the back of // the queue. This loops forever if there are still any loops present in the program. - if let crate::ssa::ir::function::RuntimeType::Brillig = function.runtime() { + if let RuntimeType::Brillig = function.runtime() { return; } let cfg = ControlFlowGraph::with_function(function); @@ -258,7 +263,7 @@ fn flatten_function_cfg(function: &mut Function) { inserter: FunctionInserter::new(function), cfg, store_values: HashMap::default(), - local_allocations: HashSet::new(), + local_allocations: HashSet::default(), branch_ends, slice_sizes: HashMap::default(), condition_stack: Vec::new(), @@ -400,12 +405,14 @@ impl<'f> Context<'f> { let one = FieldElement::one(); let old_stores = std::mem::take(&mut self.store_values); let old_allocations = std::mem::take(&mut self.local_allocations); + let branch = ConditionalBranch { old_condition, condition: self.link_condition(then_condition), store_values: old_stores, local_allocations: old_allocations, last_block: *then_destination, + array_indexes_set: HashMap::default(), }; let cond_context = ConditionalContext { condition: then_condition, @@ -449,6 +456,7 @@ impl<'f> Context<'f> { condition: else_condition, store_values: old_stores, local_allocations: old_allocations, + array_indexes_set: HashMap::default(), last_block: *block, }; let old_condition = else_branch.old_condition; @@ -521,7 +529,7 @@ impl<'f> Context<'f> { fn inline_branch_end( &mut self, destination: BasicBlockId, - cond_context: ConditionalContext, + mut cond_context: ConditionalContext, ) -> BasicBlockId { assert_eq!(self.cfg.predecessors(destination).len(), 2); let last_then = cond_context.then_branch.last_block; @@ -550,8 +558,14 @@ impl<'f> Context<'f> { capacity_tracker.compute_slice_capacity(*else_arg, &mut self.slice_sizes); } - let mut value_merger = - ValueMerger::new(&mut self.inserter.function.dfg, block, &mut self.slice_sizes); + let array_indices_modified = self.pop_array_indices_set(&mut cond_context); + + let mut value_merger = ValueMerger::new( + &mut self.inserter.function.dfg, + block, + &mut self.slice_sizes, + &array_indices_modified, + ); // Cannot include this in the previous vecmap since it requires exclusive access to self let args = vecmap(args, |(then_arg, else_arg)| { @@ -563,7 +577,11 @@ impl<'f> Context<'f> { ) }); - self.merge_stores(cond_context.then_branch, cond_context.else_branch); + self.merge_stores( + cond_context.then_branch, + cond_context.else_branch, + &array_indices_modified, + ); self.arguments_stack.pop(); self.arguments_stack.pop(); self.arguments_stack.push(args); @@ -625,6 +643,7 @@ impl<'f> Context<'f> { &mut self, then_branch: ConditionalBranch, else_branch: Option, + modified_array_indices: &ModifiedArrayIndices, ) { // Address -> (then_value, else_value, value_before_the_if) let mut new_map = BTreeMap::new(); @@ -660,8 +679,13 @@ impl<'f> Context<'f> { }; let block = self.inserter.function.entry_block(); - let mut value_merger = - ValueMerger::new(&mut self.inserter.function.dfg, block, &mut self.slice_sizes); + let mut value_merger = ValueMerger::new( + &mut self.inserter.function.dfg, + block, + &mut self.slice_sizes, + modified_array_indices, + ); + // Merging must occur in a separate loop as we cannot borrow `self` as mutable while `value_merger` does let mut new_values = HashMap::default(); for (address, (then_case, else_case, _)) in &new_map { @@ -817,6 +841,32 @@ impl<'f> Context<'f> { _ => Instruction::Call { func, arguments }, }, + Instruction::ArraySet { array, index, value, mutable } => { + if let Some(condition) = self.condition_stack.last_mut() { + let branch = + condition.else_branch.as_mut().unwrap_or(&mut condition.then_branch); + + let typ = self.inserter.function.dfg.type_of_value(array); + let length = typ.array_length(); + let (is_constant, indices) = branch.array_indexes_set.entry(typ).or_default(); + + let constant = self.inserter.function.dfg.get_numeric_constant(index); + *is_constant = *is_constant || constant.is_none(); + + if let Some(constant) = constant.and_then(|value| value.try_to_u64()) { + if let Some(length) = length { + if constant >= length as u64 { + // index is out of bounds, don't add it to `indices` + return Instruction::ArraySet { array, index, value, mutable }; + } + } + } + + let element_type = self.inserter.function.dfg.type_of_value(value); + indices.insert((index, element_type)); + } + Instruction::ArraySet { array, index, value, mutable } + } other => other, } } else { @@ -831,6 +881,34 @@ impl<'f> Context<'f> { self.insert_instruction_with_typevars(Instruction::Store { address, value }, None); } } + + fn pop_array_indices_set( + &mut self, + cond_context: &mut ConditionalContext, + ) -> ModifiedArrayIndices { + let mut modified = std::mem::take(&mut cond_context.then_branch.array_indexes_set); + + if let Some(else_branch) = cond_context.else_branch.as_mut() { + let else_indices = std::mem::take(&mut else_branch.array_indexes_set); + modified.extend(else_indices); + } + + // Remember all the indices changed in the now finished if branches as also + // being changed for any conditions that encompass the whole if. E.g. + // if foo { + // if a { array[1] = 2; } else { array[3] = 4; } + // } + // This check will ensure the outer if sees the array modifications as well. + if let Some(current_condition) = self.condition_stack.last_mut() { + let branch = current_condition + .else_branch + .as_mut() + .unwrap_or(&mut current_condition.then_branch); + branch.array_indexes_set.extend(modified.clone()); + } + + modified + } } #[cfg(test)] @@ -841,7 +919,7 @@ mod test { function_builder::FunctionBuilder, ir::{ dfg::DataFlowGraph, - function::{Function, InlineType, RuntimeType}, + function::Function, instruction::{BinaryOp, Instruction, Intrinsic, TerminatorInstruction}, map::Id, types::Type, diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/value_merger.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/value_merger.rs index 0a351148fa3..d518ebf0490 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/value_merger.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/value_merger.rs @@ -1,5 +1,5 @@ use acvm::FieldElement; -use fxhash::FxHashMap as HashMap; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use crate::ssa::ir::{ basic_block::BasicBlockId, @@ -9,12 +9,17 @@ use crate::ssa::ir::{ value::ValueId, }; +use super::ModifiedArrayIndices; + pub(crate) struct ValueMerger<'a> { dfg: &'a mut DataFlowGraph, block: BasicBlockId, + // Maps SSA array values with a slice type to their size. // This must be computed before merging values. slice_sizes: &'a mut HashMap, + + modified_array_indices: &'a ModifiedArrayIndices, } impl<'a> ValueMerger<'a> { @@ -22,8 +27,9 @@ impl<'a> ValueMerger<'a> { dfg: &'a mut DataFlowGraph, block: BasicBlockId, slice_sizes: &'a mut HashMap, + modified_array_indices: &'a ModifiedArrayIndices, ) -> Self { - ValueMerger { dfg, block, slice_sizes } + ValueMerger { dfg, block, slice_sizes, modified_array_indices } } /// Merge two values a and b from separate basic blocks to a single value. @@ -76,41 +82,27 @@ impl<'a> ValueMerger<'a> { let else_call_stack = self.dfg.get_value_call_stack(else_value); let call_stack = if then_call_stack.is_empty() { else_call_stack } else { then_call_stack }; + let dfg = &mut self.dfg; // We must cast the bool conditions to the actual numeric type used by each value. - let then_condition = self - .dfg - .insert_instruction_and_results( - Instruction::Cast(then_condition, then_type), - self.block, - None, - call_stack.clone(), - ) - .first(); - let else_condition = self - .dfg - .insert_instruction_and_results( - Instruction::Cast(else_condition, else_type), - self.block, - None, - call_stack.clone(), - ) - .first(); + let cast = Instruction::Cast(then_condition, then_type); + let then_condition = + dfg.insert_instruction_and_results(cast, self.block, None, call_stack.clone()).first(); + + let cast = Instruction::Cast(else_condition, else_type); + let else_condition = + dfg.insert_instruction_and_results(cast, self.block, None, call_stack.clone()).first(); let mul = Instruction::binary(BinaryOp::Mul, then_condition, then_value); - let then_value = self - .dfg - .insert_instruction_and_results(mul, self.block, None, call_stack.clone()) - .first(); + let then_value = + dfg.insert_instruction_and_results(mul, self.block, None, call_stack.clone()).first(); let mul = Instruction::binary(BinaryOp::Mul, else_condition, else_value); - let else_value = self - .dfg - .insert_instruction_and_results(mul, self.block, None, call_stack.clone()) - .first(); + let else_value = + dfg.insert_instruction_and_results(mul, self.block, None, call_stack.clone()).first(); let add = Instruction::binary(BinaryOp::Add, then_value, else_value); - self.dfg.insert_instruction_and_results(add, self.block, None, call_stack).first() + dfg.insert_instruction_and_results(add, self.block, None, call_stack).first() } /// Given an if expression that returns an array: `if c { array1 } else { array2 }`, @@ -131,6 +123,18 @@ impl<'a> ValueMerger<'a> { _ => panic!("Expected array type"), }; + if let Some((invalid, modified_indices)) = self.modified_array_indices.get(&typ) { + if !*invalid && modified_indices.len() < len { + return self.merge_only_modified_indices( + then_condition, + else_condition, + then_value, + else_value, + modified_indices, + ); + } + } + for i in 0..len { for (element_index, element_type) in element_types.iter().enumerate() { let index = ((i * element_types.len() + element_index) as u128).into(); @@ -260,4 +264,51 @@ impl<'a> ValueMerger<'a> { } } } + + fn merge_only_modified_indices( + &mut self, + then_condition: ValueId, + else_condition: ValueId, + then_value: ValueId, + else_value: ValueId, + modified_indices: &HashSet<(ValueId, Type)>, + ) -> ValueId { + // We're going to merge all potentially modified indices, so the choice + // to start with the then or else value is arbitrary. + let mut new_array = then_value; + + eprintln!("array set opt triggered! Modified indices are:"); + + for (index, element_type) in modified_indices { + let index = *index; + let typevars = Some(vec![element_type.clone()]); + + if let Some(c) = self.dfg.get_numeric_constant(index) { + let i = c.try_to_u64().unwrap(); + eprintln!(" index {}", i); + } + + let mut get_element = |array, typevars| { + let get = Instruction::ArrayGet { array, index }; + self.dfg + .insert_instruction_and_results(get, self.block, typevars, CallStack::new()) + .first() + }; + + let then_element = get_element(then_value, typevars.clone()); + let else_element = get_element(else_value, typevars); + + let value = + self.merge_values(then_condition, else_condition, then_element, else_element); + + let array_set = + Instruction::ArraySet { array: new_array, index, value, mutable: false }; + new_array = self + .dfg + .insert_instruction_and_results(array_set, self.block, None, CallStack::new()) + .first(); + } + + new_array + } } diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs index f1a38585bd6..7b87142d824 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -414,7 +414,6 @@ mod tests { ir::{ basic_block::BasicBlockId, dfg::DataFlowGraph, - function::{InlineType, RuntimeType}, instruction::{BinaryOp, Instruction, Intrinsic, TerminatorInstruction}, map::Id, types::Type,