From 2a4a4178b6b25d8633143361686fa283655a53f7 Mon Sep 17 00:00:00 2001
From: Tudor Bosman <tudorb@fb.com>
Date: Tue, 3 Jul 2012 17:02:05 -0700
Subject: [PATCH] Add multi-bit operations to folly/experimental/Bits.h

Summary: Add ability to set and get N consecutive bits from the bit sequence.

Test Plan: test added

Reviewed By: lucian@fb.com

FB internal diff: D511438
---
 folly/experimental/Bits.h            | 67 ++++++++++++++++++++++++++++
 folly/experimental/test/BitsTest.cpp | 26 +++++++++++
 2 files changed, 93 insertions(+)

diff --git a/folly/experimental/Bits.h b/folly/experimental/Bits.h
index 31b12376..249450eb 100644
--- a/folly/experimental/Bits.h
+++ b/folly/experimental/Bits.h
@@ -95,12 +95,34 @@ struct Bits {
    */
   static bool test(const T* p, size_t bit);
 
+  /**
+   * Set count contiguous bits starting at bitStart to the values
+   * from the least significant count bits of value; little endian.
+   * (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);
+
+  /**
+   * Get count contiguous bits starting at bitStart.
+   * Precondition: count <= sizeof(T) * 8
+   */
+  static T get(const T* p, size_t bitStart, size_t count);
+
   /**
    * Count the number of bits set in a range of blocks.
    */
   static size_t count(const T* begin, const T* end);
 
  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);
+
+  // 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 constexpr T one = T(1);
 };
 
@@ -119,6 +141,51 @@ inline bool Bits<T>::test(const T* p, size_t bit) {
   return 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) ||
+         (value & ~((one << count) - 1)) == 0);
+  size_t idx = blockIndex(bitStart);
+  size_t offset = bitOffset(bitStart);
+  if (offset + count <= bitsPerBlock) {
+    innerSet(p + idx, offset, count, value);
+  } else {
+    size_t countInThisBlock = bitsPerBlock - offset;
+    size_t countInNextBlock = count - countInThisBlock;
+    innerSet(p + idx, offset, countInThisBlock,
+             value & ((one << countInThisBlock) - 1));
+    innerSet(p + idx + 1, 0, countInNextBlock, value >> countInThisBlock);
+  }
+}
+
+template <class T>
+inline T Bits<T>::get(const T* p, size_t bitStart, size_t count) {
+  assert(count <= sizeof(T) * 8);
+  size_t idx = blockIndex(bitStart);
+  size_t offset = bitOffset(bitStart);
+  if (offset + count <= bitsPerBlock) {
+    return innerGet(p + idx, offset, 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);
+    return (nextBlockValue << countInThisBlock) | thisBlockValue;
+  }
+}
+
+template <class T>
+inline void Bits<T>::innerSet(T* p, size_t offset, size_t count, T value) {
+  // Mask out bits and set new value
+  *p = (*p & ~(((one << count) - 1) << offset)) | (value << offset);
+}
+
+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>
 inline size_t Bits<T>::count(const T* begin, const T* end) {
   size_t n = 0;
diff --git a/folly/experimental/test/BitsTest.cpp b/folly/experimental/test/BitsTest.cpp
index 32dad6a2..09eb4ae6 100644
--- a/folly/experimental/test/BitsTest.cpp
+++ b/folly/experimental/test/BitsTest.cpp
@@ -51,6 +51,32 @@ TEST(Bits, Simple) {
   EXPECT_EQ(2, Bits<uint8_t>::count(buf, buf + 256));
 }
 
+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]);
+}
+
 int main(int argc, char *argv[]) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);
-- 
2.34.1