Index: include/llvm/ADT/SmallSet.h =================================================================== --- include/llvm/ADT/SmallSet.h +++ include/llvm/ADT/SmallSet.h @@ -17,6 +17,7 @@ #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 @@ -25,6 +26,56 @@ 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; + + /// 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) {} + SmallSetIterator(VecIterTy VecIter) : VecIter(VecIter), isSmall(true) {} + + ~SmallSetIterator() { + // Destruct underlying iterator object depending on size. + if (isSmall) + VecIter.~VecIterTy(); + else + SetIter.~SetIterTy(); + } + + 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 +101,7 @@ public: using size_type = size_t; + using const_iterator = SmallSetIterator; SmallSet() = default; @@ -121,6 +173,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: unittests/ADT/SmallSetTest.cpp =================================================================== --- unittests/ADT/SmallSetTest.cpp +++ unittests/ADT/SmallSetTest.cpp @@ -31,10 +31,6 @@ for (int i = 0; i < 4; i++) EXPECT_EQ(1u, s1.count(i)); -<<<<<<< 458a792747d59ffe24a712aadf91104d2d9cd1ee -======= - EXPECT_EQ(0u, s1.count(99)); ->>>>>>> [SmallSet] Add some simple unit tests. EXPECT_EQ(0u, s1.count(4)); } @@ -49,12 +45,7 @@ for (int i = 0; i < 8; i++) EXPECT_EQ(1u, s1.count(i)); -<<<<<<< 458a792747d59ffe24a712aadf91104d2d9cd1ee EXPECT_EQ(0u, s1.count(8)); -======= - EXPECT_EQ(0u, s1.count(99)); - EXPECT_EQ(0u, s1.count(10)); ->>>>>>> [SmallSet] Add some simple unit tests. } TEST(SmallSetTest, Erase) { @@ -77,3 +68,30 @@ EXPECT_EQ(0u, s1.count(8)); } + +TEST(SmallSetTest, Iterator) { + 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); + { + 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 < 6; i++) + EXPECT_EQ(i, V[i]); + } +}