Index: llvm/trunk/include/llvm/ADT/SmallSet.h =================================================================== --- llvm/trunk/include/llvm/ADT/SmallSet.h +++ llvm/trunk/include/llvm/ADT/SmallSet.h @@ -17,14 +17,70 @@ #include "llvm/ADT/None.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/iterator.h" #include "llvm/Support/Compiler.h" #include #include #include +#include #include namespace llvm { +/// SmallSetIterator - This class implements a const_iterator for SmallSet by +/// delegating to the underlying SmallVector or Set iterators. +template +class SmallSetIterator + : public iterator_facade_base, + std::forward_iterator_tag, T> { +private: + using SetIterTy = typename std::set::const_iterator; + using VecIterTy = typename SmallVector::const_iterator; + using SelfTy = SmallSetIterator; + + /// Iterators to the parts of the SmallSet containing the data. They are set + /// depending on isSmall. + union { + SetIterTy SetIter; + VecIterTy VecIter; + }; + + bool isSmall; + +public: + SmallSetIterator(SetIterTy SetIter) : SetIter(SetIter), isSmall(false) { + // Use static_assert here, as the SmallSetIterator type is incomplete in the + // class scope. + static_assert(std::is_trivially_destructible::value && + std::is_trivially_copy_constructible::value && + std::is_trivially_copy_assignable::value && + std::is_move_constructible::value && + std::is_trivially_move_assignable::value, + "SelfTy needs to be trivial copy and move constructors and " + "assignments"); + } + + SmallSetIterator(VecIterTy VecIter) : VecIter(VecIter), isSmall(true) {} + + bool operator==(const SmallSetIterator &RHS) const { + if (isSmall != RHS.isSmall) + return false; + if (isSmall) + return VecIter == RHS.VecIter; + return SetIter == RHS.SetIter; + } + + SmallSetIterator &operator++() { // Preincrement + if (isSmall) + VecIter++; + else + SetIter++; + return *this; + } + + const T &operator*() const { return isSmall ? *VecIter : *SetIter; } +}; + /// SmallSet - This maintains a set of unique values, optimizing for the case /// when the set is small (less than N). In this case, the set can be /// maintained with no mallocs. If the set gets large, we expand to using an @@ -50,6 +106,7 @@ public: using size_type = size_t; + using const_iterator = SmallSetIterator; SmallSet() = default; @@ -121,6 +178,18 @@ Set.clear(); } + const_iterator begin() const { + if (isSmall()) + return {Vector.begin()}; + return {Set.begin()}; + } + + const_iterator end() const { + if (isSmall()) + return {Vector.end()}; + return {Set.end()}; + } + private: bool isSmall() const { return Set.empty(); } Index: llvm/trunk/unittests/ADT/SmallSetTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/SmallSetTest.cpp +++ llvm/trunk/unittests/ADT/SmallSetTest.cpp @@ -13,6 +13,7 @@ #include "llvm/ADT/SmallSet.h" #include "gtest/gtest.h" +#include using namespace llvm; @@ -68,3 +69,57 @@ EXPECT_EQ(0u, s1.count(8)); } + +TEST(SmallSetTest, IteratorInt) { + SmallSet s1; + + // Test the 'small' case. + for (int i = 0; i < 3; i++) + s1.insert(i); + + std::vector V(s1.begin(), s1.end()); + // Make sure the elements are in the expected order. + std::sort(V.begin(), V.end()); + for (int i = 0; i < 3; i++) + EXPECT_EQ(i, V[i]); + + // Test the 'big' case by adding a few more elements to switch to std::set + // internally. + for (int i = 3; i < 6; i++) + s1.insert(i); + + V.assign(s1.begin(), s1.end()); + // Make sure the elements are in the expected order. + std::sort(V.begin(), V.end()); + for (int i = 0; i < 6; i++) + EXPECT_EQ(i, V[i]); +} + +TEST(SmallSetTest, IteratorString) { + // Test SmallSetIterator for SmallSet with a type with non-trivial + // ctors/dtors. + SmallSet s1; + + s1.insert("str 1"); + s1.insert("str 2"); + s1.insert("str 1"); + + std::vector V(s1.begin(), s1.end()); + std::sort(V.begin(), V.end()); + EXPECT_EQ(2u, s1.size()); + EXPECT_EQ("str 1", V[0]); + EXPECT_EQ("str 2", V[1]); + + s1.insert("str 4"); + s1.insert("str 0"); + s1.insert("str 4"); + + V.assign(s1.begin(), s1.end()); + // Make sure the elements are in the expected order. + std::sort(V.begin(), V.end()); + EXPECT_EQ(4u, s1.size()); + EXPECT_EQ("str 0", V[0]); + EXPECT_EQ("str 1", V[1]); + EXPECT_EQ("str 2", V[2]); + EXPECT_EQ("str 4", V[3]); +}