diff --git a/llvm/include/llvm/ADT/SmallSet.h b/llvm/include/llvm/ADT/SmallSet.h --- a/llvm/include/llvm/ADT/SmallSet.h +++ b/llvm/include/llvm/ADT/SmallSet.h @@ -141,6 +141,7 @@ std::set Set; using VIterator = typename SmallVector::const_iterator; + using SIterator = typename std::set::const_iterator; using mutable_iterator = typename SmallVector::iterator; // In small mode SmallPtrSet uses linear search for the elements, so it is @@ -177,18 +178,20 @@ /// The first value of the returned pair is unused and provided for /// partial compatibility with the standard library self-associative container /// concept. - // FIXME: Add iterators that abstract over the small and large form, and then - // return those here. - std::pair insert(const T &V) { - if (!isSmall()) - return std::make_pair(None, Set.insert(V).second); + std::pair insert(const T &V) { + if (!isSmall()) { + SIterator SIt; + bool Inserted; + std::tie(SIt, Inserted) = Set.insert(V); + return std::make_pair(const_iterator(SIt), Inserted); + } - VIterator I = vfind(V); - if (I != Vector.end()) // Don't reinsert if it already exists. - return std::make_pair(None, false); + VIterator VIt = vfind(V); + if (VIt != Vector.end()) // Don't reinsert if it already exists. + return std::make_pair(const_iterator(VIt), false); if (Vector.size() < N) { Vector.push_back(V); - return std::make_pair(None, true); + return std::make_pair(const_iterator(std::prev(Vector.end())), true); } // Otherwise, grow from vector to set. @@ -196,8 +199,7 @@ Set.insert(Vector.back()); Vector.pop_back(); } - Set.insert(V); - return std::make_pair(None, true); + return std::make_pair(const_iterator(Set.insert(V).first), true); } template diff --git a/llvm/unittests/ADT/SmallSetTest.cpp b/llvm/unittests/ADT/SmallSetTest.cpp --- a/llvm/unittests/ADT/SmallSetTest.cpp +++ b/llvm/unittests/ADT/SmallSetTest.cpp @@ -21,11 +21,17 @@ SmallSet s1; - for (int i = 0; i < 4; i++) - s1.insert(i); + for (int i = 0; i < 4; i++) { + auto InsertResult = s1.insert(i); + EXPECT_EQ(*InsertResult.first, i); + EXPECT_EQ(InsertResult.second, true); + } - for (int i = 0; i < 4; i++) - s1.insert(i); + for (int i = 0; i < 4; i++) { + auto InsertResult = s1.insert(i); + EXPECT_EQ(*InsertResult.first, i); + EXPECT_EQ(InsertResult.second, false); + } EXPECT_EQ(4u, s1.size()); @@ -38,8 +44,17 @@ TEST(SmallSetTest, Grow) { SmallSet s1; - for (int i = 0; i < 8; i++) - s1.insert(i); + for (int i = 0; i < 8; i++) { + auto InsertResult = s1.insert(i); + EXPECT_EQ(*InsertResult.first, i); + EXPECT_EQ(InsertResult.second, true); + } + + for (int i = 0; i < 8; i++) { + auto InsertResult = s1.insert(i); + EXPECT_EQ(*InsertResult.first, i); + EXPECT_EQ(InsertResult.second, false); + } EXPECT_EQ(8u, s1.size());