From f1ea6b4c62abc8e761676b60e6a07c446277a541 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Thu, 15 Jun 2023 14:05:44 -0500 Subject: [PATCH] Fix #1688 --- .../src/ssa_refactor/acir_gen/mod.rs | 2 +- crates/noirc_evaluator/src/ssa_refactor/ir.rs | 1 + .../src/ssa_refactor/ir/function_inserter.rs | 129 +++++++++++++++++ .../src/ssa_refactor/opt/flatten_cfg.rs | 133 +++++++----------- .../src/ssa_refactor/opt/unrolling.rs | 110 ++++----------- 5 files changed, 204 insertions(+), 171 deletions(-) create mode 100644 crates/noirc_evaluator/src/ssa_refactor/ir/function_inserter.rs diff --git a/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs b/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs index 3341991325..68d80a36de 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs @@ -346,7 +346,7 @@ impl Context { Value::Intrinsic(..) => todo!(), Value::Function(..) => unreachable!("ICE: All functions should have been inlined"), Value::Instruction { .. } | Value::Param { .. } => { - unreachable!("ICE: Should have been in cache {value:?}") + unreachable!("ICE: Should have been in cache {value_id} {value:?}") } }; self.ssa_values.insert(value_id, acir_value.clone()); diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir.rs b/crates/noirc_evaluator/src/ssa_refactor/ir.rs index 6d19dc0344..3ef680dda0 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/ir.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/ir.rs @@ -3,6 +3,7 @@ pub(crate) mod cfg; pub(crate) mod dfg; pub(crate) mod dom; pub(crate) mod function; +pub(crate) mod function_inserter; pub(crate) mod instruction; pub(crate) mod map; pub(crate) mod post_order; diff --git a/crates/noirc_evaluator/src/ssa_refactor/ir/function_inserter.rs b/crates/noirc_evaluator/src/ssa_refactor/ir/function_inserter.rs new file mode 100644 index 0000000000..31f236765c --- /dev/null +++ b/crates/noirc_evaluator/src/ssa_refactor/ir/function_inserter.rs @@ -0,0 +1,129 @@ +use std::collections::HashMap; + +use iter_extended::vecmap; + +use super::{ + basic_block::BasicBlockId, + dfg::InsertInstructionResult, + function::Function, + instruction::{Instruction, InstructionId}, + value::ValueId, +}; + +/// The FunctionInserter can be used to help modify existing Functions +/// and map old values to new values after re-inserting optimized versions +/// of old instructions. +pub(crate) struct FunctionInserter<'f> { + pub(crate) function: &'f mut Function, + + values: HashMap, +} + +impl<'f> FunctionInserter<'f> { + pub(crate) fn new(function: &'f mut Function) -> FunctionInserter<'f> { + Self { function, values: HashMap::new() } + } + + /// Resolves a ValueId to its new, updated value. + /// If there is no updated value for this id, this returns the same + /// ValueId that was passed in. + pub(crate) fn resolve(&mut self, value: ValueId) -> ValueId { + match self.values.get(&value) { + Some(value) => *value, + None => match &self.function.dfg[self.function.dfg.resolve(value)] { + super::value::Value::Array { array, element_type } => { + let array = array.clone(); + let element_type = element_type.clone(); + let new_array = array.iter().map(|id| self.resolve(*id)).collect(); + let new_id = self.function.dfg.make_array(new_array, element_type); + self.values.insert(value, new_id); + new_id + } + _ => value, + }, + } + } + + /// 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); + } + + /// Insert a key, value pair in the map + pub(crate) fn map_value(&mut self, key: ValueId, value: ValueId) { + self.values.insert(key, value); + } + + pub(crate) fn map_instruction(&mut self, id: InstructionId) -> Instruction { + self.function.dfg[id].clone().map_values(|id| self.resolve(id)) + } + + pub(crate) fn push_instruction(&mut self, id: InstructionId, block: BasicBlockId) { + let instruction = self.map_instruction(id); + self.push_instruction_value(instruction, id, block); + } + + pub(crate) fn push_instruction_value( + &mut self, + instruction: Instruction, + id: InstructionId, + block: BasicBlockId, + ) { + let results = self.function.dfg.instruction_results(id); + let results = vecmap(results, |id| self.function.dfg.resolve(*id)); + + let ctrl_typevars = instruction + .requires_ctrl_typevars() + .then(|| vecmap(&results, |result| self.function.dfg.type_of_value(*result))); + + let new_results = + self.function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars); + + Self::insert_new_instruction_results(&mut self.values, &results, new_results); + } + + /// Modify the values HashMap to remember the mapping between an instruction result's previous + /// ValueId (from the source_function) and its new ValueId in the destination function. + pub(crate) fn insert_new_instruction_results( + values: &mut HashMap, + old_results: &[ValueId], + new_results: InsertInstructionResult, + ) { + assert_eq!(old_results.len(), new_results.len()); + + match new_results { + InsertInstructionResult::SimplifiedTo(new_result) => { + values.insert(old_results[0], new_result); + } + InsertInstructionResult::Results(new_results) => { + for (old_result, new_result) in old_results.iter().zip(new_results) { + values.insert(*old_result, *new_result); + } + } + InsertInstructionResult::InstructionRemoved => (), + } + } + + pub(crate) fn remember_block_params(&mut self, block: BasicBlockId, new_values: &[ValueId]) { + let old_parameters = self.function.dfg.block_parameters(block); + + for (param, new_param) in old_parameters.iter().zip(new_values) { + // Don't overwrite any existing entries to avoid overwriting the induction variable + self.values.entry(*param).or_insert(*new_param); + } + } + + pub(crate) fn remember_block_params_from_block( + &mut self, + block: BasicBlockId, + new_block: BasicBlockId, + ) { + let old_parameters = self.function.dfg.block_parameters(block); + let new_parameters = self.function.dfg.block_parameters(new_block); + + for (param, new_param) in old_parameters.iter().zip(new_parameters) { + // Don't overwrite any existing entries to avoid overwriting the induction variable + self.values.entry(*param).or_insert(*new_param); + } + } +} diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/flatten_cfg.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/flatten_cfg.rs index 207e6e08cd..73f411a0ef 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/flatten_cfg.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/flatten_cfg.rs @@ -143,6 +143,7 @@ use crate::ssa_refactor::{ dfg::InsertInstructionResult, dom::DominatorTree, function::Function, + function_inserter::FunctionInserter, instruction::{BinaryOp, Instruction, InstructionId, TerminatorInstruction}, post_order::PostOrder, types::Type, @@ -167,7 +168,7 @@ impl Ssa { } struct Context<'f> { - function: &'f mut Function, + inserter: FunctionInserter<'f>, /// This ControlFlowGraph is the graph from before the function was modified by this flattening pass. cfg: ControlFlowGraph, @@ -185,11 +186,6 @@ struct Context<'f> { /// condition. If we are under multiple conditions (a nested if), the topmost condition is /// the most recent condition combined with all previous conditions via `And` instructions. conditions: Vec<(BasicBlockId, ValueId)>, - - /// A map of values from the unmodified function to their values given from this pass. - /// In particular, this pass will remove all block arguments except for function parameters. - /// Each value in the function's entry block is also left unchanged. - values: HashMap, } struct Store { @@ -215,11 +211,10 @@ fn flatten_function_cfg(function: &mut Function) { } let mut context = Context { cfg: ControlFlowGraph::with_function(function), - function, + inserter: FunctionInserter::new(function), store_values: HashMap::new(), branch_ends: HashMap::new(), conditions: Vec::new(), - values: HashMap::new(), }; context.flatten(); } @@ -230,19 +225,19 @@ impl<'f> Context<'f> { // Start with following the terminator of the entry block since we don't // need to flatten the entry block into itself. - self.handle_terminator(self.function.entry_block()); + self.handle_terminator(self.inserter.function.entry_block()); } /// Visits every block in the current function to find all blocks with a jmpif instruction and /// all blocks which terminate the jmpif by having each of its branches as a predecessor. fn analyze_function(&mut self) { - let post_order = PostOrder::with_function(self.function); + let post_order = PostOrder::with_function(self.inserter.function); let dom_tree = DominatorTree::with_cfg_and_post_order(&self.cfg, &post_order); let mut branch_beginnings = Vec::new(); let mut visited = HashSet::new(); let mut queue = VecDeque::new(); - queue.push_front(self.function.entry_block()); + queue.push_front(self.inserter.function.entry_block()); while let Some(block_id) = queue.pop_front() { // If multiple blocks branch to the same successor before we visit it we can end up in @@ -278,7 +273,7 @@ impl<'f> Context<'f> { self.branch_ends.insert(branch_beginning, block_id); } - let block = &self.function.dfg[block_id]; + let block = &self.inserter.function.dfg[block_id]; if let Some(TerminatorInstruction::JmpIf { .. }) = block.terminator() { branch_beginnings.push(block_id); } @@ -300,12 +295,12 @@ impl<'f> Context<'f> { /// Returns the last block to be inlined. This is either the return block of the function or, /// if self.conditions is not empty, the end block of the most recent condition. fn handle_terminator(&mut self, block: BasicBlockId) -> BasicBlockId { - match self.function.dfg[block].unwrap_terminator() { + match self.inserter.function.dfg[block].unwrap_terminator() { TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { let old_condition = *condition; - let then_condition = self.translate_value(old_condition); let then_block = *then_destination; let else_block = *else_destination; + let then_condition = self.inserter.resolve(old_condition); let one = FieldElement::one(); let then_branch = @@ -331,26 +326,21 @@ impl<'f> Context<'f> { return block; } } - let arguments = vecmap(arguments, |value| self.translate_value(*value)); - self.inline_block(*destination, &arguments) + let destination = *destination; + let arguments = vecmap(arguments.clone(), |value| self.inserter.resolve(value)); + self.inline_block(destination, &arguments) } TerminatorInstruction::Return { return_values } => { - let return_values = vecmap(return_values, |value| self.translate_value(*value)); - let entry = self.function.entry_block(); + let return_values = + vecmap(return_values.clone(), |value| self.inserter.resolve(value)); + let entry = self.inserter.function.entry_block(); let new_return = TerminatorInstruction::Return { return_values }; - self.function.dfg.set_block_terminator(entry, new_return); + self.inserter.function.dfg.set_block_terminator(entry, new_return); block } } } - /// Translate a value id from before the function was modified to one from after it has been - /// flattened. In particular, all block parameters should be removed, having been mapped to - /// their (merged) arguments, and all values from the entry block are unchanged. - fn translate_value(&self, value: ValueId) -> ValueId { - self.values.get(&value).copied().unwrap_or(value) - } - /// Push a condition to the stack of conditions. /// /// This condition should be present while we're inlining each block reachable from the 'then' @@ -373,8 +363,8 @@ impl<'f> Context<'f> { /// Unlike push_instruction, this function will not map any ValueIds. /// within the given instruction, nor will it modify self.values in any way. fn insert_instruction(&mut self, instruction: Instruction) -> ValueId { - let block = self.function.entry_block(); - self.function.dfg.insert_instruction_and_results(instruction, block, None).first() + let block = self.inserter.function.entry_block(); + self.inserter.function.dfg.insert_instruction_and_results(instruction, block, None).first() } /// Inserts a new instruction into the function's entry block, using the given @@ -386,8 +376,8 @@ impl<'f> Context<'f> { instruction: Instruction, ctrl_typevars: Option>, ) -> InsertInstructionResult { - let block = self.function.entry_block(); - self.function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars) + let block = self.inserter.function.entry_block(); + self.inserter.function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars) } /// Checks the branch condition on the top of the stack and uses it to build and insert an @@ -398,7 +388,9 @@ impl<'f> Context<'f> { fn insert_current_side_effects_enabled(&mut self) { let condition = match self.conditions.last() { Some((_, cond)) => *cond, - None => self.function.dfg.make_constant(FieldElement::one(), Type::unsigned(1)), + None => { + self.inserter.function.dfg.make_constant(FieldElement::one(), Type::unsigned(1)) + } }; let enable_side_effects = Instruction::EnableSideEffects { condition }; self.insert_instruction_with_typevars(enable_side_effects, None); @@ -413,15 +405,17 @@ impl<'f> Context<'f> { then_value: ValueId, else_value: ValueId, ) -> ValueId { - let block = self.function.entry_block(); + let block = self.inserter.function.entry_block(); let mul = Instruction::binary(BinaryOp::Mul, then_condition, then_value); - let then_value = self.function.dfg.insert_instruction_and_results(mul, block, None).first(); + let then_value = + self.inserter.function.dfg.insert_instruction_and_results(mul, block, None).first(); let mul = Instruction::binary(BinaryOp::Mul, else_condition, else_value); - let else_value = self.function.dfg.insert_instruction_and_results(mul, block, None).first(); + let else_value = + self.inserter.function.dfg.insert_instruction_and_results(mul, block, None).first(); let add = Instruction::binary(BinaryOp::Add, then_value, else_value); - self.function.dfg.insert_instruction_and_results(add, block, None).first() + self.inserter.function.dfg.insert_instruction_and_results(add, block, None).first() } /// Inline one branch of a jmpif instruction. @@ -455,8 +449,9 @@ impl<'f> Context<'f> { let old_stores = std::mem::take(&mut self.store_values); // Remember the old condition value is now known to be true/false within this branch - let known_value = self.function.dfg.make_constant(condition_value, Type::bool()); - self.values.insert(old_condition, known_value); + let known_value = + self.inserter.function.dfg.make_constant(condition_value, Type::bool()); + self.inserter.map_value(old_condition, known_value); let final_block = self.inline_block(destination, &[]); @@ -488,15 +483,17 @@ impl<'f> Context<'f> { ) -> BasicBlockId { assert_eq!(self.cfg.predecessors(destination).len(), 2); - let then_args = self.function.dfg[then_branch.last_block].terminator_arguments(); - let else_args = self.function.dfg[else_branch.last_block].terminator_arguments(); + let then_args = + self.inserter.function.dfg[then_branch.last_block].terminator_arguments().to_vec(); + let else_args = + self.inserter.function.dfg[else_branch.last_block].terminator_arguments().to_vec(); - let params = self.function.dfg.block_parameters(destination); + let params = self.inserter.function.dfg.block_parameters(destination); assert_eq!(params.len(), then_args.len()); assert_eq!(params.len(), else_args.len()); let args = vecmap(then_args.iter().zip(else_args), |(then_arg, else_arg)| { - (self.translate_value(*then_arg), self.translate_value(*else_arg)) + (self.inserter.resolve(*then_arg), self.inserter.resolve(else_arg)) }); // Cannot include this in the previous vecmap since it requires exclusive access to self @@ -543,7 +540,7 @@ impl<'f> Context<'f> { store_value.new_value = new_value; } else { let load = Instruction::Load { address }; - let load_type = Some(vec![self.function.dfg.type_of_value(new_value)]); + let load_type = Some(vec![self.inserter.function.dfg.type_of_value(new_value)]); let old_value = self.insert_instruction_with_typevars(load, load_type).first(); self.store_values.insert(address, Store { old_value, new_value }); @@ -555,20 +552,15 @@ impl<'f> Context<'f> { /// /// Returns the final block that was inlined. /// - /// Expects that the `arguments` given are already translated via self.translate_value. + /// Expects that the `arguments` given are already translated via self.inserter.resolve. /// If they are not, it is possible some values which no longer exist, such as block /// parameters, will be kept in the program. fn inline_block(&mut self, destination: BasicBlockId, arguments: &[ValueId]) -> BasicBlockId { - let parameters = self.function.dfg.block_parameters(destination); - Self::insert_new_instruction_results( - &mut self.values, - parameters, - InsertInstructionResult::Results(arguments), - ); + self.inserter.remember_block_params(destination, arguments); // If this is not a separate variable, clippy gets confused and says the to_vec is // unnecessary, when removing it actually causes an aliasing/mutability error. - let instructions = self.function.dfg[destination].instructions().to_vec(); + let instructions = self.inserter.function.dfg[destination].instructions().to_vec(); for instruction in instructions { self.push_instruction(instruction); } @@ -578,24 +570,15 @@ impl<'f> Context<'f> { /// Push the given instruction to the end of the entry block of the current function. /// - /// Note that each ValueId of the instruction will be mapped via self.translate_value. + /// Note that each ValueId of the instruction will be mapped via self.inserter.resolve. /// As a result, the instruction that will be pushed will actually be a new instruction /// with a different InstructionId from the original. The results of the given instruction /// will also be mapped to the results of the new instruction. fn push_instruction(&mut self, id: InstructionId) { - let instruction = self.function.dfg[id].map_values(|id| self.translate_value(id)); + let instruction = self.inserter.map_instruction(id); let instruction = self.handle_instruction_side_effects(instruction); - let results = self.function.dfg.instruction_results(id); - let results = vecmap(results, |id| self.function.dfg.resolve(*id)); - - let ctrl_typevars = instruction - .requires_ctrl_typevars() - .then(|| vecmap(&results, |result| self.function.dfg.type_of_value(*result))); - - let block = self.function.entry_block(); - let new_results = - self.function.dfg.insert_instruction_and_results(instruction, block, ctrl_typevars); - Self::insert_new_instruction_results(&mut self.values, &results, new_results); + let entry = self.inserter.function.entry_block(); + self.inserter.push_instruction_value(instruction, id, entry); } /// If we are currently in a branch, we need to modify constrain instructions @@ -623,26 +606,6 @@ impl<'f> Context<'f> { instruction } } - - fn insert_new_instruction_results( - values: &mut HashMap, - old_results: &[ValueId], - new_results: InsertInstructionResult, - ) { - assert_eq!(old_results.len(), new_results.len()); - - match new_results { - InsertInstructionResult::SimplifiedTo(new_result) => { - values.insert(old_results[0], new_result); - } - InsertInstructionResult::Results(new_results) => { - for (old_result, new_result) in old_results.iter().zip(new_results) { - values.insert(*old_result, *new_result); - } - } - InsertInstructionResult::InstructionRemoved => (), - } - } } #[cfg(test)] @@ -654,6 +617,7 @@ mod test { cfg::ControlFlowGraph, dfg::DataFlowGraph, function::RuntimeType, + function_inserter::FunctionInserter, instruction::{BinaryOp, Instruction, Intrinsic, TerminatorInstruction}, map::Id, types::Type, @@ -964,11 +928,10 @@ mod test { let function = ssa.main_mut(); let mut context = super::Context { cfg: ControlFlowGraph::with_function(function), - function, + inserter: FunctionInserter::new(function), store_values: HashMap::new(), branch_ends: HashMap::new(), conditions: Vec::new(), - values: HashMap::new(), }; context.analyze_function(); assert_eq!(context.branch_ends.len(), 2); diff --git a/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs b/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs index 2960afe101..9612099920 100644 --- a/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs +++ b/crates/noirc_evaluator/src/ssa_refactor/opt/unrolling.rs @@ -14,18 +14,11 @@ //! program that will need to be removed by a later simplify cfg pass. use std::collections::{HashMap, HashSet}; -use iter_extended::vecmap; - use crate::ssa_refactor::{ ir::{ - basic_block::BasicBlockId, - cfg::ControlFlowGraph, - dfg::InsertInstructionResult, - dom::DominatorTree, - function::Function, - instruction::{InstructionId, TerminatorInstruction}, - post_order::PostOrder, - value::ValueId, + basic_block::BasicBlockId, cfg::ControlFlowGraph, dfg::DataFlowGraph, dom::DominatorTree, + function::Function, function_inserter::FunctionInserter, + instruction::TerminatorInstruction, post_order::PostOrder, value::ValueId, }, ssa_gen::Ssa, }; @@ -222,22 +215,22 @@ fn unroll_loop_header<'a>( let fresh_block = function.dfg.make_block(); let mut context = LoopIteration::new(function, loop_, fresh_block, loop_.header); - let source_block = &context.function.dfg[context.source_block]; + let source_block = &context.dfg()[context.source_block]; assert_eq!(source_block.parameters().len(), 1, "Expected only 1 argument in loop header"); // Insert the current value of the loop induction variable into our context. let first_param = source_block.parameters()[0]; - context.values.insert(first_param, induction_value); + context.inserter.try_map_value(first_param, induction_value); context.inline_instructions_from_block(); - match context.function.dfg[fresh_block].unwrap_terminator() { + match context.dfg()[fresh_block].unwrap_terminator() { TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { let next_blocks = context.handle_jmpif(*condition, *then_destination, *else_destination); // If there is only 1 next block the jmpif evaluated to a single known block. // This is the expected case and lets us know if we should loop again or not. if next_blocks.len() == 1 { - context.function.dfg.inline_block(fresh_block, unroll_into); + context.dfg_mut().inline_block(fresh_block, unroll_into); // The fresh block is gone now so we're committing to insert into the original // unroll_into block from now on. @@ -257,14 +250,9 @@ fn unroll_loop_header<'a>( /// The context object for each loop iteration. /// Notably each loop iteration maps each loop block to a fresh, unrolled block. struct LoopIteration<'f> { - function: &'f mut Function, + inserter: FunctionInserter<'f>, loop_: &'f Loop, - /// Maps pre-unrolled ValueIds to unrolled ValueIds. - /// These will often be the exact same as before, unless the ValueId was - /// dependent on the loop induction variable which is changing on each iteration. - values: HashMap, - /// Maps pre-unrolled block ids from within the loop to new block ids of each loop /// block for each loop iteration. blocks: HashMap, @@ -291,11 +279,10 @@ impl<'f> LoopIteration<'f> { source_block: BasicBlockId, ) -> Self { Self { - function, + inserter: FunctionInserter::new(function), loop_, insert_block, source_block, - values: HashMap::new(), blocks: HashMap::new(), original_blocks: HashMap::new(), visited_blocks: HashSet::new(), @@ -341,7 +328,7 @@ impl<'f> LoopIteration<'f> { self.inline_instructions_from_block(); self.visited_blocks.insert(self.source_block); - match self.function.dfg[self.insert_block].unwrap_terminator() { + match self.inserter.function.dfg[self.insert_block].unwrap_terminator() { TerminatorInstruction::JmpIf { condition, then_destination, else_destination } => { self.handle_jmpif(*condition, *then_destination, *else_destination) } @@ -365,9 +352,9 @@ impl<'f> LoopIteration<'f> { then_destination: BasicBlockId, else_destination: BasicBlockId, ) -> Vec { - let condition = self.get_value(condition); + let condition = self.inserter.resolve(condition); - match self.function.dfg.get_numeric_constant(condition) { + match self.dfg().get_numeric_constant(condition) { Some(constant) => { let destination = if constant.is_zero() { else_destination } else { then_destination }; @@ -375,22 +362,13 @@ impl<'f> LoopIteration<'f> { self.source_block = self.get_original_block(destination); let jmp = TerminatorInstruction::Jmp { destination, arguments: Vec::new() }; - self.function.dfg.set_block_terminator(self.insert_block, jmp); + self.inserter.function.dfg.set_block_terminator(self.insert_block, jmp); vec![destination] } None => vec![then_destination, else_destination], } } - /// Map a ValueId in the original pre-unrolled ssa to its new id in the unrolled SSA. - /// This is often the same ValueId as most values don't change while unrolling. The main - /// exception is instructions referencing the induction variable (or the variable itself) - /// which may have been simplified to another form. Block parameters or values outside the - /// loop shouldn't change at all and won't be present inside self.values. - fn get_value(&self, value: ValueId) -> ValueId { - self.values.get(&value).copied().unwrap_or(value) - } - /// Translate a block id to a block id in the unrolled loop. If the given /// block id is not within the loop, it is returned as-is. fn get_or_insert_block(&mut self, block: BasicBlockId) -> BasicBlockId { @@ -400,15 +378,8 @@ impl<'f> LoopIteration<'f> { // If the block is in the loop we create a fresh block for each iteration if self.loop_.blocks.contains(&block) { - let new_block = self.function.dfg.make_block_with_parameters_from_block(block); - - let old_parameters = self.function.dfg.block_parameters(block); - let new_parameters = self.function.dfg.block_parameters(new_block); - - for (param, new_param) in old_parameters.iter().zip(new_parameters) { - // Don't overwrite any existing entries to avoid overwriting the induction variable - self.values.entry(*param).or_insert(*new_param); - } + let new_block = self.dfg_mut().make_block_with_parameters_from_block(block); + self.inserter.remember_block_params_from_block(block, new_block); self.blocks.insert(block, new_block); self.original_blocks.insert(new_block, block); @@ -423,62 +394,31 @@ impl<'f> LoopIteration<'f> { } fn inline_instructions_from_block(&mut self) { - let source_block = &self.function.dfg[self.source_block]; + let source_block = &self.dfg()[self.source_block]; let instructions = source_block.instructions().to_vec(); // We cannot directly append each instruction since we need to substitute any // instances of the induction variable or any values that were changed as a result // of the new induction variable value. for instruction in instructions { - self.push_instruction(instruction); + self.inserter.push_instruction(instruction, self.insert_block); } - let mut terminator = self.function.dfg[self.source_block] + let mut terminator = self.dfg()[self.source_block] .unwrap_terminator() - .map_values(|value| self.get_value(value)); + .clone() + .map_values(|value| self.inserter.resolve(value)); terminator.mutate_blocks(|block| self.get_or_insert_block(block)); - self.function.dfg.set_block_terminator(self.insert_block, terminator); + self.inserter.function.dfg.set_block_terminator(self.insert_block, terminator); } - fn push_instruction(&mut self, id: InstructionId) { - let instruction = self.function.dfg[id].map_values(|id| self.get_value(id)); - let results = self.function.dfg.instruction_results(id); - let results = vecmap(results, |id| self.function.dfg.resolve(*id)); - - let ctrl_typevars = instruction - .requires_ctrl_typevars() - .then(|| vecmap(&results, |result| self.function.dfg.type_of_value(*result))); - - let new_results = self.function.dfg.insert_instruction_and_results( - instruction, - self.insert_block, - ctrl_typevars, - ); - - Self::insert_new_instruction_results(&mut self.values, &results, new_results); + fn dfg(&self) -> &DataFlowGraph { + &self.inserter.function.dfg } - /// Modify the values HashMap to remember the mapping between an instruction result's previous - /// ValueId (from the source_function) and its new ValueId in the destination function. - fn insert_new_instruction_results( - values: &mut HashMap, - old_results: &[ValueId], - new_results: InsertInstructionResult, - ) { - assert_eq!(old_results.len(), new_results.len()); - - match new_results { - InsertInstructionResult::SimplifiedTo(new_result) => { - values.insert(old_results[0], new_result); - } - InsertInstructionResult::Results(new_results) => { - for (old_result, new_result) in old_results.iter().zip(new_results) { - values.insert(*old_result, *new_result); - } - } - InsertInstructionResult::InstructionRemoved => (), - } + fn dfg_mut(&mut self) -> &mut DataFlowGraph { + &mut self.inserter.function.dfg } }