lib/Transforms/IPO/ExtractFunction.cpp

   1 //===-- ExtractFunction.cpp - Function extraction pass --------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file was developed by the LLVM research group and is distributed under
   6 // the University of Illinois Open Source License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This pass extracts
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Instructions.h"
  15 #include "llvm/Module.h"
  16 #include "llvm/Pass.h"
  17 #include "llvm/Transforms/IPO.h"
  18 #include "llvm/Support/Compiler.h"
  19 using namespace llvm;
  20
  21 namespace {
  22   /// @brief A pass to extract specific functions and their dependencies.
  23   class VISIBILITY_HIDDEN FunctionExtractorPass : public ModulePass {
  24     Function *Named;
  25     bool deleteFunc;
  26     bool reLink;
  27   public:
  28     /// FunctionExtractorPass - If deleteFn is true, this pass deletes as the
  29     /// specified function. Otherwise, it deletes as much of the module as
  30     /// possible, except for the function specified.
  31     ///
  32     FunctionExtractorPass(Function *F = 0, bool deleteFn = true,
  33                           bool relinkCallees = false)
  34       : Named(F), deleteFunc(deleteFn), reLink(relinkCallees) {}
  35
  36     bool runOnModule(Module &M) {
  37       if (Named == 0) {
  38         Named = M.getFunction("main");
  39         if (Named == 0) return false;  // No function to extract
  40       }
  41
  42       if (deleteFunc)
  43         return deleteFunction();
  44       M.setModuleInlineAsm("");
  45       return isolateFunction(M);
  46     }
  47
  48     bool deleteFunction() {
  49       // If we're in relinking mode, set linkage of all internal callees to
  50       // external. This will allow us extract function, and then - link
  51       // everything together
  52       if (reLink) {
  53         for (Function::iterator B = Named->begin(), BE = Named->end();
  54              B != BE; ++B) {
  55           for (BasicBlock::iterator I = B->begin(), E = B->end();
  56                I != E; ++I) {
  57             if (CallInst* callInst = dyn_cast<CallInst>(&*I)) {
  58               Function* Callee = callInst->getCalledFunction();
  59               if (Callee && Callee->hasInternalLinkage())
  60                 Callee->setLinkage(GlobalValue::ExternalLinkage);
  61             }
  62           }
  63         }
  64       }
  65
  66       Named->setLinkage(GlobalValue::ExternalLinkage);
  67       Named->deleteBody();
  68       assert(Named->isDeclaration() && "This didn't make the function external!");
  69       return true;
  70     }
  71
  72     bool isolateFunction(Module &M) {
  73       // Make sure our result is globally accessible...
  74       Named->setLinkage(GlobalValue::ExternalLinkage);
  75
  76       // Mark all global variables internal
  77       for (Module::global_iterator I = M.global_begin(), E = M.global_end(); I != E; ++I)
  78         if (!I->isDeclaration()) {
  79           I->setInitializer(0);  // Make all variables external
  80           I->setLinkage(GlobalValue::ExternalLinkage);
  81         }
  82
  83       // All of the functions may be used by global variables or the named
  84       // function.  Loop through them and create a new, external functions that
  85       // can be "used", instead of ones with bodies.
  86       std::vector<Function*> NewFunctions;
  87
  88       Function *Last = --M.end();  // Figure out where the last real fn is.
  89
  90       for (Module::iterator I = M.begin(); ; ++I) {
  91         if (&*I != Named) {
  92           Function *New = new Function(I->getFunctionType(),
  93                                        GlobalValue::ExternalLinkage,
  94                                        I->getName());
  95           New->setCallingConv(I->getCallingConv());
  96           I->setName("");  // Remove Old name
  97
  98           // If it's not the named function, delete the body of the function
  99           I->dropAllReferences();
 100
 101           M.getFunctionList().push_back(New);
 102           NewFunctions.push_back(New);
 103         }
 104
 105         if (&*I == Last) break;  // Stop after processing the last function
 106       }
 107
 108       // Now that we have replacements all set up, loop through the module,
 109       // deleting the old functions, replacing them with the newly created
 110       // functions.
 111       if (!NewFunctions.empty()) {
 112         unsigned FuncNum = 0;
 113         Module::iterator I = M.begin();
 114         do {
 115           if (&*I != Named) {
 116             // Make everything that uses the old function use the new dummy fn
 117             I->replaceAllUsesWith(NewFunctions[FuncNum++]);
 118
 119             Function *Old = I;
 120             ++I;  // Move the iterator to the new function
 121
 122             // Delete the old function!
 123             M.getFunctionList().erase(Old);
 124
 125           } else {
 126             ++I;  // Skip the function we are extracting
 127           }
 128         } while (&*I != NewFunctions[0]);
 129       }
 130
 131       return true;
 132     }
 133   };
 134
 135   RegisterPass<FunctionExtractorPass> X("extract", "Function Extractor");
 136 }
 137
 138 ModulePass *llvm::createFunctionExtractionPass(Function *F, bool deleteFn,
 139                                                bool relinkCallees) {
 140   return new FunctionExtractorPass(F, deleteFn, relinkCallees);
 141 }