Bug 1976207 - Allow mozilla::BitSet to use atomic storage r=glandium a=RyanVM

This patch relaxes the restriction on BitSet so that it can take
mozilla::Atomic as its storage type.

The Reference assignment operator is changed so that it results in a single
update rather than a separate read and update that would not be atomic when
used with atomic storage.

Tests are added for atomic storage and to exercise the assignment operator.

Differential Revision: https://phabricator.services.mozilla.com/D259230

mfbt/Atomics.h

@@ -65,7 +65,7 @@ namespace mozilla {
  * the modes we provide below, or not relevant for the CPUs we support
  * in Gecko.  These three modes are confusing enough as it is!
  */
-enum MemoryOrdering {
+enum MemoryOrdering : uint8_t {
   /*
    * Relaxed ordering is the simplest memory ordering: none at all.
    * When the result of a write is observed, nothing may be inferred

mfbt/BitSet.h

@@ -8,24 +8,47 @@
 #define mozilla_BitSet_h
 
 #include "mozilla/Array.h"
-#include "mozilla/ArrayUtils.h"
 #include "mozilla/MathAlgorithms.h"
-#include "mozilla/PodOperations.h"
 #include "mozilla/Span.h"
 
 namespace mozilla {
 
+enum MemoryOrdering : uint8_t;
+template <typename T, MemoryOrdering Order, typename Enable>
+class Atomic;
+
+namespace detail {
+
+template <typename T>
+struct UnwrapMaybeAtomic {
+  using Type = T;
+};
+template <typename T, MemoryOrdering Order, typename Enable>
+struct UnwrapMaybeAtomic<mozilla::Atomic<T, Order, Enable>> {
+  using Type = T;
+};
+
+}  // namespace detail
+
 /**
  * An object like std::bitset but which provides access to the underlying
  * storage.
  *
+ * The type |StorageType| must be an unsigned integer or a mozilla::Atomic
+ * wrapping an unsigned integer. Use of atomic types makes word access atomic,
+ * but does not make operations that operate on the whole bitset atomic.
+ *
  * The limited API is due to expedience only; feel free to flesh out any
  * std::bitset-like members.
  */
-template <size_t N, typename Word = size_t>
+template <size_t N, typename StorageType = size_t>
 class BitSet {
-  static_assert(std::is_unsigned_v<Word>,
+ public:
-                "The Word type must be an unsigned integral type");
+  using Word = typename detail::UnwrapMaybeAtomic<StorageType>::Type;
+  static_assert(sizeof(Word) == sizeof(StorageType));
+  static_assert(
+      std::is_unsigned_v<Word>,
+      "StorageType must be an unsigned integral type, or equivalent Atomic");
   static_assert(N != 0);
 
  private:
@@ -35,7 +58,7 @@ class BitSet {
   static constexpr Word kPaddingMask = Word(-1) >> kPaddingBits;
 
   // The zeroth bit in the bitset is the least significant bit of mStorage[0].
-  Array<Word, kNumWords> mStorage;
+  Array<StorageType, kNumWords> mStorage;
 
   constexpr void ResetPaddingBits() {
     if constexpr (kPaddingBits != 0) {
@@ -46,20 +69,24 @@ class BitSet {
  public:
   class Reference {
    public:
-    Reference(BitSet<N, Word>& aBitSet, size_t aPos)
+    Reference(BitSet<N, StorageType>& aBitSet, size_t aPos)
         : mBitSet(aBitSet), mPos(aPos) {}
 
     Reference& operator=(bool aValue) {
       auto bit = Word(1) << (mPos % kBitsPerWord);
       auto& word = mBitSet.mStorage[mPos / kBitsPerWord];
-      word = (word & ~bit) | (aValue ? bit : 0);
+      if (aValue) {
+        word |= bit;
+      } else {
+        word &= ~bit;
+      }
       return *this;
     }
 
     MOZ_IMPLICIT operator bool() const { return mBitSet.test(mPos); }
 
    private:
-    BitSet<N, Word>& mBitSet;
+    BitSet<N, StorageType>& mBitSet;
     size_t mPos;
   };
 
@@ -68,12 +95,16 @@ class BitSet {
   BitSet(const BitSet& aOther) { *this = aOther; }
 
   BitSet& operator=(const BitSet& aOther) {
-    PodCopy(mStorage.begin(), aOther.mStorage.begin(), kNumWords);
+    for (size_t i = 0; i < std::size(mStorage); i++) {
+      mStorage[i] = Word(aOther.mStorage[i]);
+    }
     return *this;
   }
 
-  explicit BitSet(Span<Word, kNumWords> aStorage) {
+  explicit BitSet(Span<StorageType, kNumWords> aStorage) {
-    PodCopy(mStorage.begin(), aStorage.Elements(), kNumWords);
+    for (size_t i = 0; i < std::size(mStorage); i++) {
+      mStorage[i] = Word(aStorage[i]);
+    }
   }
 
   static constexpr size_t size() { return N; }
@@ -84,7 +115,7 @@ class BitSet {
   }
 
   constexpr bool IsEmpty() const {
-    for (const Word& word : mStorage) {
+    for (const StorageType& word : mStorage) {
       if (word) {
         return false;
       }
@@ -101,13 +132,13 @@ class BitSet {
     return {*this, aPos};
   }
 
-  BitSet operator|(const BitSet<N, Word>& aOther) {
+  BitSet operator|(const BitSet<N, StorageType>& aOther) {
     BitSet result = *this;
     result |= aOther;
     return result;
   }
 
-  BitSet& operator|=(const BitSet<N, Word>& aOther) {
+  BitSet& operator|=(const BitSet<N, StorageType>& aOther) {
     for (size_t i = 0; i < std::size(mStorage); i++) {
       mStorage[i] |= aOther.mStorage[i];
     }
@@ -120,27 +151,27 @@ class BitSet {
     return result;
   }
 
-  BitSet& operator&=(const BitSet<N, Word>& aOther) {
+  BitSet& operator&=(const BitSet<N, StorageType>& aOther) {
     for (size_t i = 0; i < std::size(mStorage); i++) {
       mStorage[i] &= aOther.mStorage[i];
     }
     return *this;
   }
 
-  BitSet operator&(const BitSet<N, Word>& aOther) const {
+  BitSet operator&(const BitSet<N, StorageType>& aOther) const {
     BitSet result = *this;
     result &= aOther;
     return result;
   }
 
-  bool operator==(const BitSet<N, Word>& aOther) const {
+  bool operator==(const BitSet<N, StorageType>& aOther) const {
     return mStorage == aOther.mStorage;
   }
 
   size_t Count() const {
     size_t count = 0;
 
-    for (const Word& word : mStorage) {
+    for (const StorageType& word : mStorage) {
       if constexpr (kBitsPerWord > 32) {
         count += CountPopulation64(word);
       } else {
@@ -152,16 +183,22 @@ class BitSet {
   }
 
   // Set all bits to false.
-  void ResetAll() { PodArrayZero(mStorage); }
+  void ResetAll() {
+    for (StorageType& word : mStorage) {
+      word = Word(0);
+    }
+  }
 
   // Set all bits to true.
   void SetAll() {
-    memset(mStorage.begin(), 0xff, kNumWords * sizeof(Word));
+    for (StorageType& word : mStorage) {
+      word = ~Word(0);
+    }
     ResetPaddingBits();
   }
 
   void Flip() {
-    for (Word& word : mStorage) {
+    for (StorageType& word : mStorage) {
       word = ~word;
     }
 
@@ -217,9 +254,9 @@ class BitSet {
     return wordIndex * kBitsPerWord + pos;
   }
 
-  Span<Word> Storage() { return mStorage; }
+  Span<StorageType> Storage() { return mStorage; }
 
-  Span<const Word> Storage() const { return mStorage; }
+  Span<const StorageType> Storage() const { return mStorage; }
 };
 
 }  // namespace mozilla

mfbt/tests/TestBitSet.cpp

@@ -5,8 +5,10 @@
  * You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 #include "mozilla/Assertions.h"
+#include "mozilla/Atomics.h"
 #include "mozilla/BitSet.h"
 
+using mozilla::Atomic;
 using mozilla::BitSet;
 
 template <typename Storage>
@@ -14,9 +16,11 @@ class BitSetSuite {
   template <size_t N>
   using TestBitSet = BitSet<N, Storage>;
 
+  using Word = typename TestBitSet<1>::Word;
+
   static constexpr size_t kBitsPerWord = sizeof(Storage) * 8;
 
-  static constexpr Storage kAllBitsSet = ~Storage{0};
+  static constexpr Word kAllBitsSet = ~Word{0};
 
  public:
   void testLength() {
@@ -28,7 +32,7 @@ class BitSetSuite {
     MOZ_RELEASE_ASSERT(TestBitSet<kBitsPerWord + 1>().Storage().Length() == 2);
   }
 
-  void testConstruct() {
+  void testConstructAndAssign() {
     MOZ_RELEASE_ASSERT(TestBitSet<1>().Storage()[0] == 0);
     MOZ_RELEASE_ASSERT(TestBitSet<kBitsPerWord>().Storage()[0] == 0);
     MOZ_RELEASE_ASSERT(TestBitSet<kBitsPerWord + 1>().Storage()[0] == 0);
@@ -63,6 +67,18 @@ class BitSetSuite {
         kAllBitsSet);
     MOZ_RELEASE_ASSERT(
         TestBitSet<kBitsPerWord + 1>(bitsetW1.Storage()).Storage()[1] == 1);
+
+    TestBitSet<1> bitset1Copy;
+    bitset1Copy = bitset1;
+    TestBitSet<kBitsPerWord> bitsetWCopy;
+    bitsetWCopy = bitsetW;
+    TestBitSet<kBitsPerWord + 1> bitsetW1Copy;
+    bitsetW1Copy = bitsetW1;
+
+    MOZ_RELEASE_ASSERT(bitset1Copy.Storage()[0] == 1);
+    MOZ_RELEASE_ASSERT(bitsetWCopy.Storage()[0] == kAllBitsSet);
+    MOZ_RELEASE_ASSERT(bitsetW1Copy.Storage()[0] == kAllBitsSet);
+    MOZ_RELEASE_ASSERT(bitsetW1Copy.Storage()[1] == 1);
   }
 
   void testSetBit() {
@@ -164,7 +180,7 @@ class BitSetSuite {
 
   void runTests() {
     testLength();
-    testConstruct();
+    testConstructAndAssign();
     testSetBit();
     testFindBits();
   }
@@ -174,6 +190,8 @@ int main() {
   BitSetSuite<uint8_t>().runTests();
   BitSetSuite<uint32_t>().runTests();
   BitSetSuite<uint64_t>().runTests();
+  BitSetSuite<Atomic<uint32_t>>().runTests();
+  BitSetSuite<Atomic<uint64_t>>().runTests();
 
   return 0;
 }