diff --git a/cranelift/codegen/src/ir/function.rs b/cranelift/codegen/src/ir/function.rs index 92911c8a59..1833af27f5 100644 --- a/cranelift/codegen/src/ir/function.rs +++ b/cranelift/codegen/src/ir/function.rs @@ -7,8 +7,9 @@ use crate::binemit::CodeOffset; use crate::entity::{PrimaryMap, SecondaryMap}; use crate::ir; use crate::ir::{ - Block, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Heap, HeapData, Inst, JumpTable, - JumpTableData, Opcode, SigRef, StackSlot, StackSlotData, Table, TableData, + instructions::BranchInfo, Block, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Heap, + HeapData, Inst, InstructionData, JumpTable, JumpTableData, Opcode, SigRef, StackSlot, + StackSlotData, Table, TableData, }; use crate::ir::{BlockOffsets, InstEncodings, SourceLocs, StackSlots, ValueLocations}; use crate::ir::{DataFlowGraph, ExternalName, Layout, Signature}; @@ -270,6 +271,8 @@ impl Function { /// Changes the destination of a jump or branch instruction. /// Does nothing if called with a non-jump or non-branch instruction. + /// + /// Note that this method ignores multi-destination branches like `br_table`. pub fn change_branch_destination(&mut self, inst: Inst, new_dest: Block) { match self.dfg[inst].branch_destination_mut() { None => (), @@ -277,6 +280,43 @@ impl Function { } } + /// Rewrite the branch destination to `new_dest` if the destination matches `old_dest`. + /// Does nothing if called with a non-jump or non-branch instruction. + /// + /// Unlike [change_branch_destination](Function::change_branch_destination), this method rewrite the destinations of + /// multi-destination branches like `br_table`. + pub fn rewrite_branch_destination(&mut self, inst: Inst, old_dest: Block, new_dest: Block) { + match self.dfg.analyze_branch(inst) { + BranchInfo::SingleDest(dest, ..) => { + if dest == old_dest { + self.change_branch_destination(inst, new_dest); + } + } + + BranchInfo::Table(table, default_dest) => { + self.jump_tables[table].iter_mut().for_each(|entry| { + if *entry == old_dest { + *entry = new_dest; + } + }); + + if default_dest == Some(old_dest) { + match &mut self.dfg[inst] { + InstructionData::BranchTable { destination, .. } => { + *destination = new_dest; + } + _ => panic!( + "Unexpected instruction {} having default destination", + self.dfg.display_inst(inst, None) + ), + } + } + } + + BranchInfo::NotABranch => {} + } + } + /// Checks that the specified block can be encoded as a basic block. /// /// On error, returns the first invalid instruction and an error message. diff --git a/cranelift/codegen/src/ir/instructions.rs b/cranelift/codegen/src/ir/instructions.rs index 4a14b77f33..13310bc01c 100644 --- a/cranelift/codegen/src/ir/instructions.rs +++ b/cranelift/codegen/src/ir/instructions.rs @@ -277,7 +277,7 @@ impl InstructionData { ref mut destination, .. } => Some(destination), - Self::BranchTable { .. } => None, + Self::BranchTable { .. } | Self::IndirectJump { .. } => None, _ => { debug_assert!(!self.opcode().is_branch()); None diff --git a/cranelift/codegen/src/licm.rs b/cranelift/codegen/src/licm.rs index 75000b5297..5e9e0c1262 100644 --- a/cranelift/codegen/src/licm.rs +++ b/cranelift/codegen/src/licm.rs @@ -61,8 +61,8 @@ pub fn do_licm( domtree.compute(func, cfg); } -// Insert a pre-header before the header, modifying the function layout and CFG to reflect it. -// A jump instruction to the header is placed at the end of the pre-header. +/// Insert a pre-header before the header, modifying the function layout and CFG to reflect it. +/// A jump instruction to the header is placed at the end of the pre-header. fn create_pre_header( isa: &dyn TargetIsa, header: Block, @@ -81,30 +81,31 @@ fn create_pre_header( for typ in header_args_types { pre_header_args_value.push(func.dfg.append_block_param(pre_header, typ), pool); } + for BlockPredecessor { inst: last_inst, .. } in cfg.pred_iter(header) { // We only follow normal edges (not the back edges) if !domtree.dominates(header, last_inst, &func.layout) { - func.change_branch_destination(last_inst, pre_header); + func.rewrite_branch_destination(last_inst, header, pre_header); } } - { - let mut pos = EncCursor::new(func, isa).at_top(header); - // Inserts the pre-header at the right place in the layout. - pos.insert_block(pre_header); - pos.next_inst(); - pos.ins().jump(header, pre_header_args_value.as_slice(pool)); - } + + // Inserts the pre-header at the right place in the layout. + let mut pos = EncCursor::new(func, isa).at_top(header); + pos.insert_block(pre_header); + pos.next_inst(); + pos.ins().jump(header, pre_header_args_value.as_slice(pool)); + pre_header } -// Detects if a loop header has a natural pre-header. -// -// A loop header has a pre-header if there is only one predecessor that the header doesn't -// dominate. -// Returns the pre-header Block and the instruction jumping to the header. +/// Detects if a loop header has a natural pre-header. +/// +/// A loop header has a pre-header if there is only one predecessor that the header doesn't +/// dominate. +/// Returns the pre-header Block and the instruction jumping to the header. fn has_pre_header( layout: &Layout, cfg: &ControlFlowGraph, @@ -176,9 +177,9 @@ fn is_loop_invariant(inst: Inst, dfg: &DataFlowGraph, loop_values: &FxHashSet