2 files changed, 74 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/PriorityWorklist.h b/llvm/include/llvm/ADT/PriorityWorklist.h
index c0b4709e98f..ef7b8930ab3 100644
--- a/llvm/include/llvm/ADT/PriorityWorklist.h
+++ b/llvm/include/llvm/ADT/PriorityWorklist.h
@@ -18,6 +18,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Compiler.h"
 #include <algorithm>
@@ -107,6 +108,35 @@ public:
     return false;
   }
 
+  /// Insert a sequence of new elements into the PriorityWorklist.
+  template <typename SequenceT>
+  typename std::enable_if<!std::is_convertible<SequenceT, T>::value>::type
+  insert(SequenceT &&Input) {
+    // First pull the input sequence into the vector as a bulk append
+    // operation.
+    ptrdiff_t StartIndex = V.size();
+    V.insert(V.end(), std::begin(Input), std::end(Input));
+    // Now walk backwards fixing up the index map and deleting any duplicates.
+    for (auto i : reverse(seq<ptrdiff_t>(StartIndex, V.size()))) {
+      auto InsertResult = M.insert({V[i], i});
+      if (InsertResult.second)
+        continue;
+
+      // If the existing index is before this insert's start, nuke that one and
+      // move it up.
+      ptrdiff_t &Index = InsertResult.first->second;
+      if (Index < StartIndex) {
+        V[Index] = T();
+        Index = i;
+        continue;
+      }
+
+      // Otherwise the existing one comes first so just clear out the value in
+      // this slot.
+      V[i] = T();
+    }
+  }
+
   /// Remove the last element of the PriorityWorklist.
   void pop_back() {
     assert(!empty() && "Cannot remove an element when empty!");
@@ -169,6 +199,11 @@ public:
     return true;
   }
 
+  /// Reverse the items in the PriorityWorklist.
+  ///
+  /// This does an in-place reversal. Other kinds of reverse aren't easy to
+  /// support in the face of the worklist semantics.
+
   /// Completely clear the PriorityWorklist
   void clear() {
     M.clear();
diff --git a/llvm/unittests/ADT/PriorityWorklistTest.cpp b/llvm/unittests/ADT/PriorityWorklistTest.cpp
index bbe026434c6..aa7a5430e21 100644
--- a/llvm/unittests/ADT/PriorityWorklistTest.cpp
+++ b/llvm/unittests/ADT/PriorityWorklistTest.cpp
@@ -13,6 +13,8 @@
 
 #include "llvm/ADT/PriorityWorklist.h"
 #include "gtest/gtest.h"
+#include <list>
+#include <vector>
 
 namespace {
 
@@ -72,6 +74,43 @@ TYPED_TEST(PriorityWorklistTest, Basic) {
   EXPECT_TRUE(W.empty());
 }
 
+TYPED_TEST(PriorityWorklistTest, InsertSequence) {
+  TypeParam W;
+  ASSERT_TRUE(W.insert(2));
+  ASSERT_TRUE(W.insert(4));
+  ASSERT_TRUE(W.insert(7));
+  // Insert a sequence that has internal duplicates and a duplicate among
+  // existing entries.
+  W.insert(std::vector<int>({42, 13, 42, 7, 8}));
+  EXPECT_EQ(8, W.pop_back_val());
+  EXPECT_EQ(7, W.pop_back_val());
+  EXPECT_EQ(42, W.pop_back_val());
+  EXPECT_EQ(13, W.pop_back_val());
+  EXPECT_EQ(4, W.pop_back_val());
+  EXPECT_EQ(2, W.pop_back_val());
+  ASSERT_TRUE(W.empty());
+
+  // Simpler tests with various other input types.
+  ASSERT_TRUE(W.insert(2));
+  ASSERT_TRUE(W.insert(7));
+  // Use a non-random-access container.
+  W.insert(std::list<int>({7, 5}));
+  EXPECT_EQ(5, W.pop_back_val());
+  EXPECT_EQ(7, W.pop_back_val());
+  EXPECT_EQ(2, W.pop_back_val());
+  ASSERT_TRUE(W.empty());
+
+  ASSERT_TRUE(W.insert(2));
+  ASSERT_TRUE(W.insert(7));
+  // Use a raw array.
+  int A[] = {7, 5};
+  W.insert(A);
+  EXPECT_EQ(5, W.pop_back_val());
+  EXPECT_EQ(7, W.pop_back_val());
+  EXPECT_EQ(2, W.pop_back_val());
+  ASSERT_TRUE(W.empty());
+}
+
 TYPED_TEST(PriorityWorklistTest, EraseIf) {
   TypeParam W;
   W.insert(23);