2 files changed, 29 insertions, 7 deletions
diff --git a/lib/Transforms/IPO/DeadArgumentElimination.cpp b/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 840533f968..d83f3e027d 100644
--- a/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -230,9 +230,10 @@ static inline bool CallPassesValueThoughVararg(Instruction *Call,
 // (used in a computation), MaybeLive (only passed as an argument to a call), or
 // Dead (not used).
 DAE::Liveness DAE::getArgumentLiveness(const Argument &A) {
-  // If this is the return value of a csret function, it's not really dead.
-  if (A.getParent()->getCallingConv() == CallingConv::CSRet &&
-      &*A.getParent()->arg_begin() == &A)
+  const FunctionType *FTy = A.getParent()->getFunctionType();
+  
+  // If this is the return value of a struct function, it's not really dead.
+  if (FTy->isStructReturn() && &*A.getParent()->arg_begin() == &A)
     return Live;
   
   if (A.use_empty())  // First check, directly dead?
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);
 }