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 4e5a8cde0c2..e3d8b33837c 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -622,7 +622,7 @@ impl<'block> BrilligBlock<'block> { destination_variable, ); } - Instruction::ArraySet { array, index, value, .. } => { + Instruction::ArraySet { array, index, value, mutable: _ } => { let source_variable = self.convert_ssa_value(*array, dfg); let index_register = self.convert_ssa_single_addr_value(*index, dfg); let value_variable = self.convert_ssa_value(*value, dfg); diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 870b5e602f1..f1c817d3866 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs @@ -530,7 +530,10 @@ impl<'dfg> InsertInstructionResult<'dfg> { match self { InsertInstructionResult::SimplifiedTo(value) => *value, InsertInstructionResult::SimplifiedToMultiple(values) => values[0], - InsertInstructionResult::Results(_, results) => results[0], + InsertInstructionResult::Results(_, results) => { + assert_eq!(results.len(), 1); + results[0] + } InsertInstructionResult::InstructionRemoved => { panic!("Instruction was removed, no results") } diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs index 1187ea8cb07..78b68f6d61d 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs @@ -354,7 +354,9 @@ 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 unknown = HashMap::default(); + let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes, &unknown, None); + let new_slice = value_merger.merge_values( len_not_equals_capacity, len_equals_capacity, diff --git a/compiler/noirc_evaluator/src/ssa/ir/printer.rs b/compiler/noirc_evaluator/src/ssa/ir/printer.rs index be92984441e..ef5ba39d0e4 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/printer.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/printer.rs @@ -181,7 +181,7 @@ fn display_instruction_inner( let index = show(*index); let value = show(*value); let mutable = if *mutable { " mut" } else { "" }; - writeln!(f, "array_set{mutable} {array}, index {index}, value {value}",) + writeln!(f, "array_set{mutable} {array}, index {index}, value {value}") } Instruction::IncrementRc { value } => { writeln!(f, "inc_rc {}", show(*value)) 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 ac869caf995..47953bedb72 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,20 +1,25 @@ use acvm::FieldElement; -use fxhash::FxHashMap as HashMap; +use fxhash::{FxHashMap as HashMap, FxHashSet}; use crate::ssa::ir::{ basic_block::BasicBlockId, - dfg::{CallStack, DataFlowGraph}, + dfg::{CallStack, DataFlowGraph, InsertInstructionResult}, instruction::{BinaryOp, Instruction}, types::Type, - value::ValueId, + value::{Value, ValueId}, }; pub(crate) struct ValueMerger<'a> { dfg: &'a mut DataFlowGraph, block: BasicBlockId, + + current_condition: Option, + // Maps SSA array values with a slice type to their size. // This must be computed before merging values. slice_sizes: &'a mut HashMap, + + array_set_conditionals: &'a HashMap, } impl<'a> ValueMerger<'a> { @@ -22,8 +27,10 @@ impl<'a> ValueMerger<'a> { dfg: &'a mut DataFlowGraph, block: BasicBlockId, slice_sizes: &'a mut HashMap, + array_set_conditionals: &'a HashMap, + current_condition: Option, ) -> Self { - ValueMerger { dfg, block, slice_sizes } + ValueMerger { dfg, block, slice_sizes, array_set_conditionals, current_condition } } /// Merge two values a and b from separate basic blocks to a single value. @@ -131,6 +138,18 @@ impl<'a> ValueMerger<'a> { _ => panic!("Expected array type"), }; + let actual_length = len * element_types.len(); + + if let Some(result) = self.try_merge_only_changed_indices( + then_condition, + else_condition, + then_value, + else_value, + actual_length, + ) { + return result; + } + 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(); @@ -266,4 +285,94 @@ impl<'a> ValueMerger<'a> { } } } + + fn try_merge_only_changed_indices( + &mut self, + then_condition: ValueId, + else_condition: ValueId, + then_value: ValueId, + else_value: ValueId, + array_length: usize, + ) -> Option { + let mut seen_then = FxHashSet::default(); + let mut seen_else = FxHashSet::default(); + let mut found = false; + + let mut current_then = then_value; + let mut current_else = else_value; + + let mut changed_indices = FxHashSet::default(); + + // Arbitrarily limit this to looking at at most 10 past ArraySet operations. + // If there are more than that, we assume 2 completely separate arrays are being merged. + for _ in 0..10 { + seen_then.insert(current_then); + seen_else.insert(current_else); + + if seen_then.contains(¤t_else) || seen_else.contains(¤t_then) { + found = true; + break; + } + + current_then = self.find_previous_array_set(current_then, &mut changed_indices)?; + current_else = self.find_previous_array_set(current_else, &mut changed_indices)?; + } + + if !found || changed_indices.len() >= array_length { + return None; + } + + let mut array = then_value; + + for (index, element_type, condition) in changed_indices { + let typevars = Some(vec![element_type.clone()]); + + let instruction = Instruction::EnableSideEffects { condition }; + self.insert_instruction(instruction); + + 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 instruction = Instruction::ArraySet { array, index, value, mutable: false }; + array = self.insert_instruction(instruction).first(); + } + + let instruction = Instruction::EnableSideEffects { condition: self.current_condition? }; + self.insert_instruction(instruction); + + Some(array) + } + + fn insert_instruction(&mut self, instruction: Instruction) -> InsertInstructionResult { + self.dfg.insert_instruction_and_results(instruction, self.block, None, CallStack::new()) + } + + fn find_previous_array_set( + &self, + value: ValueId, + changed_indices: &mut FxHashSet<(ValueId, Type, ValueId)>, + ) -> Option { + match &self.dfg[value] { + Value::Instruction { instruction, .. } => match &self.dfg[*instruction] { + Instruction::ArraySet { array, index, value, .. } => { + let condition = *self.array_set_conditionals.get(index)?; + let element_type = self.dfg.type_of_value(*value); + changed_indices.insert((*index, element_type, condition)); + Some(*array) + } + _ => Some(value), + }, + _ => Some(value), + } + } } diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs index f21c2a1b519..86e90559639 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs @@ -42,12 +42,20 @@ impl Ssa { #[derive(Default)] struct Context { slice_sizes: HashMap, + + // Maps array_set result -> element that was overwritten by that instruction. + // Used to undo array_sets while merging values + prev_array_set_elem_values: HashMap, + + // Maps array_set result -> enable_side_effects_if value which was active during it. + array_set_conditionals: HashMap, } impl Context { fn remove_if_else(&mut self, function: &mut Function) { let block = function.entry_block(); let instructions = function.dfg[block].take_instructions(); + let mut current_conditional = function.dfg.make_constant(FieldElement::one(), Type::bool()); for instruction in instructions { match &function.dfg[instruction] { @@ -60,8 +68,14 @@ impl Context { let typ = function.dfg.type_of_value(then_value); assert!(!matches!(typ, Type::Numeric(_))); - let mut value_merger = - ValueMerger::new(&mut function.dfg, block, &mut self.slice_sizes); + let mut value_merger = ValueMerger::new( + &mut function.dfg, + block, + &mut self.slice_sizes, + &self.array_set_conditionals, + Some(current_conditional), + ); + let value = value_merger.merge_values( then_condition, else_condition, @@ -98,10 +112,15 @@ impl Context { let results = function.dfg.instruction_results(instruction); let result = if results.len() == 2 { results[1] } else { results[0] }; + self.array_set_conditionals.insert(*array, current_conditional); let old_capacity = self.get_or_find_capacity(&function.dfg, *array); self.slice_sizes.insert(result, old_capacity); function.dfg[block].instructions_mut().push(instruction); } + Instruction::EnableSideEffects { condition } => { + current_conditional = *condition; + function.dfg[block].instructions_mut().push(instruction); + } _ => { function.dfg[block].instructions_mut().push(instruction); }