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/Module.h"
  15 #include "llvm/Pass.h"
  16 #include "llvm/Transforms/IPO.h"
  17 using namespace llvm;
  18
  19 namespace {
  20   class FunctionExtractorPass : public ModulePass {
  21     Function *Named;
  22     bool deleteFunc;
  23   public:
  24     /// FunctionExtractorPass - If deleteFn is true, this pass deletes as the
  25     /// specified function. Otherwise, it deletes as much of the module as
  26     /// possible, except for the function specified.
  27     ///
  28     FunctionExtractorPass(Function *F = 0, bool deleteFn = true)
  29       : Named(F), deleteFunc(deleteFn) {}
  30
  31     bool runOnModule(Module &M) {
  32       if (Named == 0) {
  33         Named = M.getMainFunction();
  34         if (Named == 0) return false;  // No function to extract
  35       }
  36
  37       if (deleteFunc)
  38         return deleteFunction();
  39       M.setModuleInlineAsm("");
  40       return isolateFunction(M);
  41     }
  42
  43     bool deleteFunction() {
  44       Named->setLinkage(GlobalValue::ExternalLinkage);
  45       Named->deleteBody();
  46       assert(Named->isExternal() && "This didn't make the function external!");
  47       return true;
  48     }
  49
  50     bool isolateFunction(Module &M) {
  51       // Make sure our result is globally accessible...
  52       Named->setLinkage(GlobalValue::ExternalLinkage);
  53
  54       // Mark all global variables internal
  55       for (Module::global_iterator I = M.global_begin(), E = M.global_end(); I != E; ++I)
  56         if (!I->isExternal()) {
  57           I->setInitializer(0);  // Make all variables external
  58           I->setLinkage(GlobalValue::ExternalLinkage);
  59         }
  60
  61       // All of the functions may be used by global variables or the named
  62       // function.  Loop through them and create a new, external functions that
  63       // can be "used", instead of ones with bodies.
  64       std::vector<Function*> NewFunctions;
  65
  66       Function *Last = --M.end();  // Figure out where the last real fn is.
  67
  68       for (Module::iterator I = M.begin(); ; ++I) {
  69         if (&*I != Named) {
  70           Function *New = new Function(I->getFunctionType(),
  71                                        GlobalValue::ExternalLinkage,
  72                                        I->getName());
  73           New->setCallingConv(I->getCallingConv());
  74           I->setName("");  // Remove Old name
  75
  76           // If it's not the named function, delete the body of the function
  77           I->dropAllReferences();
  78
  79           M.getFunctionList().push_back(New);
  80           NewFunctions.push_back(New);
  81         }
  82
  83         if (&*I == Last) break;  // Stop after processing the last function
  84       }
  85
  86       // Now that we have replacements all set up, loop through the module,
  87       // deleting the old functions, replacing them with the newly created
  88       // functions.
  89       if (!NewFunctions.empty()) {
  90         unsigned FuncNum = 0;
  91         Module::iterator I = M.begin();
  92         do {
  93           if (&*I != Named) {
  94             // Make everything that uses the old function use the new dummy fn
  95             I->replaceAllUsesWith(NewFunctions[FuncNum++]);
  96
  97             Function *Old = I;
  98             ++I;  // Move the iterator to the new function
  99
 100             // Delete the old function!
 101             M.getFunctionList().erase(Old);
 102
 103           } else {
 104             ++I;  // Skip the function we are extracting
 105           }
 106         } while (&*I != NewFunctions[0]);
 107       }
 108
 109       return true;
 110     }
 111   };
 112
 113   RegisterPass<FunctionExtractorPass> X("extract", "Function Extractor");
 114 }
 115
 116 ModulePass *llvm::createFunctionExtractionPass(Function *F, bool deleteFn) {
 117   return new FunctionExtractorPass(F, deleteFn);
 118 }