Index: include/llvm/ADT/BitVector.h =================================================================== --- include/llvm/ADT/BitVector.h +++ include/llvm/ADT/BitVector.h @@ -14,6 +14,7 @@ #ifndef LLVM_ADT_BITVECTOR_H #define LLVM_ADT_BITVECTOR_H +#include "llvm/ADT/iterator_range.h" #include "llvm/Support/MathExtras.h" #include #include @@ -73,6 +74,48 @@ } }; + /// ForwardIterator for the bits that are set. + /// Iterators get invalidated when resize / reserve is called. + template + class const_set_iterator_impl { + const BitVectorT *Parent = nullptr; + int Current = 0; + + public: + const_set_iterator_impl(const BitVectorT *Parent, int Current) + : Parent(Parent), Current(Current) {} + const_set_iterator_impl(const BitVectorT *Parent) + : const_set_iterator_impl(Parent, Parent->find_first()) {} + const_set_iterator_impl(const const_set_iterator_impl&) = default; + + void advance() { + assert(Current != -1 && "Trying to advance past end."); + Current = Parent->find_next(Current); + } + + const_set_iterator_impl operator++(int) { + auto Prev = *this; + advance(); + return Prev; + } + + const_set_iterator_impl &operator++() { + advance(); + return *this; + } + + int operator*() const { return Current; } + const BitVectorT *getParent() const { return Parent; } + }; + + typedef const_set_iterator_impl const_set_iterator; + typedef const_set_iterator set_iterator; + + const_set_iterator set_begin() const { return {this}; } + const_set_iterator set_end() const { return {this, -1}; } + iterator_range bits_set() const { + return make_range(set_begin(), set_end()); + } /// BitVector default ctor - Creates an empty bitvector. BitVector() : Size(0), Capacity(0) { @@ -619,6 +662,22 @@ return X.getMemorySize(); } +template +bool operator==(const BitVector::const_set_iterator_impl &This, + const BitVector::const_set_iterator_impl &Other) { + assert(This.getParent() == Other.getParent() && + "Comparing iterators from different BitVectors"); + return *This == *Other; +} + +template +bool operator!=(const BitVector::const_set_iterator_impl &This, + const BitVector::const_set_iterator_impl &Other) { + assert(This.getParent() == Other.getParent() && + "Comparing iterators from different BitVectors"); + return *This != *Other; +} + } // end namespace llvm namespace std { Index: include/llvm/ADT/SmallBitVector.h =================================================================== --- include/llvm/ADT/SmallBitVector.h +++ include/llvm/ADT/SmallBitVector.h @@ -136,6 +136,16 @@ } public: + + typedef BitVector::const_set_iterator_impl const_set_iterator; + typedef const_set_iterator set_iterator; + + const_set_iterator set_begin() const { return {this}; } + const_set_iterator set_end() const { return {this, -1}; } + iterator_range bits_set() const { + return make_range(set_begin(), set_end()); + } + /// Creates an empty bitvector. SmallBitVector() : X(1) {} Index: unittests/ADT/BitVectorTest.cpp =================================================================== --- unittests/ADT/BitVectorTest.cpp +++ unittests/ADT/BitVectorTest.cpp @@ -501,5 +501,37 @@ testEmpty(E); } +TYPED_TEST(BitVectorTest, Iterators) { + TypeParam Filled(10, true); + EXPECT_NE(Filled.set_begin(), Filled.set_end()); + int Counter = 0; + for (int Bit : Filled.bits_set()) + EXPECT_EQ(Bit, Counter++); + + TypeParam Empty; + EXPECT_EQ(Empty.set_begin(), Empty.set_end()); + for (int Bit : Empty.bits_set()) + EXPECT_FALSE(Bit != 9999); + + TypeParam ToFill(100, false); + ToFill.set(0); + EXPECT_NE(ToFill.set_begin(), ToFill.set_end()); + EXPECT_EQ(++ToFill.set_begin(), ToFill.set_end()); + EXPECT_EQ(*ToFill.set_begin(), 0); + + ToFill.set(1); + ToFill.set(10); + ToFill.set(99); + ToFill.set(25); + size_t Count = ToFill.count(); + EXPECT_EQ((int)Count, 5); + size_t ItCount = 0; + for (int Bit : ToFill.bits_set()) { + (void)Bit; + ++ItCount; + } + EXPECT_EQ(Count, ItCount); +} + } #endif