xref: /llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp (revision 7f50c15be5c020b2789f83dc3c87d28f8294664f)
1 //=- WebAssemblyFixBrTableDefaults.cpp - Fix br_table default branch targets -//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 ///
9 /// \file This file implements a pass that eliminates redundant range checks
10 /// guarding br_table instructions. Since jump tables on most targets cannot
11 /// handle out of range indices, LLVM emits these checks before most jump
12 /// tables. But br_table takes a default branch target as an argument, so it
13 /// does not need the range checks.
14 ///
15 //===----------------------------------------------------------------------===//
16 
17 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
18 #include "WebAssembly.h"
19 #include "llvm/CodeGen/MachineFunction.h"
20 #include "llvm/CodeGen/MachineFunctionPass.h"
21 #include "llvm/CodeGen/MachineRegisterInfo.h"
22 #include "llvm/Pass.h"
23 
24 using namespace llvm;
25 
26 #define DEBUG_TYPE "wasm-fix-br-table-defaults"
27 
28 namespace {
29 
30 class WebAssemblyFixBrTableDefaults final : public MachineFunctionPass {
31   StringRef getPassName() const override {
32     return "WebAssembly Fix br_table Defaults";
33   }
34 
35   bool runOnMachineFunction(MachineFunction &MF) override;
36 
37 public:
38   static char ID; // Pass identification, replacement for typeid
39   WebAssemblyFixBrTableDefaults() : MachineFunctionPass(ID) {}
40 };
41 
42 char WebAssemblyFixBrTableDefaults::ID = 0;
43 
44 // `MI` is a br_table instruction missing its default target argument. This
45 // function finds and adds the default target argument and removes any redundant
46 // range check preceding the br_table.
47 MachineBasicBlock *fixBrTable(MachineInstr &MI, MachineBasicBlock *MBB,
48                               MachineFunction &MF) {
49   // Get the header block, which contains the redundant range check.
50   assert(MBB->pred_size() == 1 && "Expected a single guard predecessor");
51   auto *HeaderMBB = *MBB->pred_begin();
52 
53   // Find the conditional jump to the default target. If it doesn't exist, the
54   // default target is unreachable anyway, so we can choose anything.
55   auto JumpMII = --HeaderMBB->end();
56   while (JumpMII->getOpcode() != WebAssembly::BR_IF &&
57          JumpMII != HeaderMBB->begin()) {
58     --JumpMII;
59   }
60   if (JumpMII->getOpcode() == WebAssembly::BR_IF) {
61     // Install the default target and remove the jumps in the header.
62     auto *DefaultMBB = JumpMII->getOperand(0).getMBB();
63     assert(DefaultMBB != MBB && "Expected conditional jump to default target");
64     MI.addOperand(MF, MachineOperand::CreateMBB(DefaultMBB));
65     HeaderMBB->erase(JumpMII, HeaderMBB->end());
66   } else {
67     // Arbitrarily choose the first jump target as the default.
68     auto *SomeMBB = MI.getOperand(1).getMBB();
69     MI.addOperand(MachineOperand::CreateMBB(SomeMBB));
70   }
71 
72   // Splice the jump table into the header.
73   HeaderMBB->splice(HeaderMBB->end(), MBB, MBB->begin(), MBB->end());
74 
75   // Update CFG to skip the old jump table block. Remove shared successors
76   // before transferring to avoid duplicated successors.
77   HeaderMBB->removeSuccessor(MBB);
78   for (auto &Succ : MBB->successors())
79     if (HeaderMBB->isSuccessor(Succ))
80       HeaderMBB->removeSuccessor(Succ);
81   HeaderMBB->transferSuccessorsAndUpdatePHIs(MBB);
82 
83   // Remove the old jump table block from the function
84   MF.erase(MBB);
85 
86   return HeaderMBB;
87 }
88 
89 bool WebAssemblyFixBrTableDefaults::runOnMachineFunction(MachineFunction &MF) {
90   LLVM_DEBUG(dbgs() << "********** Fixing br_table Default Targets **********\n"
91                        "********** Function: "
92                     << MF.getName() << '\n');
93 
94   bool Changed = false;
95   SmallPtrSet<MachineBasicBlock *, 16> MBBSet;
96   for (auto &MBB : MF)
97     MBBSet.insert(&MBB);
98 
99   while (!MBBSet.empty()) {
100     MachineBasicBlock *MBB = *MBBSet.begin();
101     MBBSet.erase(MBB);
102     for (auto &MI : *MBB) {
103       if (WebAssembly::isBrTable(MI)) {
104         auto *Fixed = fixBrTable(MI, MBB, MF);
105         MBBSet.erase(Fixed);
106         Changed = true;
107         break;
108       }
109     }
110   }
111 
112   if (Changed) {
113     // We rewrote part of the function; recompute relevant things.
114     MF.RenumberBlocks();
115     return true;
116   }
117 
118   return false;
119 }
120 
121 } // end anonymous namespace
122 
123 INITIALIZE_PASS(WebAssemblyFixBrTableDefaults, DEBUG_TYPE,
124                 "Removes range checks and sets br_table default targets", false,
125                 false);
126 
127 FunctionPass *llvm::createWebAssemblyFixBrTableDefaults() {
128   return new WebAssemblyFixBrTableDefaults();
129 }
130