BitVector: add iterators for bits that are set

[thegameg]

Summary

Add iterators to iterate on the bits of the bit vector that are set. This is useful when BitVector is used as a set<unsigned>, which usually is a set of basic block numbers or a set of registers.

D30808 would be a user of this.

This would allow us to transform:

for (int Elt = Vector.find_first(); Elt >= 0; Elt = Vector.find_next(Elt)) {
  // Stuff
}

into

for (int Elt : Vector.bits_set()) {
  // Stuff
}

[MatzeB]

Thanks for working on this. I missed this feature myself a couple of times.

We generally do not commit unused code, so before commit this we have to wait until one of the shrink-wrapping patches is accepted. Or you go around llvm and try to find a few find_next() users that can be replaced with a range based for.

include/llvm/ADT/BitVector.h

123–124

You should probably pull this out of the class, as it doesn't use any private methods and is used not just by BitVector but also by SmallBitVector.

125

Does this ever make sense to be nullptr? I'd recommend using a reference.

131

I'd recommend an explicit modifier here.

135–138

This can be private.

151

We have size_type BitVector::size() or BitVector::set(unsigned Idx). We seem to be a bit inconsistent regarding size_type vs. unsigned but the majority of functions appear to be using unsigned for bit indexes so this should be the return type here.

152

Is it necessary to make this accessible? Having access to the parent bitvector seems like an implementation detail to me that could stay private.

155–162

This could be more consistent: const_bits_set_iterator, bit_set_begin() and bit_set().

158–159

http://llvm.org/docs/CodingStandards.html#do-not-use-braced-initializer-lists-to-call-a-constructor recommends against calling constructors with curly braces.

679–693

Isn't it possible to move those into the iterator class?

unittests/ADT/BitVectorTest.cpp

514

How about EXPECT_TRUE(false); to make it more obvious?

526–533

Maybe construct an array {1,10,99,25} that you verify against during the iteration to make sure we do indeed return the right numbers?

[MatzeB]

Adding a few reviewers who often have comments on ADT changes.

[thegameg]

Thanks Matthias for the review.

[thegameg]

Fixl typo in unit tests.

[MatzeB]

include/llvm/ADT/SmallBitVector.h

139–144

This is still inconsistent bits_set vs. bit_set.

[thegameg]

Updated some passes to use this.

[thegameg]

Better name for the iterator range.
Fix typo.

[MatzeB]

This is looking pretty good, but I'd expect the range based for loops to use the same type as operator() of operator*(), i.e. `for (unsigned Bit : Empty.set_bits()) {}.

[thegameg]

int -> unsigned.

[MatzeB]

LGTM

include/llvm/ADT/BitVector.h

 //===- llvm/ADT/BitVector.h - Bit vectors -----------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the BitVector class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ADT_BITVECTOR_H
 #define LLVM_ADT_BITVECTOR_H
 
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/MathExtras.h"
 #include <algorithm>
 #include <cassert>
 #include <climits>
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <utility>
 
 namespace llvm {
 
+/// ForwardIterator for the bits that are set.
+/// Iterators get invalidated when resize / reserve is called.
+template <typename BitVectorT> class const_bits_set_iterator_impl {
+  const BitVectorT &Parent;
+  int Current = 0;
+
+  void advance() {
+    assert(Current != -1 && "Trying to advance past end.");
+    Current = Parent.find_next(Current);
+  }
+
+public:
+  const_bits_set_iterator_impl(const BitVectorT &Parent, int Current)
+      : Parent(Parent), Current(Current) {}
+  explicit const_bits_set_iterator_impl(const BitVectorT &Parent)
+      : const_bits_set_iterator_impl(Parent, Parent.find_first()) {}
+  const_bits_set_iterator_impl(const const_bits_set_iterator_impl &) = default;
+
+  const_bits_set_iterator_impl operator++(int) {
+    auto Prev = *this;
+    advance();
+    return Prev;
+  }
+
+  const_bits_set_iterator_impl &operator++() {
+    advance();
+    return *this;
+  }
+
+  unsigned operator*() const { return Current; }
+
+  bool operator==(const const_bits_set_iterator_impl &Other) const {
+    assert(&Parent == &Other.Parent &&
+           "Comparing iterators from different BitVectors");
+    return Current == Other.Current;
+  }
+
+  bool operator!=(const const_bits_set_iterator_impl &Other) const {
+    assert(&Parent == &Other.Parent &&
+           "Comparing iterators from different BitVectors");
+    return Current != Other.Current;
+  }
+};
+
 class BitVector {
   typedef unsigned long BitWord;
 
   enum { BITWORD_SIZE = (unsigned)sizeof(BitWord) * CHAR_BIT };
 
   static_assert(BITWORD_SIZE == 64 || BITWORD_SIZE == 32,
                 "Unsupported word size");
 
reference& operator=(bool t) {
       return *this;
     }
 
     operator bool() const {
       return ((*WordRef) & (BitWord(1) << BitPos)) != 0;
     }
   };
 
+  typedef const_bits_set_iterator_impl<BitVector> const_bits_set_iterator;
+  typedef const_bits_set_iterator set_iterator;
+
+  const_bits_set_iterator bit_set_begin() const {
+    return const_bits_set_iterator(*this);
+  }
+  const_bits_set_iterator bit_set_end() const {
+    return const_bits_set_iterator(*this, -1);
+  }
+  iterator_range<const_bits_set_iterator> bit_set() const {
+    return make_range(bit_set_begin(), bit_set_end());
+  }
 
   /// BitVector default ctor - Creates an empty bitvector.
   BitVector() : Size(0), Capacity(0) {
     Bits = nullptr;
   }
 
   /// BitVector ctor - Creates a bitvector of specified number of bits. All
   /// bits are initialized to the specified value.
   explicit BitVector(unsigned s, bool t = false) : Size(s) {
     Capacity = NumBitWords(s);
     Bits = (BitWord *)std::malloc(Capacity * sizeof(BitWord));
     init_words(Bits, Capacity, t);
     if (t)
       clear_unused_bits();
   }
 
   /// BitVector copy ctor.
   BitVector(const BitVector &RHS) : Size(RHS.size()) {
     if (Size == 0) {
       Bits = nullptr;
       Capacity = 0;
       return;
     }
 
     Capacity = NumBitWords(RHS.size());
     Bits = (BitWord *)std::malloc(Capacity * sizeof(BitWord));
     std::memcpy(Bits, RHS.Bits, Capacity * sizeof(BitWord));
   }
 
   BitVector(BitVector &&RHS)
     : Bits(RHS.Bits), Size(RHS.Size), Capacity(RHS.Capacity) {
     RHS.Bits = nullptr;
     RHS.Size = RHS.Capacity = 0;
   }
 
   ~BitVector() {
     std::free(Bits);
   }
public:
   /// Return the size (in bytes) of the bit vector.
   size_t getMemorySize() const { return Capacity * sizeof(BitWord); }
 };
 
 static inline size_t capacity_in_bytes(const BitVector &X) {
   return X.getMemorySize();
 }
 
 } // end namespace llvm
 
 namespace std {
   /// Implement std::swap in terms of BitVector swap.
   inline void
   swap(llvm::BitVector &LHS, llvm::BitVector &RHS) {
     LHS.swap(RHS);
   }
 } // end namespace std
 
 #endif // LLVM_ADT_BITVECTOR_H

include/llvm/ADT/SmallBitVector.h

private:
   }
 
   void setSmallBits(uintptr_t NewBits) {
     setSmallRawBits((NewBits & ~(~uintptr_t(0) << getSmallSize())) |
                     (getSmallSize() << SmallNumDataBits));
   }
 
 public:
+  typedef const_bits_set_iterator_impl<SmallBitVector> const_bits_set_iterator;
+  typedef const_bits_set_iterator set_iterator;
+
+  const_bits_set_iterator bit_set_begin() const {
+    return const_bits_set_iterator(*this);
+  }
+  const_bits_set_iterator bit_set_end() const {
+    return const_bits_set_iterator(*this, -1);
+  }
+  iterator_range<const_bits_set_iterator> bit_set() const {
+    return make_range(bit_set_begin(), bit_set_end());
+  }
+
   /// Creates an empty bitvector.
   SmallBitVector() : X(1) {}
 
   /// Creates a bitvector of specified number of bits. All bits are initialized
   /// to the specified value.
   explicit SmallBitVector(unsigned s, bool t = false) {
     if (s <= SmallNumDataBits)
       switchToSmall(t ? ~uintptr_t(0) : 0, s);

unittests/ADT/BitVectorTest.cpp

TYPED_TEST(BitVectorTest, EmptyVector) {
   E = A;
   testEmpty(E);
 
   TypeParam F;
   E.reset(A);
   testEmpty(E);
 }
 
+TYPED_TEST(BitVectorTest, Iterators) {
+  TypeParam Filled(10, true);
+  EXPECT_NE(Filled.bit_set_begin(), Filled.bit_set_end());
+  int Counter = 0;
+  for (int Bit : Filled.bit_set())
+    EXPECT_EQ(Bit, Counter++);
+
+  TypeParam Empty;
+  EXPECT_EQ(Empty.bit_set_begin(), Empty.bit_set_end());
+  for (int Bit : Empty.bit_set()) {
+    (void)Bit;
+    EXPECT_FALSE(false);
+  }
+
+  TypeParam ToFill(100, false);
+  ToFill.set(0);
+  EXPECT_NE(ToFill.bit_set_begin(), ToFill.bit_set_end());
+  EXPECT_EQ(++ToFill.bit_set_begin(), ToFill.bit_set_end());
+  EXPECT_EQ(*ToFill.bit_set_begin(), 0U);
+  ToFill.reset(0);
+  EXPECT_EQ(ToFill.bit_set_begin(), ToFill.bit_set_end());
+
+  const int List[] = {1, 10, 25, 99};
+  for (int Num : List)
+    ToFill.set(Num);
+  int i = 0;
+  for (int Bit : ToFill.bit_set())
+    EXPECT_EQ(List[i++], Bit);
+}
 }
 #endif