Move the CodeExtractor utility to a dedicated header file / source file,

and expose it as a utility class rather than as free function wrappers.

The simple free-function interface works well for the bugpoint-specific
pass's uses of code extraction, but in an upcoming patch for more
advanced code extraction, they simply don't expose a rich enough
interface. I need to expose various stages of the process of doing the
code extraction and query information to decide whether or not to
actually complete the extraction or give up.

Rather than build up a new predicate model and pass that into these
functions, just take the class that was actually implementing the
functions and lift it up into a proper interface that can be used to
perform code extraction. The interface is cleaned up and re-documented
to work better in a header. It also is now setup to accept the blocks to
be extracted in the constructor rather than in a method.

In passing this essentially reverts my previous commit here exposing
a block-level query for eligibility of extraction. That is no longer
necessary with the more rich interface as clients can query the
extraction object for eligibility directly. This will reduce the number
of walks of the input basic block sequence by quite a bit which is
useful if this enters the normal optimization pipeline.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@156163 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 225 insertions, 220 deletions

diff --git a/include/llvm/Transforms/Utils/CodeExtractor.h b/include/llvm/Transforms/Utils/CodeExtractor.h
new file mode 100644
index 00000000000..7268a3c6b62
--- /dev/null
+++ b/include/llvm/Transforms/Utils/CodeExtractor.h
@@ -0,0 +1,110 @@
+//===-- Transform/Utils/CodeExtractor.h - Code extraction util --*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// A utility to support extracting code from one function into its own
+// stand-alone function.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_UTILS_CODE_EXTRACTOR_H
+#define LLVM_TRANSFORMS_UTILS_CODE_EXTRACTOR_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SetVector.h"
+
+namespace llvm {
+  class BasicBlock;
+  class DominatorTree;
+  class Function;
+  class Loop;
+  class Module;
+  class Type;
+  class Value;
+
+  /// \brief Utility class for extracting code into a new function.
+  ///
+  /// This utility provides a simple interface for extracting some sequence of
+  /// code into its own function, replacing it with a call to that function. It
+  /// also provides various methods to query about the nature and result of
+  /// such a transformation.
+  ///
+  /// The rough algorithm used is:
+  /// 1) Find both the inputs and outputs for the extracted region.
+  /// 2) Pass the inputs as arguments, remapping them within the extracted
+  ///    function to arguments.
+  /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
+  ///    as arguments, and inserting stores to the arguments for any scalars.
+  class CodeExtractor {
+    typedef SetVector<Value *> ValueSet;
+
+    // Various bits of state computed on construction.
+    DominatorTree *const DT;
+    const bool AggregateArgs;
+
+    // Bits of intermediate state computed at various phases of extraction.
+    SetVector<BasicBlock *> Blocks;
+    unsigned NumExitBlocks;
+    Type *RetTy;
+
+  public:
+    /// \brief Create a code extractor for a single basic block.
+    ///
+    /// In this formation, we don't require a dominator tree. The given basic
+    /// block is set up for extraction.
+    CodeExtractor(BasicBlock *BB, bool AggregateArgs = false);
+
+    /// \brief Create a code extractor for a sequence of blocks.
+    ///
+    /// Given a sequence of basic blocks where the first block in the sequence
+    /// dominates the rest, prepare a code extractor object for pulling this
+    /// sequence out into its new function. When a DominatorTree is also given,
+    /// extra checking and transformations are enabled.
+    CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = 0,
+                  bool AggregateArgs = false);
+
+    /// \brief Create a code extractor for a loop body.
+    ///
+    /// Behaves just like the generic code sequence constructor, but uses the
+    /// block sequence of the loop.
+    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false);
+
+    /// \brief Perform the extraction, returning the new function.
+    ///
+    /// Returns zero when called on a CodeExtractor instance where isEligible
+    /// returns false.
+    Function *extractCodeRegion();
+
+    /// \brief Test whether this code extractor is eligible.
+    ///
+    /// Based on the blocks used when constructing the code extractor,
+    /// determine whether it is eligible for extraction.
+    bool isEligible() { return !Blocks.empty(); };
+
+  private:
+    void severSplitPHINodes(BasicBlock *&Header);
+    void splitReturnBlocks();
+    void findInputsOutputs(ValueSet &inputs, ValueSet &outputs);
+
+    Function *constructFunction(const ValueSet &inputs,
+                                const ValueSet &outputs,
+                                BasicBlock *header,
+                                BasicBlock *newRootNode, BasicBlock *newHeader,
+                                Function *oldFunction, Module *M);
+
+    void moveCodeToFunction(Function *newFunction);
+
+    void emitCallAndSwitchStatement(Function *newFunction,
+                                    BasicBlock *newHeader,
+                                    ValueSet &inputs,
+                                    ValueSet &outputs);
+
+  };
+}
+
+#endif
diff --git a/include/llvm/Transforms/Utils/FunctionUtils.h b/include/llvm/Transforms/Utils/FunctionUtils.h
deleted file mode 100644
index 821a61e99b7..00000000000
--- a/include/llvm/Transforms/Utils/FunctionUtils.h
+++ /dev/null
@@ -1,53 +0,0 @@
-//===-- Transform/Utils/FunctionUtils.h - Function Utils --------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-//
-// This family of transformations manipulate LLVM functions.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef LLVM_TRANSFORMS_UTILS_FUNCTION_H
-#define LLVM_TRANSFORMS_UTILS_FUNCTION_H
-
-#include "llvm/ADT/ArrayRef.h"
-#include <vector>
-
-namespace llvm {
-  class BasicBlock;
-  class DominatorTree;
-  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.
-  ///
-  Function* ExtractCodeRegion(DominatorTree& DT,
-                              ArrayRef<BasicBlock*> code,
-                              bool AggregateArgs = false);
-
-  /// ExtractLoop - Rip out a natural loop into a new function.
-  ///
-  Function* ExtractLoop(DominatorTree& DT, Loop *L,
-                        bool AggregateArgs = false);
-
-  /// ExtractBasicBlock - Rip out a basic block (and the associated landing pad)
-  /// into a new function.
-  ///
-  Function* ExtractBasicBlock(ArrayRef<BasicBlock*> BBs,
-                              bool AggregateArgs = false);
-}
-
-#endif
diff --git a/lib/Transforms/IPO/LoopExtractor.cpp b/lib/Transforms/IPO/LoopExtractor.cpp
index 4f96afe44c0..97d7cdced0e 100644
--- a/lib/Transforms/IPO/LoopExtractor.cpp
+++ b/lib/Transforms/IPO/LoopExtractor.cpp
@@ -24,7 +24,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
-#include "llvm/Transforms/Utils/FunctionUtils.h"
+#include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/ADT/Statistic.h"
 #include <fstream>
 #include <set>
@@ -132,7 +132,8 @@ bool LoopExtractor::runOnLoop(Loop *L, LPPassManager &LPM) {
   if (ShouldExtractLoop) {
     if (NumLoops == 0) return Changed;
     --NumLoops;
-    if (ExtractLoop(DT, L) != 0) {
+    CodeExtractor Extractor(DT, *L);
+    if (Extractor.extractCodeRegion() != 0) {
       Changed = true;
       // After extraction, the loop is replaced by a function call, so
       // we shouldn't try to run any more loop passes on it.
@@ -296,7 +297,7 @@ bool BlockExtractorPass::runOnModule(Module &M) {
     if (const InvokeInst *II =
         dyn_cast<InvokeInst>(BlocksToExtract[i]->getTerminator()))
       BlocksToExtractVec.push_back(II->getUnwindDest());
-    ExtractBasicBlock(BlocksToExtractVec);
+    CodeExtractor(BlocksToExtractVec).extractCodeRegion();
   }
 
   return !BlocksToExtract.empty();
diff --git a/lib/Transforms/IPO/PartialInlining.cpp b/lib/Transforms/IPO/PartialInlining.cpp
index d9d1d106111..9c9910bd5cc 100644
--- a/lib/Transforms/IPO/PartialInlining.cpp
+++ b/lib/Transforms/IPO/PartialInlining.cpp
@@ -19,7 +19,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Analysis/Dominators.h"
 #include "llvm/Transforms/Utils/Cloning.h"
-#include "llvm/Transforms/Utils/FunctionUtils.h"
+#include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/CFG.h"
 using namespace llvm;
@@ -122,7 +122,8 @@ Function* PartialInliner::unswitchFunction(Function* F) {
   DT.runOnFunction(*duplicateFunction);
   
   // Extract the body of the if.
-  Function* extractedFunction = ExtractCodeRegion(DT, toExtract);
+  Function* extractedFunction
+    = CodeExtractor(toExtract, &DT).extractCodeRegion();
   
   InlineFunctionInfo IFI;
   
diff --git a/lib/Transforms/Utils/CodeExtractor.cpp b/lib/Transforms/Utils/CodeExtractor.cpp
index b8cea45178c..50eb8a27e07 100644
--- a/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/lib/Transforms/Utils/CodeExtractor.cpp
@@ -13,7 +13,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Transforms/Utils/FunctionUtils.h"
+#include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/Constants.h"
 #include "llvm/DerivedTypes.h"
 #include "llvm/Instructions.h"
@@ -43,61 +43,78 @@ static cl::opt<bool>
 AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
                  cl::desc("Aggregate arguments to code-extracted functions"));
 
-namespace {
-  class CodeExtractor {
-    typedef SetVector<Value*> Values;
-    SetVector<BasicBlock*> BlocksToExtract;
-    DominatorTree* DT;
-    bool AggregateArgs;
-    unsigned NumExitBlocks;
-    Type *RetTy;
-  public:
-    CodeExtractor(DominatorTree* dt = 0, bool AggArgs = false)
-      : DT(dt), AggregateArgs(AggArgs||AggregateArgsOpt), NumExitBlocks(~0U) {}
-
-    Function *ExtractCodeRegion(ArrayRef<BasicBlock*> code);
-
-    bool isEligible(ArrayRef<BasicBlock*> code);
-
-  private:
-    /// definedInRegion - Return true if the specified value is defined in the
-    /// extracted region.
-    bool definedInRegion(Value *V) const {
-      if (Instruction *I = dyn_cast<Instruction>(V))
-        if (BlocksToExtract.count(I->getParent()))
-          return true;
-      return false;
-    }
+/// \brief Test whether a block is valid for extraction.
+static bool isBlockValidForExtraction(const BasicBlock &BB) {
+  // Landing pads must be in the function where they were inserted for cleanup.
+  if (BB.isLandingPad())
+    return false;
 
-    /// definedInCaller - Return true if the specified value is defined in the
-    /// function being code extracted, but not in the region being extracted.
-    /// These values must be passed in as live-ins to the function.
-    bool definedInCaller(Value *V) const {
-      if (isa<Argument>(V)) return true;
-      if (Instruction *I = dyn_cast<Instruction>(V))
-        if (!BlocksToExtract.count(I->getParent()))
-          return true;
+  // 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;
+}
+
+/// \brief Build a set of blocks to extract if the input blocks are viable.
+static SetVector<BasicBlock *>
+buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs) {
+  SetVector<BasicBlock *> Result;
+
+  // Loop over the blocks, adding them to our set-vector, and aborting with an
+  // empty set if we encounter invalid blocks.
+  for (ArrayRef<BasicBlock *>::iterator I = BBs.begin(), E = BBs.end();
+       I != E; ++I) {
+    if (!Result.insert(*I))
+      continue;
+
+    if (!isBlockValidForExtraction(**I)) {
+      Result.clear();
+      break;
     }
+  }
 
-    void severSplitPHINodes(BasicBlock *&Header);
-    void splitReturnBlocks();
-    void findInputsOutputs(Values &inputs, Values &outputs);
+  return Result;
+}
+
+CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs)
+  : DT(0), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    Blocks(buildExtractionBlockSet(BB)), NumExitBlocks(~0U) {}
 
-    Function *constructFunction(const Values &inputs,
-                                const Values &outputs,
-                                BasicBlock *header,
-                                BasicBlock *newRootNode, BasicBlock *newHeader,
-                                Function *oldFunction, Module *M);
+CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
+                             bool AggregateArgs)
+  : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    Blocks(buildExtractionBlockSet(BBs)), NumExitBlocks(~0U) {}
 
-    void moveCodeToFunction(Function *newFunction);
+CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs)
+  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    Blocks(buildExtractionBlockSet(L.getBlocks())), NumExitBlocks(~0U) {}
 
-    void emitCallAndSwitchStatement(Function *newFunction,
-                                    BasicBlock *newHeader,
-                                    Values &inputs,
-                                    Values &outputs);
 
-  };
+/// definedInRegion - Return true if the specified value is defined in the
+/// extracted region.
+static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
+  if (Instruction *I = dyn_cast<Instruction>(V))
+    if (Blocks.count(I->getParent()))
+      return true;
+  return false;
+}
+
+/// definedInCaller - Return true if the specified value is defined in the
+/// function being code extracted, but not in the region being extracted.
+/// These values must be passed in as live-ins to the function.
+static bool definedInCaller(const SetVector<BasicBlock *> &Blocks, Value *V) {
+  if (isa<Argument>(V)) return true;
+  if (Instruction *I = dyn_cast<Instruction>(V))
+    if (!Blocks.count(I->getParent()))
+      return true;
+  return false;
 }
 
 /// severSplitPHINodes - If a PHI node has multiple inputs from outside of the
@@ -115,7 +132,7 @@ void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
     // than one entry from outside the region.  If so, we need to sever the
     // header block into two.
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
-      if (BlocksToExtract.count(PN->getIncomingBlock(i)))
+      if (Blocks.count(PN->getIncomingBlock(i)))
         ++NumPredsFromRegion;
       else
         ++NumPredsOutsideRegion;
@@ -136,8 +153,8 @@ void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
   // We only want to code extract the second block now, and it becomes the new
   // header of the region.
   BasicBlock *OldPred = Header;
-  BlocksToExtract.remove(OldPred);
-  BlocksToExtract.insert(NewBB);
+  Blocks.remove(OldPred);
+  Blocks.insert(NewBB);
   Header = NewBB;
 
   // Okay, update dominator sets. The blocks that dominate the new one are the
@@ -152,7 +169,7 @@ void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
     // Loop over all of the predecessors of OldPred that are in the region,
     // changing them to branch to NewBB instead.
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
-      if (BlocksToExtract.count(PN->getIncomingBlock(i))) {
+      if (Blocks.count(PN->getIncomingBlock(i))) {
         TerminatorInst *TI = PN->getIncomingBlock(i)->getTerminator();
         TI->replaceUsesOfWith(OldPred, NewBB);
       }
@@ -170,7 +187,7 @@ void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
       // Loop over all of the incoming value in PN, moving them to NewPN if they
       // are from the extracted region.
       for (unsigned i = 0; i != PN->getNumIncomingValues(); ++i) {
-        if (BlocksToExtract.count(PN->getIncomingBlock(i))) {
+        if (Blocks.count(PN->getIncomingBlock(i))) {
           NewPN->addIncoming(PN->getIncomingValue(i), PN->getIncomingBlock(i));
           PN->removeIncomingValue(i);
           --i;
@@ -181,8 +198,8 @@ void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
 }
 
 void CodeExtractor::splitReturnBlocks() {
-  for (SetVector<BasicBlock*>::iterator I = BlocksToExtract.begin(),
-         E = BlocksToExtract.end(); I != E; ++I)
+  for (SetVector<BasicBlock *>::iterator I = Blocks.begin(), E = Blocks.end();
+       I != E; ++I)
     if (ReturnInst *RI = dyn_cast<ReturnInst>((*I)->getTerminator())) {
       BasicBlock *New = (*I)->splitBasicBlock(RI, (*I)->getName()+".ret");
       if (DT) {
@@ -205,23 +222,23 @@ void CodeExtractor::splitReturnBlocks() {
 
 // findInputsOutputs - Find inputs to, outputs from the code region.
 //
-void CodeExtractor::findInputsOutputs(Values &inputs, Values &outputs) {
+void CodeExtractor::findInputsOutputs(ValueSet &inputs, ValueSet &outputs) {
   std::set<BasicBlock*> ExitBlocks;
-  for (SetVector<BasicBlock*>::const_iterator ci = BlocksToExtract.begin(),
-       ce = BlocksToExtract.end(); ci != ce; ++ci) {
+  for (SetVector<BasicBlock*>::const_iterator ci = Blocks.begin(),
+       ce = Blocks.end(); ci != ce; ++ci) {
     BasicBlock *BB = *ci;
 
     for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I) {
       // If a used value is defined outside the region, it's an input.  If an
       // instruction is used outside the region, it's an output.
       for (User::op_iterator O = I->op_begin(), E = I->op_end(); O != E; ++O)
-        if (definedInCaller(*O))
+        if (definedInCaller(Blocks, *O))
           inputs.insert(*O);
 
       // Consider uses of this instruction (outputs).
       for (Value::use_iterator UI = I->use_begin(), E = I->use_end();
            UI != E; ++UI)
-        if (!definedInRegion(*UI)) {
+        if (!definedInRegion(Blocks, *UI)) {
           outputs.insert(I);
           break;
         }
@@ -230,7 +247,7 @@ void CodeExtractor::findInputsOutputs(Values &inputs, Values &outputs) {
     // Keep track of the exit blocks from the region.
     TerminatorInst *TI = BB->getTerminator();
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
-      if (!BlocksToExtract.count(TI->getSuccessor(i)))
+      if (!Blocks.count(TI->getSuccessor(i)))
         ExitBlocks.insert(TI->getSuccessor(i));
   } // for: basic blocks
 
@@ -240,8 +257,8 @@ void CodeExtractor::findInputsOutputs(Values &inputs, Values &outputs) {
 /// constructFunction - make a function based on inputs and outputs, as follows:
 /// f(in0, ..., inN, out0, ..., outN)
 ///
-Function *CodeExtractor::constructFunction(const Values &inputs,
-                                           const Values &outputs,
+Function *CodeExtractor::constructFunction(const ValueSet &inputs,
+                                           const ValueSet &outputs,
                                            BasicBlock *header,
                                            BasicBlock *newRootNode,
                                            BasicBlock *newHeader,
@@ -261,15 +278,15 @@ Function *CodeExtractor::constructFunction(const Values &inputs,
   std::vector<Type*> paramTy;
 
   // Add the types of the input values to the function's argument list
-  for (Values::const_iterator i = inputs.begin(),
-         e = inputs.end(); i != e; ++i) {
+  for (ValueSet::const_iterator i = inputs.begin(), e = inputs.end();
+       i != e; ++i) {
     const Value *value = *i;
     DEBUG(dbgs() << "value used in func: " << *value << "\n");
     paramTy.push_back(value->getType());
   }
 
   // Add the types of the output values to the function's argument list.
-  for (Values::const_iterator I = outputs.begin(), E = outputs.end();
+  for (ValueSet::const_iterator I = outputs.begin(), E = outputs.end();
        I != E; ++I) {
     DEBUG(dbgs() << "instr used in func: " << **I << "\n");
     if (AggregateArgs)
@@ -326,7 +343,7 @@ Function *CodeExtractor::constructFunction(const Values &inputs,
     for (std::vector<User*>::iterator use = Users.begin(), useE = Users.end();
          use != useE; ++use)
       if (Instruction* inst = dyn_cast<Instruction>(*use))
-        if (BlocksToExtract.count(inst->getParent()))
+        if (Blocks.count(inst->getParent()))
           inst->replaceUsesOfWith(inputs[i], RewriteVal);
   }
 
@@ -347,7 +364,7 @@ Function *CodeExtractor::constructFunction(const Values &inputs,
     // The BasicBlock which contains the branch is not in the region
     // modify the branch target to a new block
     if (TerminatorInst *TI = dyn_cast<TerminatorInst>(Users[i]))
-      if (!BlocksToExtract.count(TI->getParent()) &&
+      if (!Blocks.count(TI->getParent()) &&
           TI->getParent()->getParent() == oldFunction)
         TI->replaceUsesOfWith(header, newHeader);
 
@@ -373,7 +390,7 @@ static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {
 /// necessary.
 void CodeExtractor::
 emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
-                           Values &inputs, Values &outputs) {
+                           ValueSet &inputs, ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
   std::vector<Value*> params, StructValues, ReloadOutputs, Reloads;
@@ -381,14 +398,14 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
   LLVMContext &Context = newFunction->getContext();
 
   // Add inputs as params, or to be filled into the struct
-  for (Values::iterator i = inputs.begin(), e = inputs.end(); i != e; ++i)
+  for (ValueSet::iterator i = inputs.begin(), e = inputs.end(); i != e; ++i)
     if (AggregateArgs)
       StructValues.push_back(*i);
     else
       params.push_back(*i);
 
   // Create allocas for the outputs
-  for (Values::iterator i = outputs.begin(), e = outputs.end(); i != e; ++i) {
+  for (ValueSet::iterator i = outputs.begin(), e = outputs.end(); i != e; ++i) {
     if (AggregateArgs) {
       StructValues.push_back(*i);
     } else {
@@ -403,7 +420,7 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
   AllocaInst *Struct = 0;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     std::vector<Type*> ArgTypes;
-    for (Values::iterator v = StructValues.begin(),
+    for (ValueSet::iterator v = StructValues.begin(),
            ve = StructValues.end(); v != ve; ++v)
       ArgTypes.push_back((*v)->getType());
 
@@ -458,7 +475,7 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
     std::vector<User*> Users(outputs[i]->use_begin(), outputs[i]->use_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
-      if (!BlocksToExtract.count(inst->getParent()))
+      if (!Blocks.count(inst->getParent()))
         inst->replaceUsesOfWith(outputs[i], load);
     }
   }
@@ -476,11 +493,11 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
   std::map<BasicBlock*, BasicBlock*> ExitBlockMap;
 
   unsigned switchVal = 0;
-  for (SetVector<BasicBlock*>::const_iterator i = BlocksToExtract.begin(),
-         e = BlocksToExtract.end(); i != e; ++i) {
+  for (SetVector<BasicBlock*>::const_iterator i = Blocks.begin(),
+         e = Blocks.end(); i != e; ++i) {
     TerminatorInst *TI = (*i)->getTerminator();
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
-      if (!BlocksToExtract.count(TI->getSuccessor(i))) {
+      if (!Blocks.count(TI->getSuccessor(i))) {
         BasicBlock *OldTarget = TI->getSuccessor(i);
         // add a new basic block which returns the appropriate value
         BasicBlock *&NewTarget = ExitBlockMap[OldTarget];
@@ -624,12 +641,12 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
 }
 
 void CodeExtractor::moveCodeToFunction(Function *newFunction) {
-  Function *oldFunc = (*BlocksToExtract.begin())->getParent();
+  Function *oldFunc = (*Blocks.begin())->getParent();
   Function::BasicBlockListType &oldBlocks = oldFunc->getBasicBlockList();
   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
 
-  for (SetVector<BasicBlock*>::const_iterator i = BlocksToExtract.begin(),
-         e = BlocksToExtract.end(); i != e; ++i) {
+  for (SetVector<BasicBlock*>::const_iterator i = Blocks.begin(),
+         e = Blocks.end(); i != e; ++i) {
     // Delete the basic block from the old function, and the list of blocks
     oldBlocks.remove(*i);
 
@@ -638,45 +655,22 @@ void CodeExtractor::moveCodeToFunction(Function *newFunction) {
   }
 }
 
-/// ExtractRegion - Removes a loop from a function, replaces it with a call to
-/// new function. Returns pointer to the new function.
-///
-/// algorithm:
-///
-/// find inputs and outputs for the region
-///
-/// for inputs: add to function as args, map input instr* to arg#
-/// for outputs: add allocas for scalars,
-///             add to func as args, map output instr* to arg#
-///
-/// rewrite func to use argument #s instead of instr*
-///
-/// for each scalar output in the function: at every exit, store intermediate
-/// computed result back into memory.
-///
-Function *CodeExtractor::
-ExtractCodeRegion(ArrayRef<BasicBlock*> code) {
-  if (!isEligible(code))
+Function *CodeExtractor::extractCodeRegion() {
+  if (!isEligible())
     return 0;
 
-  // 1) Find inputs, outputs
-  // 2) Construct new function
-  //  * Add allocas for defs, pass as args by reference
-  //  * Pass in uses as args
-  // 3) Move code region, add call instr to func
-  //
-  BlocksToExtract.insert(code.begin(), code.end());
-
-  Values inputs, outputs;
+  ValueSet inputs, outputs;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
-  BasicBlock *header = code[0];
+  BasicBlock *header = *Blocks.begin();
 
-  for (unsigned i = 1, e = code.size(); i != e; ++i)
-    for (pred_iterator PI = pred_begin(code[i]), E = pred_end(code[i]);
+  for (SetVector<BasicBlock *>::iterator BI = llvm::next(Blocks.begin()),
+                                         BE = Blocks.end();
+       BI != BE; ++BI)
+    for (pred_iterator PI = pred_begin(*BI), E = pred_end(*BI);
          PI != E; ++PI)
-      assert(BlocksToExtract.count(*PI) &&
+      assert(Blocks.count(*PI) &&
              "No blocks in this region may have entries from outside the region"
              " except for the first block!");
 
@@ -718,7 +712,7 @@ ExtractCodeRegion(ArrayRef<BasicBlock*> code) {
   for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PN = cast<PHINode>(I);
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
-      if (!BlocksToExtract.count(PN->getIncomingBlock(i)))
+      if (!Blocks.count(PN->getIncomingBlock(i)))
         PN->setIncomingBlock(i, newFuncRoot);
   }
 
@@ -732,7 +726,7 @@ ExtractCodeRegion(ArrayRef<BasicBlock*> code) {
       PHINode *PN = cast<PHINode>(I);
       std::set<BasicBlock*> ProcessedPreds;
       for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
-        if (BlocksToExtract.count(PN->getIncomingBlock(i))) {
+        if (Blocks.count(PN->getIncomingBlock(i))) {
           if (ProcessedPreds.insert(PN->getIncomingBlock(i)).second)
             PN->setIncomingBlock(i, codeReplacer);
           else {
@@ -754,51 +748,3 @@ ExtractCodeRegion(ArrayRef<BasicBlock*> code) {
         report_fatal_error("verifyFunction failed!"));
   return newFunction;
 }
-
-bool CodeExtractor::isEligible(ArrayRef<BasicBlock*> code) {
-  for (ArrayRef<BasicBlock*>::iterator I = code.begin(), E = code.end();
-       I != E; ++I)
-    if (!isBlockViableForExtraction(**I))
-      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.
-///
-Function* llvm::ExtractCodeRegion(DominatorTree &DT,
-                                  ArrayRef<BasicBlock*> code,
-                                  bool AggregateArgs) {
-  return CodeExtractor(&DT, AggregateArgs).ExtractCodeRegion(code);
-}
-
-/// ExtractLoop - Slurp a natural loop into a brand new function.
-///
-Function* llvm::ExtractLoop(DominatorTree &DT, Loop *L, bool AggregateArgs) {
-  return CodeExtractor(&DT, AggregateArgs).ExtractCodeRegion(L->getBlocks());
-}
-
-/// ExtractBasicBlock - Slurp a basic block into a brand new function.
-///
-Function* llvm::ExtractBasicBlock(ArrayRef<BasicBlock*> BBs, bool AggregateArgs){
-  return CodeExtractor(0, AggregateArgs).ExtractCodeRegion(BBs);
-}
diff --git a/tools/bugpoint/ExtractFunction.cpp b/tools/bugpoint/ExtractFunction.cpp
index ac8e1597115..888d2c8e926 100644
--- a/tools/bugpoint/ExtractFunction.cpp
+++ b/tools/bugpoint/ExtractFunction.cpp
@@ -24,7 +24,7 @@
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Cloning.h"
-#include "llvm/Transforms/Utils/FunctionUtils.h"
+#include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/Target/TargetData.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"