1 //===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
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 /// \file
9 ///
10 /// Implements an expansion pass to turn `indirectbr` instructions in the IR
11 /// into `switch` instructions. This works by enumerating the basic blocks in
12 /// a dense range of integers, replacing each `blockaddr` constant with the
13 /// corresponding integer constant, and then building a switch that maps from
14 /// the integers to the actual blocks. All of the indirectbr instructions in the
15 /// function are redirected to this common switch.
16 ///
17 /// While this is generically useful if a target is unable to codegen
18 /// `indirectbr` natively, it is primarily useful when there is some desire to
19 /// get the builtin non-jump-table lowering of a switch even when the input
20 /// source contained an explicit indirect branch construct.
21 ///
22 /// Note that it doesn't make any sense to enable this pass unless a target also
23 /// disables jump-table lowering of switches. Doing that is likely to pessimize
24 /// the code.
25 ///
26 //===----------------------------------------------------------------------===//
27 
28 #include "llvm/ADT/STLExtras.h"
29 #include "llvm/ADT/Sequence.h"
30 #include "llvm/ADT/SmallVector.h"
31 #include "llvm/Analysis/DomTreeUpdater.h"
32 #include "llvm/CodeGen/TargetPassConfig.h"
33 #include "llvm/CodeGen/TargetSubtargetInfo.h"
34 #include "llvm/IR/BasicBlock.h"
35 #include "llvm/IR/Dominators.h"
36 #include "llvm/IR/Function.h"
37 #include "llvm/IR/IRBuilder.h"
38 #include "llvm/IR/InstIterator.h"
39 #include "llvm/IR/Instruction.h"
40 #include "llvm/IR/Instructions.h"
41 #include "llvm/InitializePasses.h"
42 #include "llvm/Pass.h"
43 #include "llvm/Support/Debug.h"
44 #include "llvm/Support/ErrorHandling.h"
45 #include "llvm/Support/raw_ostream.h"
46 #include "llvm/Target/TargetMachine.h"
47 
48 using namespace llvm;
49 
50 #define DEBUG_TYPE "indirectbr-expand"
51 
52 namespace {
53 
54 class IndirectBrExpandPass : public FunctionPass {
55   const TargetLowering *TLI = nullptr;
56 
57 public:
58   static char ID; // Pass identification, replacement for typeid
59 
60   IndirectBrExpandPass() : FunctionPass(ID) {
61     initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry());
62   }
63 
64   void getAnalysisUsage(AnalysisUsage &AU) const override {
65     AU.addPreserved<DominatorTreeWrapperPass>();
66   }
67 
68   bool runOnFunction(Function &F) override;
69 };
70 
71 } // end anonymous namespace
72 
73 char IndirectBrExpandPass::ID = 0;
74 
75 INITIALIZE_PASS_BEGIN(IndirectBrExpandPass, DEBUG_TYPE,
76                       "Expand indirectbr instructions", false, false)
77 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
78 INITIALIZE_PASS_END(IndirectBrExpandPass, DEBUG_TYPE,
79                     "Expand indirectbr instructions", false, false)
80 
81 FunctionPass *llvm::createIndirectBrExpandPass() {
82   return new IndirectBrExpandPass();
83 }
84 
85 bool IndirectBrExpandPass::runOnFunction(Function &F) {
86   auto &DL = F.getParent()->getDataLayout();
87   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
88   if (!TPC)
89     return false;
90 
91   auto &TM = TPC->getTM<TargetMachine>();
92   auto &STI = *TM.getSubtargetImpl(F);
93   if (!STI.enableIndirectBrExpand())
94     return false;
95   TLI = STI.getTargetLowering();
96 
97   Optional<DomTreeUpdater> DTU;
98   if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
99     DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
100 
101   SmallVector<IndirectBrInst *, 1> IndirectBrs;
102 
103   // Set of all potential successors for indirectbr instructions.
104   SmallPtrSet<BasicBlock *, 4> IndirectBrSuccs;
105 
106   // Build a list of indirectbrs that we want to rewrite.
107   for (BasicBlock &BB : F)
108     if (auto *IBr = dyn_cast<IndirectBrInst>(BB.getTerminator())) {
109       // Handle the degenerate case of no successors by replacing the indirectbr
110       // with unreachable as there is no successor available.
111       if (IBr->getNumSuccessors() == 0) {
112         (void)new UnreachableInst(F.getContext(), IBr);
113         IBr->eraseFromParent();
114         continue;
115       }
116 
117       IndirectBrs.push_back(IBr);
118       for (BasicBlock *SuccBB : IBr->successors())
119         IndirectBrSuccs.insert(SuccBB);
120     }
121 
122   if (IndirectBrs.empty())
123     return false;
124 
125   // If we need to replace any indirectbrs we need to establish integer
126   // constants that will correspond to each of the basic blocks in the function
127   // whose address escapes. We do that here and rewrite all the blockaddress
128   // constants to just be those integer constants cast to a pointer type.
129   SmallVector<BasicBlock *, 4> BBs;
130 
131   for (BasicBlock &BB : F) {
132     // Skip blocks that aren't successors to an indirectbr we're going to
133     // rewrite.
134     if (!IndirectBrSuccs.count(&BB))
135       continue;
136 
137     auto IsBlockAddressUse = [&](const Use &U) {
138       return isa<BlockAddress>(U.getUser());
139     };
140     auto BlockAddressUseIt = llvm::find_if(BB.uses(), IsBlockAddressUse);
141     if (BlockAddressUseIt == BB.use_end())
142       continue;
143 
144     assert(std::find_if(std::next(BlockAddressUseIt), BB.use_end(),
145                         IsBlockAddressUse) == BB.use_end() &&
146            "There should only ever be a single blockaddress use because it is "
147            "a constant and should be uniqued.");
148 
149     auto *BA = cast<BlockAddress>(BlockAddressUseIt->getUser());
150 
151     // Skip if the constant was formed but ended up not being used (due to DCE
152     // or whatever).
153     if (!BA->isConstantUsed())
154       continue;
155 
156     // Compute the index we want to use for this basic block. We can't use zero
157     // because null can be compared with block addresses.
158     int BBIndex = BBs.size() + 1;
159     BBs.push_back(&BB);
160 
161     auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
162     ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);
163 
164     // Now rewrite the blockaddress to an integer constant based on the index.
165     // FIXME: This part doesn't properly recognize other uses of blockaddress
166     // expressions, for instance, where they are used to pass labels to
167     // asm-goto. This part of the pass needs a rework.
168     BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));
169   }
170 
171   if (BBs.empty()) {
172     // There are no blocks whose address is taken, so any indirectbr instruction
173     // cannot get a valid input and we can replace all of them with unreachable.
174     SmallVector<DominatorTree::UpdateType, 8> Updates;
175     if (DTU)
176       Updates.reserve(IndirectBrSuccs.size());
177     for (auto *IBr : IndirectBrs) {
178       if (DTU) {
179         for (BasicBlock *SuccBB : IBr->successors())
180           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
181       }
182       (void)new UnreachableInst(F.getContext(), IBr);
183       IBr->eraseFromParent();
184     }
185     if (DTU) {
186       assert(Updates.size() == IndirectBrSuccs.size() &&
187              "Got unexpected update count.");
188       DTU->applyUpdates(Updates);
189     }
190     return true;
191   }
192 
193   BasicBlock *SwitchBB;
194   Value *SwitchValue;
195 
196   // Compute a common integer type across all the indirectbr instructions.
197   IntegerType *CommonITy = nullptr;
198   for (auto *IBr : IndirectBrs) {
199     auto *ITy =
200         cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));
201     if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
202       CommonITy = ITy;
203   }
204 
205   auto GetSwitchValue = [DL, CommonITy](IndirectBrInst *IBr) {
206     return CastInst::CreatePointerCast(
207         IBr->getAddress(), CommonITy,
208         Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr);
209   };
210 
211   SmallVector<DominatorTree::UpdateType, 8> Updates;
212 
213   if (IndirectBrs.size() == 1) {
214     // If we only have one indirectbr, we can just directly replace it within
215     // its block.
216     IndirectBrInst *IBr = IndirectBrs[0];
217     SwitchBB = IBr->getParent();
218     SwitchValue = GetSwitchValue(IBr);
219     if (DTU) {
220       Updates.reserve(IndirectBrSuccs.size());
221       for (BasicBlock *SuccBB : IBr->successors())
222         Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
223       assert(Updates.size() == IndirectBrSuccs.size() &&
224              "Got unexpected update count.");
225     }
226     IBr->eraseFromParent();
227   } else {
228     // Otherwise we need to create a new block to hold the switch across BBs,
229     // jump to that block instead of each indirectbr, and phi together the
230     // values for the switch.
231     SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
232     auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
233                                      "switch_value_phi", SwitchBB);
234     SwitchValue = SwitchPN;
235 
236     // Now replace the indirectbr instructions with direct branches to the
237     // switch block and fill out the PHI operands.
238     if (DTU)
239       Updates.reserve(IndirectBrs.size() + 2 * IndirectBrSuccs.size());
240     for (auto *IBr : IndirectBrs) {
241       SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
242       BranchInst::Create(SwitchBB, IBr);
243       if (DTU) {
244         Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
245         for (BasicBlock *SuccBB : IBr->successors())
246           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
247       }
248       IBr->eraseFromParent();
249     }
250   }
251 
252   // Now build the switch in the block. The block will have no terminator
253   // already.
254   auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);
255 
256   // Add a case for each block.
257   for (int i : llvm::seq<int>(1, BBs.size()))
258     SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);
259 
260   if (DTU) {
261     // If there were multiple indirectbr's, they may have common successors,
262     // but in the dominator tree, we only track unique edges.
263     SmallPtrSet<BasicBlock *, 8> UniqueSuccessors(BBs.begin(), BBs.end());
264     Updates.reserve(Updates.size() + UniqueSuccessors.size());
265     for (BasicBlock *BB : UniqueSuccessors)
266       Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
267     DTU->applyUpdates(Updates);
268   }
269 
270   return true;
271 }
272