2 files changed, 109 insertions, 0 deletions

diff --git a/include/llvm/ADT/BitVector.h b/include/llvm/ADT/BitVector.h
index ac1cf0c..7d7afc3 100644
--- a/include/llvm/ADT/BitVector.h
+++ b/include/llvm/ADT/BitVector.h
@@ -365,6 +365,42 @@ public:
     std::swap(Capacity, RHS.Capacity);
   }
 
+  //===--------------------------------------------------------------------===//
+  // Portable bit mask operations.
+  //===--------------------------------------------------------------------===//
+  //
+  // These methods all operate on arrays of uint32_t, each holding 32 bits. The
+  // fixed word size makes it easier to work with literal bit vector constants
+  // in portable code.
+  //
+  // The LSB in each word is the lowest numbered bit.  The size of a portable
+  // bit mask is always a whole multiple of 32 bits.  If no bit mask size is
+  // given, the bit mask is assumed to cover the entire BitVector.
+
+  /// setBitsInMask - Add '1' bits from Mask to this vector. Don't resize.
+  /// This computes "*this |= Mask".
+  void setBitsInMask(const uint32_t *Mask, unsigned MaskWords = ~0u) {
+    applyMask<true, false>(Mask, MaskWords);
+  }
+
+  /// clearBitsInMask - Clear any bits in this vector that are set in Mask.
+  /// Don't resize. This computes "*this &= ~Mask".
+  void clearBitsInMask(const uint32_t *Mask, unsigned MaskWords = ~0u) {
+    applyMask<false, false>(Mask, MaskWords);
+  }
+
+  /// setBitsNotInMask - Add a bit to this vector for every '0' bit in Mask.
+  /// Don't resize.  This computes "*this |= ~Mask".
+  void setBitsNotInMask(const uint32_t *Mask, unsigned MaskWords = ~0u) {
+    applyMask<true, true>(Mask, MaskWords);
+  }
+
+  /// clearBitsNotInMask - Clear a bit in this vector for every '0' bit in Mask.
+  /// Don't resize.  This computes "*this &= Mask".
+  void clearBitsNotInMask(const uint32_t *Mask, unsigned MaskWords = ~0u) {
+    applyMask<false, true>(Mask, MaskWords);
+  }
+
 private:
   unsigned NumBitWords(unsigned S) const {
     return (S + BITWORD_SIZE-1) / BITWORD_SIZE;
@@ -400,6 +436,33 @@ private:
   void init_words(BitWord *B, unsigned NumWords, bool t) {
     memset(B, 0 - (int)t, NumWords*sizeof(BitWord));
   }
+
+  template<bool AddBits, bool InvertMask>
+  void applyMask(const uint32_t *Mask, unsigned MaskWords) {
+    assert(BITWORD_SIZE % 32 == 0 && "Unsupported BitWord size.");
+    MaskWords = std::min(MaskWords, (size() + 31) / 32);
+    const unsigned Scale = BITWORD_SIZE / 32;
+    unsigned i;
+    for (i = 0; MaskWords >= Scale; ++i, MaskWords -= Scale) {
+      BitWord BW = Bits[i];
+      // This inner loop should unroll completely when BITWORD_SIZE > 32.
+      for (unsigned b = 0; b != BITWORD_SIZE; b += 32) {
+        uint32_t M = *Mask++;
+        if (InvertMask) M = ~M;
+        if (AddBits) BW |=   BitWord(M) << b;
+        else         BW &= ~(BitWord(M) << b);
+      }
+      Bits[i] = BW;
+    }
+    for (unsigned b = 0; MaskWords; b += 32, --MaskWords) {
+      uint32_t M = *Mask++;
+      if (InvertMask) M = ~M;
+      if (AddBits) Bits[i] |=   BitWord(M) << b;
+      else         Bits[i] &= ~(BitWord(M) << b);
+    }
+    if (AddBits)
+      clear_unused_bits();
+  }
 };
 
 inline BitVector operator&(const BitVector &LHS, const BitVector &RHS) {
diff --git a/unittests/ADT/BitVectorTest.cpp b/unittests/ADT/BitVectorTest.cpp
index fa66312..f733e13 100644
--- a/unittests/ADT/BitVectorTest.cpp
+++ b/unittests/ADT/BitVectorTest.cpp
@@ -196,6 +196,52 @@ TEST(BitVectorTest, ProxyIndex) {
   EXPECT_TRUE(Vec.none());
 }
 
+TEST(BitVectorTest, PortableBitMask) {
+  BitVector A;
+  const uint32_t Mask1[] = { 0x80000000, 6, 5 };
+
+  A.resize(10);
+  A.setBitsInMask(Mask1, 3);
+  EXPECT_EQ(10u, A.size());
+  EXPECT_FALSE(A.test(0));
+
+  A.resize(32);
+  A.setBitsInMask(Mask1, 3);
+  EXPECT_FALSE(A.test(0));
+  EXPECT_TRUE(A.test(31));
+  EXPECT_EQ(1u, A.count());
+
+  A.resize(33);
+  A.setBitsInMask(Mask1, 1);
+  EXPECT_EQ(1u, A.count());
+  A.setBitsInMask(Mask1, 2);
+  EXPECT_EQ(1u, A.count());
+
+  A.resize(34);
+  A.setBitsInMask(Mask1, 2);
+  EXPECT_EQ(2u, A.count());
+
+  A.resize(65);
+  A.setBitsInMask(Mask1, 3);
+  EXPECT_EQ(4u, A.count());
+
+  A.setBitsNotInMask(Mask1, 1);
+  EXPECT_EQ(32u+3u, A.count());
+
+  A.setBitsNotInMask(Mask1, 3);
+  EXPECT_EQ(65u, A.count());
+
+  A.resize(96);
+  EXPECT_EQ(65u, A.count());
+
+  A.clear();
+  A.resize(128);
+  A.setBitsNotInMask(Mask1, 3);
+  EXPECT_EQ(96u-5u, A.count());
+
+  A.clearBitsNotInMask(Mask1, 1);
+  EXPECT_EQ(64-4u, A.count());
+}
 }
 
 #endif