2 files changed, 29 insertions, 14 deletions
diff --git a/include/llvm/Transforms/Utils/FunctionUtils.h b/include/llvm/Transforms/Utils/FunctionUtils.h
index 8d71e43aa9..821a61e99b 100644
--- a/include/llvm/Transforms/Utils/FunctionUtils.h
+++ b/include/llvm/Transforms/Utils/FunctionUtils.h
@@ -23,6 +23,14 @@ namespace llvm {
   class Function;
   class Loop;
 
+  /// \brief Test whether a basic block is a viable candidate for extraction.
+  ///
+  /// This tests whether a particular basic block is viable for extraction into
+  /// a separate function. It can be used to prune code extraction eagerly
+  /// rather than waiting for one of the Extract* methods below to reject
+  /// a request.
+  bool isBlockViableForExtraction(const BasicBlock &BB);
+
   /// ExtractCodeRegion - Rip out a sequence of basic blocks into a new
   /// function.
   ///
diff --git a/lib/Transforms/Utils/CodeExtractor.cpp b/lib/Transforms/Utils/CodeExtractor.cpp
index e8c0b80c21..b8cea45178 100644
--- a/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/lib/Transforms/Utils/CodeExtractor.cpp
@@ -756,24 +756,31 @@ ExtractCodeRegion(ArrayRef<BasicBlock*> code) {
 }
 
 bool CodeExtractor::isEligible(ArrayRef<BasicBlock*> code) {
-  // Deny a single basic block that's a landing pad block.
-  if (code.size() == 1 && code[0]->isLandingPad())
-    return false;
+  for (ArrayRef<BasicBlock*>::iterator I = code.begin(), E = code.end();
+       I != E; ++I)
+    if (!isBlockViableForExtraction(**I))
+      return false;
 
-  // Deny code region if it contains allocas or vastarts.
-  for (ArrayRef<BasicBlock*>::iterator BB = code.begin(), e=code.end();
-       BB != e; ++BB)
-    for (BasicBlock::const_iterator I = (*BB)->begin(), Ie = (*BB)->end();
-         I != Ie; ++I)
-      if (isa<AllocaInst>(*I))
-        return false;
-      else if (const CallInst *CI = dyn_cast<CallInst>(I))
-        if (const Function *F = CI->getCalledFunction())
-          if (F->getIntrinsicID() == Intrinsic::vastart)
-            return false;
   return true;
 }
 
+bool llvm::isBlockViableForExtraction(const BasicBlock &BB) {
+  // Landing pads must be in the function where they were inserted for cleanup.
+  if (BB.isLandingPad())
+    return false;
+
+  // Don't hoist code containing allocas, invokes, or vastarts.
+  for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
+    if (isa<AllocaInst>(I) || isa<InvokeInst>(I))
+      return false;
+    if (const CallInst *CI = dyn_cast<CallInst>(I))
+      if (const Function *F = CI->getCalledFunction())
+        if (F->getIntrinsicID() == Intrinsic::vastart)
+          return false;
+  }
+
+  return true;
+}
 
 /// ExtractCodeRegion - Slurp a sequence of basic blocks into a brand new
 /// function.