diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp --- a/llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp +++ b/llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp @@ -52,24 +52,32 @@ // Find the conditional jump to the default target. If it doesn't exist, the // default target is unreachable anyway, so we can choose anything. - auto JumpMII = --HeaderMBB->end(); - while (JumpMII->getOpcode() != WebAssembly::BR_IF && - JumpMII != HeaderMBB->begin()) { - --JumpMII; - } - if (JumpMII->getOpcode() == WebAssembly::BR_IF) { - // Install the default target and remove the jumps in the header. - auto *DefaultMBB = JumpMII->getOperand(0).getMBB(); - assert(DefaultMBB != MBB && "Expected conditional jump to default target"); - MI.addOperand(MF, MachineOperand::CreateMBB(DefaultMBB)); - HeaderMBB->erase(JumpMII, HeaderMBB->end()); + MachineBasicBlock *TBB = nullptr, *FBB = nullptr; + SmallVector Cond; + const auto &TII = *MF.getSubtarget().getInstrInfo(); + TII.analyzeBranch(*HeaderMBB, TBB, FBB, Cond); + + // Here are the possible outcomes. '_' is nullptr, `J` is the jump table block + // aka MBB, 'D' is the default block. + // + // TBB | FBB | Meaning + // _ | _ | No default block, header falls through to jump table + // J | _ | No default block, header jumps to the jump table + // D | _ | Header jumps to the default and falls through to the jump table + // D | J | Header jumps to the default and also to the jump table + if (TBB && TBB != MBB) { + // Install the default target. + assert((FBB == nullptr || FBB == MBB) && + "Expected jump or fallthrough to br_table block"); + MI.addOperand(MF, MachineOperand::CreateMBB(TBB)); } else { // Arbitrarily choose the first jump target as the default. auto *SomeMBB = MI.getOperand(1).getMBB(); MI.addOperand(MachineOperand::CreateMBB(SomeMBB)); } - // Splice the jump table into the header. + // Remove any branches from the header and splice in the jump table instead + TII.removeBranch(*HeaderMBB, nullptr); HeaderMBB->splice(HeaderMBB->end(), MBB, MBB->begin(), MBB->end()); // Update CFG to skip the old jump table block. Remove shared successors diff --git a/llvm/test/CodeGen/WebAssembly/switch-unreachable-default.ll b/llvm/test/CodeGen/WebAssembly/switch-unreachable-default.ll --- a/llvm/test/CodeGen/WebAssembly/switch-unreachable-default.ll +++ b/llvm/test/CodeGen/WebAssembly/switch-unreachable-default.ll @@ -36,3 +36,52 @@ unreachable: unreachable } + +; CHECK-LABEL: split: +; CHECK: .functype split (i32) -> () +; CHECK: block +; CHECK: br_if 0 +; CHECK: block +; CHECK: block +; CHECK: br_table {1, 1, 1, 1, 1, 1, 1, 0} +; CHECK: .LBB1_2 +; CHECK: end_block +; CHECK: br_table {0, 0, 0} +; CHECK: .LBB1_3 +; CHECK: end_block +; CHECK: unreachable +; CHECK: .LBB1_4 +; CHECK: end_block +; CHECK: end_function +define void @split(i8 %c) { +entry: + switch i8 %c, label %sw.default [ + i8 114, label %return + i8 103, label %sw.bb1 + i8 98, label %sw.bb2 + i8 97, label %sw.bb3 + i8 48, label %sw.bb4 + i8 49, label %sw.bb5 + ] + +sw.bb1: + unreachable + +sw.bb2: + unreachable + +sw.bb3: + unreachable + +sw.bb4: + unreachable + +sw.bb5: + unreachable + +sw.default: + unreachable + +return: + ret void +}