1 //===- BlockExtractor.cpp - Extracts blocks into their own functions ------===//
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 // This pass extracts the specified basic blocks from the module into their
10 // own functions.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "llvm/ADT/STLExtras.h"
15 #include "llvm/ADT/Statistic.h"
16 #include "llvm/IR/Instructions.h"
17 #include "llvm/IR/Module.h"
18 #include "llvm/Pass.h"
19 #include "llvm/Support/CommandLine.h"
20 #include "llvm/Support/Debug.h"
21 #include "llvm/Support/MemoryBuffer.h"
22 #include "llvm/Transforms/IPO.h"
23 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
24 #include "llvm/Transforms/Utils/CodeExtractor.h"
25 using namespace llvm;
26 
27 #define DEBUG_TYPE "block-extractor"
28 
29 STATISTIC(NumExtracted, "Number of basic blocks extracted");
30 
31 static cl::opt<std::string> BlockExtractorFile(
32     "extract-blocks-file", cl::value_desc("filename"),
33     cl::desc("A file containing list of basic blocks to extract"), cl::Hidden);
34 
35 cl::opt<bool> BlockExtractorEraseFuncs("extract-blocks-erase-funcs",
36                                        cl::desc("Erase the existing functions"),
37                                        cl::Hidden);
38 
39 namespace {
40 class BlockExtractor : public ModulePass {
41   SmallVector<BasicBlock *, 16> Blocks;
42   bool EraseFunctions;
43   SmallVector<std::pair<std::string, std::string>, 32> BlocksByName;
44 
45 public:
46   static char ID;
47   BlockExtractor(const SmallVectorImpl<BasicBlock *> &BlocksToExtract,
48                  bool EraseFunctions)
49       : ModulePass(ID), Blocks(BlocksToExtract.begin(), BlocksToExtract.end()),
50         EraseFunctions(EraseFunctions) {
51     if (!BlockExtractorFile.empty())
52       loadFile();
53   }
54   BlockExtractor() : BlockExtractor(SmallVector<BasicBlock *, 0>(), false) {}
55   bool runOnModule(Module &M) override;
56 
57 private:
58   void loadFile();
59   void splitLandingPadPreds(Function &F);
60 };
61 } // end anonymous namespace
62 
63 char BlockExtractor::ID = 0;
64 INITIALIZE_PASS(BlockExtractor, "extract-blocks",
65                 "Extract basic blocks from module", false, false)
66 
67 ModulePass *llvm::createBlockExtractorPass() { return new BlockExtractor(); }
68 ModulePass *llvm::createBlockExtractorPass(
69     const SmallVectorImpl<BasicBlock *> &BlocksToExtract, bool EraseFunctions) {
70   return new BlockExtractor(BlocksToExtract, EraseFunctions);
71 }
72 
73 /// Gets all of the blocks specified in the input file.
74 void BlockExtractor::loadFile() {
75   auto ErrOrBuf = MemoryBuffer::getFile(BlockExtractorFile);
76   if (ErrOrBuf.getError())
77     report_fatal_error("BlockExtractor couldn't load the file.");
78   // Read the file.
79   auto &Buf = *ErrOrBuf;
80   SmallVector<StringRef, 16> Lines;
81   Buf->getBuffer().split(Lines, '\n', /*MaxSplit=*/-1,
82                          /*KeepEmpty=*/false);
83   for (const auto &Line : Lines) {
84     auto FBPair = Line.split(' ');
85     BlocksByName.push_back({FBPair.first, FBPair.second});
86   }
87 }
88 
89 /// Extracts the landing pads to make sure all of them have only one
90 /// predecessor.
91 void BlockExtractor::splitLandingPadPreds(Function &F) {
92   for (BasicBlock &BB : F) {
93     for (Instruction &I : BB) {
94       if (!isa<InvokeInst>(&I))
95         continue;
96       InvokeInst *II = cast<InvokeInst>(&I);
97       BasicBlock *Parent = II->getParent();
98       BasicBlock *LPad = II->getUnwindDest();
99 
100       // Look through the landing pad's predecessors. If one of them ends in an
101       // 'invoke', then we want to split the landing pad.
102       bool Split = false;
103       for (auto PredBB : predecessors(LPad)) {
104         if (PredBB->isLandingPad() && PredBB != Parent &&
105             isa<InvokeInst>(Parent->getTerminator())) {
106           Split = true;
107           break;
108         }
109       }
110 
111       if (!Split)
112         continue;
113 
114       SmallVector<BasicBlock *, 2> NewBBs;
115       SplitLandingPadPredecessors(LPad, Parent, ".1", ".2", NewBBs);
116     }
117   }
118 }
119 
120 bool BlockExtractor::runOnModule(Module &M) {
121 
122   bool Changed = false;
123 
124   // Get all the functions.
125   SmallVector<Function *, 4> Functions;
126   for (Function &F : M) {
127     splitLandingPadPreds(F);
128     Functions.push_back(&F);
129   }
130 
131   // Get all the blocks specified in the input file.
132   for (const auto &BInfo : BlocksByName) {
133     Function *F = M.getFunction(BInfo.first);
134     if (!F)
135       report_fatal_error("Invalid function name specified in the input file");
136     auto Res = llvm::find_if(*F, [&](const BasicBlock &BB) {
137       return BB.getName().equals(BInfo.second);
138     });
139     if (Res == F->end())
140       report_fatal_error("Invalid block name specified in the input file");
141     Blocks.push_back(&*Res);
142   }
143 
144   // Extract basic blocks.
145   for (BasicBlock *BB : Blocks) {
146     // Check if the module contains BB.
147     if (BB->getParent()->getParent() != &M)
148       report_fatal_error("Invalid basic block");
149     LLVM_DEBUG(dbgs() << "BlockExtractor: Extracting "
150                       << BB->getParent()->getName() << ":" << BB->getName()
151                       << "\n");
152     SmallVector<BasicBlock *, 2> BlocksToExtractVec;
153     BlocksToExtractVec.push_back(BB);
154     if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
155       BlocksToExtractVec.push_back(II->getUnwindDest());
156     CodeExtractor(BlocksToExtractVec).extractCodeRegion();
157     ++NumExtracted;
158     Changed = true;
159   }
160 
161   // Erase the functions.
162   if (EraseFunctions || BlockExtractorEraseFuncs) {
163     for (Function *F : Functions) {
164       LLVM_DEBUG(dbgs() << "BlockExtractor: Trying to delete " << F->getName()
165                         << "\n");
166       F->deleteBody();
167     }
168     // Set linkage as ExternalLinkage to avoid erasing unreachable functions.
169     for (Function &F : M)
170       F.setLinkage(GlobalValue::ExternalLinkage);
171     Changed = true;
172   }
173 
174   return Changed;
175 }
176