Index: include/llvm/ADT/SmallSet.h =================================================================== --- include/llvm/ADT/SmallSet.h +++ include/llvm/ADT/SmallSet.h @@ -25,6 +25,64 @@ namespace llvm { +/// SmallSetIterator - This class implements a const_iterator for SmallSet by +/// delegating to the underlying SmallVector or Set iterators. +template class SmallSetIterator { +private: + using SetIterTy = typename std::set::const_iterator; + using VecIterTy = typename SmallVector::const_iterator; + + bool isSmall; + /// Iterators to the parts of the SmallSet containing the data. They are set + /// depending on isSmall. + union { + SetIterTy SetIter; + VecIterTy VecIter; + }; + +public: + using value_type = T; + using reference = T; + using pointer = T; + using difference_type = std::ptrdiff_t; + using iterator_category = std::forward_iterator_tag; + + SmallSetIterator(SetIterTy SetIter) : isSmall(false), SetIter(SetIter) {} + SmallSetIterator(VecIterTy VecIter) : isSmall(true), VecIter(VecIter) {} + + ~SmallSetIterator() { + // Destruct underlying iterator object depending on size. + if (isSmall) + VecIter.~VecIterTy(); + else + SetIter.~SetIterTy(); + } + + bool operator==(const SmallSetIterator &RHS) const { + if (isSmall) + return VecIter == RHS.VecIter; + return SetIter == RHS.SetIter; + } + + bool operator!=(const SmallSetIterator &RHS) const { return !(*this == RHS); } + + SmallSetIterator &operator++() { // Preincrement + if (isSmall) + VecIter++; + else + SetIter++; + return *this; + } + + SmallSetIterator operator++(int) { // Postincrement + SmallSetIterator tmp = *this; + ++*this; + return tmp; + } + + 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 +108,7 @@ public: using size_type = size_t; + using const_iterator = SmallSetIterator; SmallSet() = default; @@ -121,6 +180,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 @@ -71,3 +71,17 @@ EXPECT_EQ(0u, s1.count(99)); EXPECT_EQ(0u, s1.count(10)); } + +TEST(SmallSetTest, Iterator) { + SmallSet s1; + + for (int i = 0; i < 8; 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 < 8; i++) + EXPECT_EQ(i, V[i]); +}