lib/Transforms/IPO/ExtractFunction.cpp

   1
   2 #include "llvm/Transforms/IPO.h"
   3 #include "llvm/Pass.h"
   4 #include "llvm/Module.h"
   5
   6 namespace {
   7   class FunctionExtractorPass : public Pass {
   8     Function *Named;
   9   public:
  10     FunctionExtractorPass(Function *F = 0) : Named(F) {}
  11
  12     bool run(Module &M) {
  13       if (Named == 0) {
  14         Named = M.getMainFunction();
  15         if (Named == 0) return false;  // No function to extract
  16       }
  17
  18       // Make sure our result is globally accessable...
  19       Named->setInternalLinkage(false);
  20
  21       // Mark all global variables internal
  22       for (Module::giterator I = M.gbegin(), E = M.gend(); I != E; ++I)
  23         if (!I->isExternal()) {
  24           I->setInitializer(0);  // Make all variables external
  25           I->setInternalLinkage(false); // Make sure it's not internal
  26         }
  27
  28       // All of the functions may be used by global variables or the named
  29       // function.  Loop through them and create a new, external functions that
  30       // can be "used", instead of ones with bodies.
  31       //
  32       std::vector<Function*> NewFunctions;
  33
  34       Function *Last = &M.back();  // Figure out where the last real fn is...
  35
  36       for (Module::iterator I = M.begin(); ; ++I) {
  37         if (&*I != Named) {
  38           Function *New = new Function(I->getFunctionType(),false,I->getName());
  39           I->setName("");  // Remove Old name
  40
  41           // If it's not the named function, delete the body of the function
  42           I->dropAllReferences();
  43
  44           M.getFunctionList().push_back(New);
  45           NewFunctions.push_back(New);
  46         }
  47
  48         if (&*I == Last) break;  // Stop after processing the last function
  49       }
  50
  51       // Now that we have replacements all set up, loop through the module,
  52       // deleting the old functions, replacing them with the newly created
  53       // functions.
  54       if (!NewFunctions.empty()) {
  55         unsigned FuncNum = 0;
  56         Module::iterator I = M.begin();
  57         do {
  58           if (&*I != Named) {
  59             // Make everything that uses the old function use the new dummy fn
  60             I->replaceAllUsesWith(NewFunctions[FuncNum++]);
  61
  62             Function *Old = I;
  63             ++I;  // Move the iterator to the new function
  64
  65             // Delete the old function!
  66             M.getFunctionList().erase(Old);
  67
  68           } else {
  69             ++I;  // Skip the function we are extracting
  70           }
  71         } while (&*I != NewFunctions[0]);
  72       }
  73
  74       return true;
  75     }
  76   };
  77
  78   RegisterPass<FunctionExtractorPass> X("extract", "Function Extractor");
  79 }
  80
  81 Pass *createFunctionExtractionPass(Function *F) {
  82   return new FunctionExtractorPass(F);
  83 }