Fix missing modification of jump table in licm

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.

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

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.
+    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<Va
     true
 }
 
-// Traverses a loop in reverse post-order from a header block and identify loop-invariant
-// instructions. These loop-invariant instructions are then removed from the code and returned
-// (in reverse post-order) for later use.
+/// Traverses a loop in reverse post-order from a header block and identify loop-invariant
+/// instructions. These loop-invariant instructions are then removed from the code and returned
+/// (in reverse post-order) for later use.
 fn remove_loop_invariant_instructions(
     lp: Loop,
     func: &mut Function,

cranelift/filetests/filetests/licm/br-table.clif

@@ -0,0 +1,19 @@
+test compile
+set opt_level=speed_and_size
+target x86_64
+
+function %br_table_opt() {
+    jt0 = jump_table [block1, block2]
+
+    block0:
+        v0 = iconst.i32 1
+        br_table v0, block2, jt0
+
+    block1: 
+        return
+
+    block2:
+        v1 = iconst.i32 1
+        jump block2
+
+}

cranelift/filetests/filetests/licm/rewrite-jump-table.clif

@@ -0,0 +1,28 @@
+test licm
+target aarch64
+
+function %rewrite_jump_table() {
+    jt0 = jump_table [block1, block2]
+
+    block0:
+        v0 = iconst.i64 1
+        v1 = jump_table_base.i64 jt0
+        v2 = jump_table_entry.i64 v0, v1, 4, jt0
+        v3 = iadd v1, v2
+        indirect_jump_table_br v3, jt0
+
+    block1:
+        return
+
+    block2:
+        v4 = bconst.b1 false
+        jump block2
+}
+
+; sameln: function
+; nextln: jt0 = jump_table [block1, block3]
+; check: block3:
+; nextln: v4 = bconst.b1 false
+; nextln: jump block2
+; check: block2:
+; nextln: jump block2