Fixed-size atomic bitset

Summary: No magic; what it says on the package.

Test Plan: atomic_bitset_test

Reviewed By: delong.j@fb.com

FB internal diff: D971875

folly/AtomicBitSet.h

+/*
+ * Copyright 2013 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FOLLY_ATOMICBITSET_H_
+#define FOLLY_ATOMICBITSET_H_
+
+#include <array>
+#include <atomic>
+#include <cassert>
+#include <cstddef>
+#include <limits>
+
+#include <boost/noncopyable.hpp>
+
+namespace folly {
+
+/**
+ * An atomic bitset of fixed size (specified at compile time).
+ */
+template <size_t N>
+class AtomicBitSet : private boost::noncopyable {
+ public:
+  /**
+   * Construct an AtomicBitSet; all bits are initially false.
+   */
+  AtomicBitSet();
+
+  /**
+   * Set bit idx to true, using the given memory order. Returns the
+   * previous value of the bit.
+   *
+   * Note that the operation is a read-modify-write operation due to the use
+   * of fetch_or.
+   */
+  bool set(size_t idx, std::memory_order order = std::memory_order_seq_cst);
+
+  /**
+   * Set bit idx to false, using the given memory order. Returns the
+   * previous value of the bit.
+   *
+   * Note that the operation is a read-modify-write operation due to the use
+   * of fetch_and.
+   */
+  bool reset(size_t idx, std::memory_order order = std::memory_order_seq_cst);
+
+  /**
+   * Set bit idx to the given value, using the given memory order. Returns
+   * the previous value of the bit.
+   *
+   * Note that the operation is a read-modify-write operation due to the use
+   * of fetch_and or fetch_or.
+   *
+   * Yes, this is an overload of set(), to keep as close to std::bitset's
+   * interface as possible.
+   */
+  bool set(size_t idx,
+           bool value,
+           std::memory_order order = std::memory_order_seq_cst);
+
+  /**
+   * Read bit idx.
+   */
+  bool test(size_t idx,
+            std::memory_order order = std::memory_order_seq_cst) const;
+
+  /**
+   * Same as test() with the default memory order.
+   */
+  bool operator[](size_t idx) const;
+
+  /**
+   * Return the size of the bitset.
+   */
+  constexpr size_t size() const {
+    return N;
+  }
+
+ private:
+  // Pick the largest lock-free type available
+#if (ATOMIC_LLONG_LOCK_FREE == 2)
+  typedef unsigned long long BlockType;
+#elif (ATOMIC_LONG_LOCK_FREE == 2)
+  typedef unsigned long BlockType;
+#else
+  // Even if not lock free, what can we do?
+  typedef unsigned int BlockType;
+#endif
+  typedef std::atomic<BlockType> AtomicBlockType;
+
+  static constexpr size_t kBitsPerBlock =
+    std::numeric_limits<BlockType>::digits;
+
+  static constexpr size_t blockIndex(size_t bit) {
+    return bit / kBitsPerBlock;
+  }
+
+  static constexpr size_t bitOffset(size_t bit) {
+    return bit % kBitsPerBlock;
+  }
+
+  // avoid casts
+  static constexpr BlockType kOne = 1;
+
+  std::array<AtomicBlockType, N> data_;
+};
+
+// value-initialize to zero
+template <size_t N>
+inline AtomicBitSet<N>::AtomicBitSet() : data_() {
+}
+
+template <size_t N>
+inline bool AtomicBitSet<N>::set(size_t idx, std::memory_order order) {
+  assert(idx < N * kBitsPerBlock);
+  BlockType mask = kOne << bitOffset(idx);
+  return data_[blockIndex(idx)].fetch_or(mask, order) & mask;
+}
+
+template <size_t N>
+inline bool AtomicBitSet<N>::reset(size_t idx, std::memory_order order) {
+  assert(idx < N * kBitsPerBlock);
+  BlockType mask = kOne << bitOffset(idx);
+  return data_[blockIndex(idx)].fetch_and(~mask, order) & mask;
+}
+
+template <size_t N>
+inline bool AtomicBitSet<N>::set(size_t idx,
+                                 bool value,
+                                 std::memory_order order) {
+  return value ? set(idx, order) : reset(idx, order);
+}
+
+template <size_t N>
+inline bool AtomicBitSet<N>::test(size_t idx, std::memory_order order) const {
+  assert(idx < N * kBitsPerBlock);
+  BlockType mask = kOne << bitOffset(idx);
+  return data_[blockIndex(idx)].load(order) & mask;
+}
+
+template <size_t N>
+inline bool AtomicBitSet<N>::operator[](size_t idx) const {
+  return test(idx);
+}
+
+}  // namespaces
+
+#endif /* FOLLY_ATOMICBITSET_H_ */
+

folly/test/AtomicBitSetTest.cpp

+/*
+ * Copyright 2013 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "folly/AtomicBitSet.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+namespace folly { namespace test {
+
+TEST(AtomicBitSet, Simple) {
+  constexpr size_t kSize = 1000;
+  AtomicBitSet<kSize> bs;
+
+  EXPECT_EQ(kSize, bs.size());
+
+  for (size_t i = 0; i < kSize; ++i) {
+    EXPECT_FALSE(bs[i]);
+  }
+
+  bs.set(42);
+  for (size_t i = 0; i < kSize; ++i) {
+    EXPECT_EQ(i == 42, bs[i]);
+  }
+
+  bs.set(43);
+  for (size_t i = 0; i < kSize; ++i) {
+    EXPECT_EQ((i == 42 || i == 43), bs[i]);
+  }
+
+  bs.reset(42);
+  for (size_t i = 0; i < kSize; ++i) {
+    EXPECT_EQ((i == 43), bs[i]);
+  }
+
+  bs.reset(43);
+  for (size_t i = 0; i < kSize; ++i) {
+    EXPECT_FALSE(bs[i]);
+  }
+}
+
+}}  // namespaces
+
+int main(int argc, char *argv[]) {
+  testing::InitGoogleTest(&argc, argv);
+  google::ParseCommandLineFlags(&argc, &argv, true);
+  return RUN_ALL_TESTS();
+}
+