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