[llvm] [ADT] Refactor Bitset to Be More Constexpr-Usable and Add More Member Functions (PR #172062)

[Jakub Kuderski via llvm-commits]

================
@@ -144,25 +166,90 @@ class Bitset {
 
   constexpr Bitset operator~() const {
     Bitset Result = *this;
-    for (auto &B : Result.Bits)
+    for (BitWord &B : Result.Bits)
       B = ~B;
+    Result.maskLastWord();
     return Result;
   }
 
-  bool operator==(const Bitset &RHS) const {
-    return std::equal(std::begin(Bits), std::end(Bits), std::begin(RHS.Bits));
+  constexpr bool operator==(const Bitset &RHS) const {
+    for (unsigned I = 0; I < NumWords; ++I)
+      if (Bits[I] != RHS.Bits[I])
+        return false;
+    return true;
   }
 
-  bool operator!=(const Bitset &RHS) const { return !(*this == RHS); }
+  constexpr bool operator!=(const Bitset &RHS) const { return !(*this == RHS); }
 
-  bool operator < (const Bitset &Other) const {
+  constexpr bool operator<(const Bitset &Other) const {
     for (unsigned I = 0, E = size(); I != E; ++I) {
       bool LHS = test(I), RHS = Other.test(I);
       if (LHS != RHS)
         return LHS < RHS;
     }
     return false;
   }
+
+  constexpr Bitset &operator<<=(unsigned N) {
+    if (N == 0)
+      return *this;
+    if (N >= NumBits) {
+      return *this = Bitset();
+    }
+    const unsigned WordShift = N / BitwordBits;
+    const unsigned BitShift = N % BitwordBits;
+    if (BitShift == 0) {
+      for (int I = getLastWordIndex(); I >= static_cast<int>(WordShift); --I)
+        Bits[I] = Bits[I - WordShift];
+    } else {
+      const unsigned CarryShift = BitwordBits - BitShift;
+      for (int I = getLastWordIndex(); I > static_cast<int>(WordShift); --I) {
+        Bits[I] = (Bits[I - WordShift] << BitShift) |
+                  (Bits[I - WordShift - 1] >> CarryShift);
+      }
+      Bits[WordShift] = Bits[0] << BitShift;
+    }
+    for (unsigned I = 0; I < WordShift; ++I)
+      Bits[I] = 0;
+    maskLastWord();
+    return *this;
+  }
+
+  constexpr Bitset operator<<(unsigned N) const {
+    Bitset Result(*this);
+    Result <<= N;
+    return Result;
+  }
+
+  constexpr Bitset &operator>>=(unsigned N) {
+    if (N == 0)
+      return *this;
+    if (N >= NumBits) {
+      return *this = Bitset();
+    }
+    const unsigned WordShift = N / BitwordBits;
+    const unsigned BitShift = N % BitwordBits;
+    if (BitShift == 0) {
+      for (unsigned I = 0; I < NumWords - WordShift; ++I)
+        Bits[I] = Bits[I + WordShift];
+    } else {
+      const unsigned CarryShift = BitwordBits - BitShift;
+      for (unsigned I = 0; I < NumWords - WordShift - 1; ++I) {
+        Bits[I] = (Bits[I + WordShift] >> BitShift) |
+                  (Bits[I + WordShift + 1] << CarryShift);
+      }
+      Bits[NumWords - WordShift - 1] = Bits[getLastWordIndex()] >> BitShift;
+    }
+    for (unsigned I = NumWords - WordShift; I < NumWords; ++I)
+      Bits[I] = 0;
+    return *this;
+  }
+
+  constexpr Bitset operator>>(unsigned N) const {
+    Bitset Result(*this);
+    Result >>= N;
+    return Result;
+  }
----------------
kuhar wrote:

Can you land the new operator support in a separate PR?

https://github.com/llvm/llvm-project/pull/172062