1 files changed, 38 insertions, 26 deletions

diff --git a/include/llvm/Transforms/Utils/CodeExtractor.h b/include/llvm/Transforms/Utils/CodeExtractor.h
index 682b353ab5ae8..63d34511102dc 100644
--- a/include/llvm/Transforms/Utils/CodeExtractor.h
+++ b/include/llvm/Transforms/Utils/CodeExtractor.h
@@ -1,4 +1,4 @@
-//===-- Transform/Utils/CodeExtractor.h - Code extraction util --*- C++ -*-===//
+//===- Transform/Utils/CodeExtractor.h - Code extraction util ---*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -15,22 +15,24 @@
 #ifndef LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 #define LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
+#include <limits>
 
 namespace llvm {
-template <typename T> class ArrayRef;
-  class BasicBlock;
-  class BlockFrequency;
-  class BlockFrequencyInfo;
-  class BranchProbabilityInfo;
-  class DominatorTree;
-  class Function;
-  class Instruction;
-  class Loop;
-  class Module;
-  class RegionNode;
-  class Type;
-  class Value;
+
+class BasicBlock;
+class BlockFrequency;
+class BlockFrequencyInfo;
+class BranchProbabilityInfo;
+class DominatorTree;
+class Function;
+class Instruction;
+class Loop;
+class Module;
+class Type;
+class Value;
 
   /// \brief Utility class for extracting code into a new function.
   ///
@@ -46,7 +48,7 @@ template <typename T> class ArrayRef;
   /// 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;
+    using ValueSet = SetVector<Value *>;
 
     // Various bits of state computed on construction.
     DominatorTree *const DT;
@@ -54,27 +56,27 @@ template <typename T> class ArrayRef;
     BlockFrequencyInfo *BFI;
     BranchProbabilityInfo *BPI;
 
+    // If true, varargs functions can be extracted.
+    bool AllowVarArgs;
+
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
-    unsigned NumExitBlocks;
+    unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
     Type *RetTy;
 
   public:
-
-    /// \brief Check to see if a block is valid for extraction.
-    ///
-    /// Blocks containing EHPads, allocas, invokes, or vastarts are not valid.
-    static bool isBlockValidForExtraction(const BasicBlock &BB);
-
     /// \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.
+    /// extra checking and transformations are enabled. If AllowVarArgs is true,
+    /// vararg functions can be extracted. This is safe, if all vararg handling
+    /// code is extracted, including vastart.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
-                  BranchProbabilityInfo *BPI = nullptr);
+                  BranchProbabilityInfo *BPI = nullptr,
+                  bool AllowVarArgs = false);
 
     /// \brief Create a code extractor for a loop body.
     ///
@@ -84,6 +86,14 @@ template <typename T> class ArrayRef;
                   BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr);
 
+    /// \brief Check to see if a block is valid for extraction.
+    ///
+    /// Blocks containing EHPads, allocas and invokes are not valid. If
+    /// AllowVarArgs is true, blocks with vastart can be extracted. This is
+    /// safe, if all vararg handling code is extracted, including vastart.
+    static bool isBlockValidForExtraction(const BasicBlock &BB,
+                                          bool AllowVarArgs);
+
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
@@ -112,6 +122,7 @@ template <typename T> class ArrayRef;
     ///
     /// Returns true if it is safe to do the code motion.
     bool isLegalToShrinkwrapLifetimeMarkers(Instruction *AllocaAddr) const;
+
     /// Find the set of allocas whose life ranges are contained within the
     /// outlined region.
     ///
@@ -155,6 +166,7 @@ template <typename T> class ArrayRef;
                                     ValueSet &inputs,
                                     ValueSet &outputs);
   };
-}
 
-#endif
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H