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 @@ -171,22 +172,21 @@ } /// insert - Insert an element into the set if it isn't already there. - /// Returns true if the element is inserted (it was not in the set before). - /// 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); + /// Returns a pair. The first value of it is an iterator to the inserted + /// element or the existing element in the set. The second value is true + /// if the element is inserted (it was not in the set before). + std::pair insert(const T &V) { + if (!isSmall()) { + auto [I, Inserted] = Set.insert(V); + return std::make_pair(const_iterator(I), Inserted); + } VIterator I = vfind(V); if (I != Vector.end()) // Don't reinsert if it already exists. - return std::make_pair(None, false); + return std::make_pair(const_iterator(I), 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. @@ -194,8 +194,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());