1 files changed, 25 insertions, 4 deletions

diff --git a/lib/Transforms/IPO/ExtractFunction.cpp b/lib/Transforms/IPO/ExtractFunction.cpp
index a871c4d7f2..3246cd96bb 100644
--- a/lib/Transforms/IPO/ExtractFunction.cpp
+++ b/lib/Transforms/IPO/ExtractFunction.cpp
@@ -11,6 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Instructions.h"
 #include "llvm/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/IPO.h"
@@ -20,13 +21,15 @@ namespace {
   class FunctionExtractorPass : public ModulePass {
     Function *Named;
     bool deleteFunc;
+    bool reLink;
   public:
     /// FunctionExtractorPass - If deleteFn is true, this pass deletes as the
     /// specified function. Otherwise, it deletes as much of the module as
     /// possible, except for the function specified.
     ///
-    FunctionExtractorPass(Function *F = 0, bool deleteFn = true)
-      : Named(F), deleteFunc(deleteFn) {}
+    FunctionExtractorPass(Function *F = 0, bool deleteFn = true,
+                          bool relinkCallees = false)
+      : Named(F), deleteFunc(deleteFn), reLink(relinkCallees) {}
 
     bool runOnModule(Module &M) {
       if (Named == 0) {
@@ -41,6 +44,23 @@ namespace {
     }
 
     bool deleteFunction() {
+      // If we're in relinking mode, set linkage of all internal callees to
+      // external. This will allow us extract function, and then - link
+      // everything together
+      if (reLink) {
+        for (Function::iterator B = Named->begin(), BE = Named->end();
+             B != BE; ++B) {
+          for (BasicBlock::iterator I = B->begin(), E = B->end();
+               I != E; ++I) {
+            if (CallInst* callInst = dyn_cast<CallInst>(&*I)) {
+              Function* Callee = callInst->getCalledFunction();
+              if (Callee && Callee->hasInternalLinkage())
+                Callee->setLinkage(GlobalValue::ExternalLinkage);
+            }
+          }
+        }
+      }
+      
       Named->setLinkage(GlobalValue::ExternalLinkage);
       Named->deleteBody();
       assert(Named->isExternal() && "This didn't make the function external!");
@@ -113,6 +133,7 @@ namespace {
   RegisterPass<FunctionExtractorPass> X("extract", "Function Extractor");
 }
 
-ModulePass *llvm::createFunctionExtractionPass(Function *F, bool deleteFn) {
-  return new FunctionExtractorPass(F, deleteFn);
+ModulePass *llvm::createFunctionExtractionPass(Function *F, bool deleteFn,
+                                               bool relinkCallees) {
+  return new FunctionExtractorPass(F, deleteFn, relinkCallees);
 }