diff --git a/llvm/include/llvm/ADT/BitVector.h b/llvm/include/llvm/ADT/BitVector.h --- a/llvm/include/llvm/ADT/BitVector.h +++ b/llvm/include/llvm/ADT/BitVector.h @@ -28,6 +28,45 @@ namespace llvm { +/// ForwardIterator for all bits. +/// Iterators get invalidated when resize / reserve is called. +template class const_all_bits_iterator_impl { + using size_type = typename BitVectorT::size_type; + + const BitVectorT &Parent; + size_type Current; + +public: + const_all_bits_iterator_impl(const BitVectorT &Parent, size_type Current) + : Parent(Parent), Current(Current) {} + const_all_bits_iterator_impl(const const_all_bits_iterator_impl &) = default; + + const_all_bits_iterator_impl operator++(int) { + auto Prev = *this; + ++Current; + return Prev; + } + + const_all_bits_iterator_impl &operator++() { + ++Current; + return *this; + } + + bool operator*() { return Parent[Current]; } + + bool operator==(const const_all_bits_iterator_impl &Other) const { + assert(&Parent == &Other.Parent && + "Comparing iterators from different BitVectors"); + return Current == Other.Current; + } + + bool operator!=(const const_all_bits_iterator_impl &Other) const { + assert(&Parent == &Other.Parent && + "Comparing iterators from different BitVectors"); + return Current != Other.Current; + } +}; + /// ForwardIterator for the bits that are set. /// Iterators get invalidated when resize / reserve is called. template class const_set_bits_iterator_impl { @@ -121,6 +160,11 @@ } }; + using const_all_bits_iterator = const_all_bits_iterator_impl; + + const_all_bits_iterator begin() const { return {*this, 0}; } + const_all_bits_iterator end() const { return {*this, size()}; } + typedef const_set_bits_iterator_impl const_set_bits_iterator; typedef const_set_bits_iterator set_iterator; diff --git a/llvm/include/llvm/ADT/SmallBitVector.h b/llvm/include/llvm/ADT/SmallBitVector.h --- a/llvm/include/llvm/ADT/SmallBitVector.h +++ b/llvm/include/llvm/ADT/SmallBitVector.h @@ -177,6 +177,11 @@ delete getPointer(); } + using const_all_bits_iterator = const_all_bits_iterator_impl; + + const_all_bits_iterator begin() const { return {*this, 0}; } + const_all_bits_iterator end() const { return {*this, size()}; } + using const_set_bits_iterator = const_set_bits_iterator_impl; using set_iterator = const_set_bits_iterator; diff --git a/llvm/unittests/ADT/BitVectorTest.cpp b/llvm/unittests/ADT/BitVectorTest.cpp --- a/llvm/unittests/ADT/BitVectorTest.cpp +++ b/llvm/unittests/ADT/BitVectorTest.cpp @@ -1143,7 +1143,31 @@ testEmpty(E); } -TYPED_TEST(BitVectorTest, Iterators) { +TYPED_TEST(BitVectorTest, AllBitsIterators) { + TypeParam Empty; + EXPECT_EQ(Empty.begin(), Empty.end()); + unsigned BitCount = 0; + for (bool Bit : Empty) { + (void)Bit; + BitCount++; + } + ASSERT_EQ(BitCount, 0U); + + TypeParam ToFill(100, false); + const unsigned List[] = {1, 10, 25, 99}; + for (unsigned Idx : List) + ToFill.set(Idx); + + unsigned Idx = 0; + for (bool Bit : ToFill) { + ASSERT_TRUE(Bit ? llvm::is_contained(List, Idx) + : !llvm::is_contained(List, Idx)); + ++Idx; + } + ASSERT_EQ(Idx, 100U); +} + +TYPED_TEST(BitVectorTest, SetBitsIterators) { TypeParam Filled(10, true); EXPECT_NE(Filled.set_bits_begin(), Filled.set_bits_end()); unsigned Counter = 0;