diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h index 1f8571a851e1..88826ab49a7c 100644 --- a/llvm/include/llvm/ADT/SmallPtrSet.h +++ b/llvm/include/llvm/ADT/SmallPtrSet.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -336,6 +337,10 @@ public: insert(*I); } + void insert(std::initializer_list IL) { + insert(IL.begin(), IL.end()); + } + inline iterator begin() const { return iterator(CurArray, EndPointer()); } @@ -374,6 +379,11 @@ public: this->insert(I, E); } + SmallPtrSet(std::initializer_list IL) + : BaseT(SmallStorage, SmallSizePowTwo) { + this->insert(IL.begin(), IL.end()); + } + SmallPtrSet & operator=(const SmallPtrSet &RHS) { if (&RHS != this) @@ -388,6 +398,13 @@ public: return *this; } + SmallPtrSet & + operator=(std::initializer_list IL) { + this->clear(); + this->insert(IL.begin(), IL.end()); + return *this; + } + /// swap - Swaps the elements of two sets. void swap(SmallPtrSet &RHS) { SmallPtrSetImplBase::swap(RHS); diff --git a/llvm/unittests/ADT/SmallPtrSetTest.cpp b/llvm/unittests/ADT/SmallPtrSetTest.cpp index d8d07b16cfe3..d4d963fdc5bd 100644 --- a/llvm/unittests/ADT/SmallPtrSetTest.cpp +++ b/llvm/unittests/ADT/SmallPtrSetTest.cpp @@ -21,10 +21,7 @@ TEST(SmallPtrSetTest, Assignment) { for (int i = 0; i < 8; ++i) buf[i] = 0; - SmallPtrSet s1; - s1.insert(&buf[0]); - s1.insert(&buf[1]); - + SmallPtrSet s1 = {&buf[0], &buf[1]}; SmallPtrSet s2; (s2 = s1).insert(&buf[2]); @@ -38,6 +35,15 @@ TEST(SmallPtrSetTest, Assignment) { EXPECT_TRUE(s1.count(&buf[i])); else EXPECT_FALSE(s1.count(&buf[i])); + + // Assign and insert with initializer lists, and ones that contain both + // duplicates and out-of-order elements. + (s2 = {&buf[6], &buf[7], &buf[6]}).insert({&buf[5], &buf[4]}); + for (int i = 0; i < 8; ++i) + if (i < 4) + EXPECT_FALSE(s2.count(&buf[i])); + else + EXPECT_TRUE(s2.count(&buf[i])); } TEST(SmallPtrSetTest, GrowthTest) {