Make Bits<T> work with T = Unaligned<X> (X is unsigned integral type)

Test Plan: all folly tests

Reviewed By: andrei.alexandrescu@fb.com

FB internal diff: D517118

folly/Bits.h

@@ -530,10 +530,15 @@ BitIterator<BaseIter> findFirstSet(BitIterator<BaseIter> begin,
 
 template <class T, class Enable=void> struct Unaligned;
 
+/**
+ * Representation of an unaligned value of a POD type.
+ */
 template <class T>
 struct Unaligned<
-  T,
-  typename std::enable_if<std::is_pod<T>::value>::type> {
+    T,
+    typename std::enable_if<std::is_pod<T>::value>::type> {
+  Unaligned() { }  // uninitialized
+  /* implicit */ Unaligned(T v) : value(v) { }
   T value;
 } __attribute__((packed));
 
@@ -542,6 +547,7 @@ struct Unaligned<
  */
 template <class T>
 inline T loadUnaligned(const void* p) {
+  static_assert(sizeof(Unaligned<T>) == sizeof(T), "Invalid unaligned size");
   static_assert(alignof(Unaligned<T>) == 1, "Invalid alignment");
   return static_cast<const Unaligned<T>*>(p)->value;
 }
@@ -551,8 +557,9 @@ inline T loadUnaligned(const void* p) {
  */
 template <class T>
 inline void storeUnaligned(void* p, T value) {
+  static_assert(sizeof(Unaligned<T>) == sizeof(T), "Invalid unaligned size");
   static_assert(alignof(Unaligned<T>) == 1, "Invalid alignment");
-  static_cast<Unaligned<T>*>(p)->value = value;
+  new (p) Unaligned<T>(value);
 }
 
 }  // namespace folly

folly/experimental/Bits.h

@@ -21,10 +21,15 @@
 #include <type_traits>
 #include <limits>
 
+#include "folly/Bits.h"
 #include "folly/Range.h"
 
 namespace folly {
 
+// As a general rule, bit operations work on unsigned values only;
+// right-shift is arithmetic for signed values, and that can lead to
+// unpleasant bugs.
+
 /**
  * Population count (number of bits set), using __builtin_popcount or
  * __builtin_popcountll, depending on size.
@@ -50,14 +55,51 @@ inline typename std::enable_if<
   return __builtin_popcountll(x);
 }
 
+namespace detail {
+
+/**
+ * Helper class to make Bits<T> (below) work with both aligned values
+ * (T, where T is an unsigned integral type) or unaligned values
+ * (Unaligned<T>, where T is an unsigned integral type)
+ */
+template <class T, class Enable=void> struct BitsTraits;
+
+// Partial specialization for Unaligned<T>, where T is unsigned integral
 template <class T>
-struct Bits {
-  static_assert(std::is_integral<T>::value &&
-                std::is_unsigned<T>::value,
-                "Unsigned integral type required");
+struct BitsTraits<Unaligned<T>, typename std::enable_if<
+    (std::is_integral<T>::value && std::is_unsigned<T>::value)>::type> {
+  typedef T UnderlyingType;
+  static T load(const Unaligned<T>& x) { return x.value; }
+  static void store(Unaligned<T>& x, T v) { x.value = v; }
+};
+
+// Partial specialization for T, where T is unsigned integral
+template <class T>
+struct BitsTraits<T, typename std::enable_if<
+    (std::is_integral<T>::value && std::is_unsigned<T>::value)>::type> {
+  typedef T UnderlyingType;
+  static T load(const T& x) { return x; }
+  static void store(T& x, T v) { x = v; }
+};
+}  // namespace detail
 
+/**
+ * Wrapper class with static methods for various bit-level operations,
+ * treating an array of T as an array of bits (in little-endian order).
+ * (T is either an unsigned integral type or Unaligned<X>, where X is
+ * an unsigned integral type)
+ */
+template <class T, class Traits=detail::BitsTraits<T>>
+struct Bits {
+  typedef typename Traits::UnderlyingType UnderlyingType;
   typedef T type;
-  static constexpr size_t bitsPerBlock = std::numeric_limits<T>::digits;
+  static_assert(sizeof(T) == sizeof(UnderlyingType), "Size mismatch");
+
+  /**
+   * Number of bits in a block.
+   */
+  static constexpr size_t bitsPerBlock =
+    std::numeric_limits<UnderlyingType>::digits;
 
   /**
    * Byte index of the given bit.
@@ -101,13 +143,13 @@ struct Bits {
    * (value & 1 becomes the bit at bitStart, etc)
    * Precondition: count <= sizeof(T) * 8
    */
-  static void set(T* p, size_t bitStart, size_t count, T value);
+  static void set(T* p, size_t bitStart, size_t count, UnderlyingType value);
 
   /**
    * Get count contiguous bits starting at bitStart.
    * Precondition: count <= sizeof(T) * 8
    */
-  static T get(const T* p, size_t bitStart, size_t count);
+  static UnderlyingType get(const T* p, size_t bitStart, size_t count);
 
   /**
    * Count the number of bits set in a range of blocks.
@@ -117,34 +159,38 @@ struct Bits {
  private:
   // Same as set, assumes all bits are in the same block.
   // (bitStart < sizeof(T) * 8, bitStart + count <= sizeof(T) * 8)
-  static void innerSet(T* p, size_t bitStart, size_t count, T value);
+  static void innerSet(T* p, size_t bitStart, size_t count,
+                       UnderlyingType value);
 
   // Same as get, assumes all bits are in the same block.
   // (bitStart < sizeof(T) * 8, bitStart + count <= sizeof(T) * 8)
-  static T innerGet(const T* p, size_t bitStart, size_t count);
+  static UnderlyingType innerGet(const T* p, size_t bitStart, size_t count);
 
-  static constexpr T one = T(1);
+  static constexpr UnderlyingType one = UnderlyingType(1);
 };
 
-template <class T>
-inline void Bits<T>::set(T* p, size_t bit) {
-  p[blockIndex(bit)] |= (one << bitOffset(bit));
+template <class T, class Traits>
+inline void Bits<T, Traits>::set(T* p, size_t bit) {
+  T& block = p[blockIndex(bit)];
+  Traits::store(block, Traits::load(block) | (one << bitOffset(bit)));
 }
 
-template <class T>
-inline void Bits<T>::clear(T* p, size_t bit) {
-  p[blockIndex(bit)] &= ~(one << bitOffset(bit));
+template <class T, class Traits>
+inline void Bits<T, Traits>::clear(T* p, size_t bit) {
+  T& block = p[blockIndex(bit)];
+  Traits::store(block, Traits::load(block) & ~(one << bitOffset(bit)));
 }
 
-template <class T>
-inline bool Bits<T>::test(const T* p, size_t bit) {
-  return p[blockIndex(bit)] & (one << bitOffset(bit));
+template <class T, class Traits>
+inline bool Bits<T, Traits>::test(const T* p, size_t bit) {
+  return Traits::load(p[blockIndex(bit)]) & (one << bitOffset(bit));
 }
 
-template <class T>
-inline void Bits<T>::set(T* p, size_t bitStart, size_t count, T value) {
-  assert(count <= sizeof(T) * 8);
-  assert(count == sizeof(T) ||
+template <class T, class Traits>
+inline void Bits<T, Traits>::set(T* p, size_t bitStart, size_t count,
+                                 UnderlyingType value) {
+  assert(count <= sizeof(UnderlyingType) * 8);
+  assert(count == sizeof(UnderlyingType) ||
          (value & ~((one << count) - 1)) == 0);
   size_t idx = blockIndex(bitStart);
   size_t offset = bitOffset(bitStart);
@@ -159,9 +205,10 @@ inline void Bits<T>::set(T* p, size_t bitStart, size_t count, T value) {
   }
 }
 
-template <class T>
-inline T Bits<T>::get(const T* p, size_t bitStart, size_t count) {
-  assert(count <= sizeof(T) * 8);
+template <class T, class Traits>
+inline auto Bits<T, Traits>::get(const T* p, size_t bitStart, size_t count)
+  -> UnderlyingType {
+  assert(count <= sizeof(UnderlyingType) * 8);
   size_t idx = blockIndex(bitStart);
   size_t offset = bitOffset(bitStart);
   if (offset + count <= bitsPerBlock) {
@@ -169,28 +216,33 @@ inline T Bits<T>::get(const T* p, size_t bitStart, size_t count) {
   } else {
     size_t countInThisBlock = bitsPerBlock - offset;
     size_t countInNextBlock = count - countInThisBlock;
-    T thisBlockValue = innerGet(p + idx, offset, countInThisBlock);
-    T nextBlockValue = innerGet(p + idx + 1, 0, countInNextBlock);
+    UnderlyingType thisBlockValue = innerGet(p + idx, offset, countInThisBlock);
+    UnderlyingType nextBlockValue = innerGet(p + idx + 1, 0, countInNextBlock);
     return (nextBlockValue << countInThisBlock) | thisBlockValue;
   }
 }
 
-template <class T>
-inline void Bits<T>::innerSet(T* p, size_t offset, size_t count, T value) {
+template <class T, class Traits>
+inline void Bits<T, Traits>::innerSet(T* p, size_t offset, size_t count,
+                                      UnderlyingType value) {
   // Mask out bits and set new value
-  *p = (*p & ~(((one << count) - 1) << offset)) | (value << offset);
+  UnderlyingType v = Traits::load(*p);
+  v &= ~(((one << count) - 1) << offset);
+  v |= (value << offset);
+  Traits::store(*p, v);
 }
 
-template <class T>
-inline T Bits<T>::innerGet(const T* p, size_t offset, size_t count) {
-  return (*p >> offset) & ((one << count) - 1);
+template <class T, class Traits>
+inline auto Bits<T, Traits>::innerGet(const T* p, size_t offset, size_t count)
+  -> UnderlyingType {
+  return (Traits::load(*p) >> offset) & ((one << count) - 1);
 }
 
-template <class T>
-inline size_t Bits<T>::count(const T* begin, const T* end) {
+template <class T, class Traits>
+inline size_t Bits<T, Traits>::count(const T* begin, const T* end) {
   size_t n = 0;
   for (; begin != end; ++begin) {
-    n += popcount(*begin);
+    n += popcount(Traits::load(*begin));
   }
   return n;
 }

folly/experimental/test/BitsTest.cpp

@@ -21,60 +21,149 @@
 
 using namespace folly;
 
-TEST(Bits, Simple) {
-  EXPECT_EQ(0, Bits<uint8_t>::blockCount(0));
-  EXPECT_EQ(1, Bits<uint8_t>::blockCount(1));
-  EXPECT_EQ(1, Bits<uint8_t>::blockCount(8));
-  EXPECT_EQ(2, Bits<uint8_t>::blockCount(9));
-  EXPECT_EQ(256, Bits<uint8_t>::blockCount(2048));
-  EXPECT_EQ(257, Bits<uint8_t>::blockCount(2049));
-
-  EXPECT_EQ(4, Bits<uint8_t>::blockIndex(39));
-  EXPECT_EQ(7, Bits<uint8_t>::bitOffset(39));
-  EXPECT_EQ(5, Bits<uint8_t>::blockIndex(40));
-  EXPECT_EQ(0, Bits<uint8_t>::bitOffset(40));
-
-  uint8_t buf[256];
-  memset(buf, 0, 256);
-
-  Bits<uint8_t>::set(buf, 36);
-  Bits<uint8_t>::set(buf, 39);
-  EXPECT_EQ((1 << 7) | (1 << 4), buf[4]);
-  EXPECT_EQ(0, buf[5]);
-  Bits<uint8_t>::clear(buf, 39);
-  EXPECT_EQ(1 << 4, buf[4]);
-  EXPECT_EQ(0, buf[5]);
-  Bits<uint8_t>::set(buf, 40);
-  EXPECT_EQ(1 << 4, buf[4]);
-  EXPECT_EQ(1, buf[5]);
-
-  EXPECT_EQ(2, Bits<uint8_t>::count(buf, buf + 256));
+template <class T>
+void runSimpleTest8() {
+  auto load = detail::BitsTraits<T>::load;
+
+  EXPECT_EQ(0, Bits<T>::blockCount(0));
+  EXPECT_EQ(1, Bits<T>::blockCount(1));
+  EXPECT_EQ(1, Bits<T>::blockCount(8));
+  EXPECT_EQ(2, Bits<T>::blockCount(9));
+  EXPECT_EQ(256, Bits<T>::blockCount(2048));
+  EXPECT_EQ(257, Bits<T>::blockCount(2049));
+
+  EXPECT_EQ(4, Bits<T>::blockIndex(39));
+  EXPECT_EQ(7, Bits<T>::bitOffset(39));
+  EXPECT_EQ(5, Bits<T>::blockIndex(40));
+  EXPECT_EQ(0, Bits<T>::bitOffset(40));
+
+  T buf[256];
+  std::fill(buf, buf + 256, T(0));
+
+  Bits<T>::set(buf, 36);
+  Bits<T>::set(buf, 39);
+  EXPECT_EQ((1 << 7) | (1 << 4), load(buf[4]));
+  EXPECT_EQ(0, load(buf[5]));
+  Bits<T>::clear(buf, 39);
+  EXPECT_EQ(1 << 4, load(buf[4]));
+  EXPECT_EQ(0, load(buf[5]));
+  Bits<T>::set(buf, 40);
+  EXPECT_EQ(1 << 4, load(buf[4]));
+  EXPECT_EQ(1, load(buf[5]));
+
+  EXPECT_EQ(2, Bits<T>::count(buf, buf + 256));
+}
+
+TEST(Bits, Simple8) {
+  runSimpleTest8<uint8_t>();
+}
+
+TEST(Bits, SimpleUnaligned8) {
+  runSimpleTest8<Unaligned<uint8_t>>();
+}
+
+template <class T>
+void runSimpleTest64() {
+  auto load = detail::BitsTraits<T>::load;
+
+  EXPECT_EQ(0, Bits<T>::blockCount(0));
+  EXPECT_EQ(1, Bits<T>::blockCount(1));
+  EXPECT_EQ(1, Bits<T>::blockCount(8));
+  EXPECT_EQ(1, Bits<T>::blockCount(9));
+  EXPECT_EQ(1, Bits<T>::blockCount(64));
+  EXPECT_EQ(2, Bits<T>::blockCount(65));
+  EXPECT_EQ(32, Bits<T>::blockCount(2048));
+  EXPECT_EQ(33, Bits<T>::blockCount(2049));
+
+  EXPECT_EQ(0, Bits<T>::blockIndex(39));
+  EXPECT_EQ(39, Bits<T>::bitOffset(39));
+  EXPECT_EQ(4, Bits<T>::blockIndex(319));
+  EXPECT_EQ(63, Bits<T>::bitOffset(319));
+  EXPECT_EQ(5, Bits<T>::blockIndex(320));
+  EXPECT_EQ(0, Bits<T>::bitOffset(320));
+
+  T buf[256];
+  std::fill(buf, buf + 256, T(0));
+
+  Bits<T>::set(buf, 300);
+  Bits<T>::set(buf, 319);
+  EXPECT_EQ((uint64_t(1) << 44) | (uint64_t(1) << 63), load(buf[4]));
+  EXPECT_EQ(0, load(buf[5]));
+  Bits<T>::clear(buf, 319);
+  EXPECT_EQ(uint64_t(1) << 44, load(buf[4]));
+  EXPECT_EQ(0, load(buf[5]));
+  Bits<T>::set(buf, 320);
+  EXPECT_EQ(uint64_t(1) << 44, load(buf[4]));
+  EXPECT_EQ(1, load(buf[5]));
+
+  EXPECT_EQ(2, Bits<T>::count(buf, buf + 256));
+}
+
+TEST(Bits, Simple64) {
+  runSimpleTest64<uint64_t>();
+}
+
+TEST(Bits, SimpleUnaligned64) {
+  runSimpleTest64<Unaligned<uint64_t>>();
+}
+
+template <class T>
+void runMultiBitTest8() {
+  auto load = detail::BitsTraits<T>::load;
+  T buf[] = {0x12, 0x34, 0x56, 0x78};
+
+  EXPECT_EQ(0x02, load(Bits<T>::get(buf, 0, 4)));
+  EXPECT_EQ(0x1a, load(Bits<T>::get(buf, 9, 5)));
+  EXPECT_EQ(0xb1, load(Bits<T>::get(buf, 13, 8)));
+
+  Bits<T>::set(buf, 0, 4, 0x0b);
+  EXPECT_EQ(0x1b, load(buf[0]));
+  EXPECT_EQ(0x34, load(buf[1]));
+  EXPECT_EQ(0x56, load(buf[2]));
+  EXPECT_EQ(0x78, load(buf[3]));
+
+  Bits<T>::set(buf, 9, 5, 0x0e);
+  EXPECT_EQ(0x1b, load(buf[0]));
+  EXPECT_EQ(0x1c, load(buf[1]));
+  EXPECT_EQ(0x56, load(buf[2]));
+  EXPECT_EQ(0x78, load(buf[3]));
+
+  Bits<T>::set(buf, 13, 8, 0xaa);
+  EXPECT_EQ(0x1b, load(buf[0]));
+  EXPECT_EQ(0x5c, load(buf[1]));
+  EXPECT_EQ(0x55, load(buf[2]));
+  EXPECT_EQ(0x78, load(buf[3]));
+}
+
+TEST(Bits, MultiBit8) {
+  runMultiBitTest8<uint8_t>();
+}
+
+TEST(Bits, MultiBitUnaligned8) {
+  runMultiBitTest8<Unaligned<uint8_t>>();
+}
+
+template <class T>
+void runMultiBitTest64() {
+  auto load = detail::BitsTraits<T>::load;
+  T buf[] = {0x123456789abcdef0, 0x13579bdf2468ace0};
+
+  EXPECT_EQ(0xf0, load(Bits<T>::get(buf, 0, 8)));
+  EXPECT_EQ(0x89abcdef, load(Bits<T>::get(buf, 4, 32)));
+  EXPECT_EQ(0x189abcdef, load(Bits<T>::get(buf, 4, 33)));
+
+  Bits<T>::set(buf, 4, 31, 0x55555555);
+  EXPECT_EQ(0xd5555555, load(Bits<T>::get(buf, 4, 32)));
+  EXPECT_EQ(0x1d5555555, load(Bits<T>::get(buf, 4, 33)));
+  EXPECT_EQ(0xd55555550, load(Bits<T>::get(buf, 0, 36)));
+}
+
+TEST(Bits, MultiBit64) {
+  runMultiBitTest64<uint64_t>();
 }
 
-TEST(Bits, MultiBit) {
-  uint8_t buf[] = {0x12, 0x34, 0x56, 0x78};
-
-  EXPECT_EQ(0x02, Bits<uint8_t>::get(buf, 0, 4));
-  EXPECT_EQ(0x1a, Bits<uint8_t>::get(buf, 9, 5));
-  EXPECT_EQ(0xb1, Bits<uint8_t>::get(buf, 13, 8));
-
-  Bits<uint8_t>::set(buf, 0, 4, 0x0b);
-  EXPECT_EQ(0x1b, buf[0]);
-  EXPECT_EQ(0x34, buf[1]);
-  EXPECT_EQ(0x56, buf[2]);
-  EXPECT_EQ(0x78, buf[3]);
-
-  Bits<uint8_t>::set(buf, 9, 5, 0x0e);
-  EXPECT_EQ(0x1b, buf[0]);
-  EXPECT_EQ(0x1c, buf[1]);
-  EXPECT_EQ(0x56, buf[2]);
-  EXPECT_EQ(0x78, buf[3]);
-
-  Bits<uint8_t>::set(buf, 13, 8, 0xaa);
-  EXPECT_EQ(0x1b, buf[0]);
-  EXPECT_EQ(0x5c, buf[1]);
-  EXPECT_EQ(0x55, buf[2]);
-  EXPECT_EQ(0x78, buf[3]);
+TEST(Bits, MultiBitUnaligned64) {
+  runMultiBitTest64<Unaligned<uint64_t>>();
 }
 
 int main(int argc, char *argv[]) {