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 Pass {
  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 run(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       else
  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::giterator I = M.gbegin(), E = M.gend(); 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.back();  // 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           I->setName("");  // Remove Old name
  74
  75           // If it's not the named function, delete the body of the function
  76           I->dropAllReferences();
  77
  78           M.getFunctionList().push_back(New);
  79           NewFunctions.push_back(New);
  80         }
  81
  82         if (&*I == Last) break;  // Stop after processing the last function
  83       }
  84
  85       // Now that we have replacements all set up, loop through the module,
  86       // deleting the old functions, replacing them with the newly created
  87       // functions.
  88       if (!NewFunctions.empty()) {
  89         unsigned FuncNum = 0;
  90         Module::iterator I = M.begin();
  91         do {
  92           if (&*I != Named) {
  93             // Make everything that uses the old function use the new dummy fn
  94             I->replaceAllUsesWith(NewFunctions[FuncNum++]);
  95
  96             Function *Old = I;
  97             ++I;  // Move the iterator to the new function
  98
  99             // Delete the old function!
 100             M.getFunctionList().erase(Old);
 101
 102           } else {
 103             ++I;  // Skip the function we are extracting
 104           }
 105         } while (&*I != NewFunctions[0]);
 106       }
 107
 108       return true;
 109     }
 110   };
 111
 112   RegisterPass<FunctionExtractorPass> X("extract", "Function Extractor");
 113 }
 114
 115 Pass *llvm::createFunctionExtractionPass(Function *F, bool deleteFn) {
 116   return new FunctionExtractorPass(F, deleteFn);
 117 }