Added methods to write small ints to bit streams.

Summary: Added BitStreamPutInt() and BitStreamGetInt() which take a stream of chars and can write integers of arbitrary bit sizes to that stream at arbitrary positions. There are also convenience versions of these functions that take std::strings and leveldb::Slices.

Test Plan: make check

Reviewers: sheki, vamsi, dhruba, emayanke

Reviewed By: vamsi

CC: leveldb

Differential Revision: https://reviews.facebook.net/D7071

util/coding.cc

@@ -4,6 +4,8 @@
 
 #include "util/coding.h"
 
+#include <algorithm>
+
 namespace leveldb {
 
 void EncodeFixed32(char* buf, uint32_t value) {
@@ -191,4 +193,102 @@ bool GetLengthPrefixedSlice(Slice* input, Slice* result) {
   }
 }
 
+void BitStreamPutInt(char* dst, size_t dstlen, size_t offset,
+                     uint32_t bits, uint64_t value) {
+  assert((offset + bits + 7)/8 <= dstlen);
+  assert(bits <= 64);
+
+  unsigned char* ptr = reinterpret_cast<unsigned char*>(dst);
+
+  size_t byteOffset = offset / 8;
+  size_t bitOffset = offset % 8;
+
+  // This prevents unused variable warnings when compiling.
+#ifndef NDEBUG
+  // Store truncated value.
+  uint64_t origValue = (bits < 64)?(value & (((uint64_t)1 << bits) - 1)):value;
+  uint32_t origBits = bits;
+#endif
+
+  while (bits > 0) {
+    size_t bitsToGet = std::min<size_t>(bits, 8 - bitOffset);
+    unsigned char mask = ((1 << bitsToGet) - 1);
+
+    ptr[byteOffset] = (ptr[byteOffset] & ~(mask << bitOffset)) +
+                      ((value & mask) << bitOffset);
+
+    value >>= bitsToGet;
+    byteOffset += 1;
+    bitOffset = 0;
+    bits -= bitsToGet;
+  }
+
+  assert(origValue == BitStreamGetInt(dst, dstlen, offset, origBits));
+}
+
+uint64_t BitStreamGetInt(const char* src, size_t srclen, size_t offset,
+                         uint32_t bits) {
+  assert((offset + bits + 7)/8 <= srclen);
+  assert(bits <= 64);
+
+  const unsigned char* ptr = reinterpret_cast<const unsigned char*>(src);
+
+  uint64_t result = 0;
+
+  size_t byteOffset = offset / 8;
+  size_t bitOffset = offset % 8;
+  size_t shift = 0;
+
+  while (bits > 0) {
+    size_t bitsToGet = std::min<size_t>(bits, 8 - bitOffset);
+    unsigned char mask = ((1 << bitsToGet) - 1);
+
+    result += (uint64_t)((ptr[byteOffset] >> bitOffset) & mask) << shift;
+
+    shift += bitsToGet;
+    byteOffset += 1;
+    bitOffset = 0;
+    bits -= bitsToGet;
+  }
+
+  return result;
+}
+
+void BitStreamPutInt(std::string* dst, size_t offset, uint32_t bits,
+                     uint64_t value) {
+  assert((offset + bits + 7)/8 <= dst->size());
+
+  const size_t kTmpBufLen = sizeof(value) + 1;
+  char tmpBuf[kTmpBufLen];
+
+  // Number of bytes of tmpBuf being used
+  const size_t kUsedBytes = (offset%8 + bits)/8;
+
+  // Copy relevant parts of dst to tmpBuf
+  for (size_t idx = 0; idx <= kUsedBytes; ++idx) {
+    tmpBuf[idx] = (*dst)[offset/8 + idx];
+  }
+
+  BitStreamPutInt(tmpBuf, kTmpBufLen, offset%8, bits, value);
+
+  // Copy tmpBuf back to dst
+  for (size_t idx = 0; idx <= kUsedBytes; ++idx) {
+    (*dst)[offset/8 + idx] = tmpBuf[idx];
+  }
+
+  // Do the check here too as we are working with a buffer.
+  assert(((bits < 64)?(value & (((uint64_t)1 << bits) - 1)):value) ==
+         BitStreamGetInt(dst, offset, bits));
+}
+
+uint64_t BitStreamGetInt(const std::string* src, size_t offset,
+                         uint32_t bits) {
+  return BitStreamGetInt(src->data(), src->size(), offset, bits);
+}
+
+uint64_t BitStreamGetInt(const Slice* src, size_t offset,
+                         uint32_t bits) {
+  return BitStreamGetInt(src->data(), src->size(), offset, bits);
+}
+
 }  // namespace leveldb