[ADT] Support BitVector as a key in DenseSet/Map

This patch adds DenseMapInfo<> support for BitVector and SmallBitVector.

This is part of https://reviews.llvm.org/D71775, where a BitVector is used as a thread affinity mask.

include/llvm/ADT/BitVector.h

@@ -14,6 +14,7 @@
 #define LLVM_ADT_BITVECTOR_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/MathExtras.h"
 #include <algorithm>
@@ -719,6 +720,14 @@ public:
     if (this == &RHS) return *this;
 
     Size = RHS.size();
+
+    // Handle tombstone when the BitVector is a key of a DenseHash.
+    if (RHS.isInvalid()) {
+      std::free(Bits.data());
+      Bits = None;
+      return *this;
+    }
+
     unsigned RHSWords = NumBitWords(Size);
     if (Size <= getBitCapacity()) {
       if (Size)
@@ -758,6 +767,14 @@ public:
     std::swap(Size, RHS.Size);
   }
 
+  void invalid() {
+    assert(!Size && Bits.empty());
+    Size = (unsigned)-1;
+  }
+  bool isInvalid() const { return Size == (unsigned)-1; }
+
+  ArrayRef<BitWord> getData() const { return Bits; }
+
   //===--------------------------------------------------------------------===//
   // Portable bit mask operations.
   //===--------------------------------------------------------------------===//
@@ -932,6 +949,23 @@ inline size_t capacity_in_bytes(const BitVector &X) {
   return X.getMemorySize();
 }
 
+template <> struct DenseMapInfo<BitVector> {
+  static inline BitVector getEmptyKey() { return BitVector(); }
+  static inline BitVector getTombstoneKey() {
+    BitVector V;
+    V.invalid();
+    return V;
+  }
+  static unsigned getHashValue(const BitVector &V) {
+    return DenseMapInfo<std::pair<unsigned, ArrayRef<uintptr_t>>>::getHashValue(
+        std::make_pair(V.size(), V.getData()));
+  }
+  static bool isEqual(const BitVector &LHS, const BitVector &RHS) {
+    if (LHS.isInvalid() || RHS.isInvalid())
+      return LHS.isInvalid() == RHS.isInvalid();
+    return LHS == RHS;
+  }
+};
 } // end namespace llvm
 
 namespace std {

include/llvm/ADT/SmallBitVector.h

@@ -662,6 +662,16 @@ public:
       getPointer()->clearBitsNotInMask(Mask, MaskWords);
   }
 
+  void invalid() {
+    assert(empty());
+    X = (uintptr_t)-1;
+  }
+  bool isInvalid() const { return X == (uintptr_t)-1; }
+
+  ArrayRef<uintptr_t> getData() const {
+    return isSmall() ? makeArrayRef(X) : getPointer()->getData();
+  }
+
 private:
   template <bool AddBits, bool InvertMask>
   void applyMask(const uint32_t *Mask, unsigned MaskWords) {
@@ -699,6 +709,23 @@ operator^(const SmallBitVector &LHS, const SmallBitVector &RHS) {
   return Result;
 }
 
+template <> struct DenseMapInfo<SmallBitVector> {
+  static inline SmallBitVector getEmptyKey() { return SmallBitVector(); }
+  static inline SmallBitVector getTombstoneKey() {
+    SmallBitVector V;
+    V.invalid();
+    return V;
+  }
+  static unsigned getHashValue(const SmallBitVector &V) {
+    return DenseMapInfo<std::pair<unsigned, ArrayRef<uintptr_t>>>::getHashValue(
+        std::make_pair(V.size(), V.getData()));
+  }
+  static bool isEqual(const SmallBitVector &LHS, const SmallBitVector &RHS) {
+    if (LHS.isInvalid() || RHS.isInvalid())
+      return LHS.isInvalid() == RHS.isInvalid();
+    return LHS == RHS;
+  }
+};
 } // end namespace llvm
 
 namespace std {

unittests/ADT/BitVectorTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "gtest/gtest.h"
 
@@ -1149,4 +1150,38 @@ TYPED_TEST(BitVectorTest, PushBack) {
   EXPECT_EQ(213U, Vec.size());
   EXPECT_EQ(102U, Vec.count());
 }
+
+TYPED_TEST(BitVectorTest, DenseSet) {
+  DenseSet<TypeParam> Set;
+  TypeParam A(10, true);
+  auto I = Set.insert(A);
+  EXPECT_EQ(true, I.second);
+
+  TypeParam B(5, true);
+  I = Set.insert(B);
+  EXPECT_EQ(true, I.second);
+
+  TypeParam C(20, false);
+  C.set(19);
+  I = Set.insert(C);
+  EXPECT_EQ(true, I.second);
+
+  TypeParam D;
+  EXPECT_DEATH(Set.insert(D),
+               "Empty/Tombstone value shouldn't be inserted into map!");
+
+  EXPECT_EQ(3U, Set.size());
+  EXPECT_EQ(1U, Set.count(A));
+  EXPECT_EQ(1U, Set.count(B));
+  EXPECT_EQ(1U, Set.count(C));
+
+  EXPECT_EQ(true, Set.erase(B));
+  EXPECT_EQ(2U, Set.size());
+
+  EXPECT_EQ(true, Set.erase(C));
+  EXPECT_EQ(1U, Set.size());
+
+  EXPECT_EQ(true, Set.erase(A));
+  EXPECT_EQ(0U, Set.size());
 }
+} // namespace