folly/GroupVarint.h

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #if !defined(__GNUC__) && !defined(_MSC_VER)
  20 #error GroupVarint.h requires GCC or MSVC
  21 #endif
  22
  23 #include <folly/Portability.h>
  24
  25 #if FOLLY_X64 || defined(__i386__) || FOLLY_PPC64 || FOLLY_A64
  26 #define HAVE_GROUP_VARINT 1
  27
  28 #include <cstdint>
  29 #include <limits>
  30 #include <folly/detail/GroupVarintDetail.h>
  31 #include <folly/Bits.h>
  32 #include <folly/Range.h>
  33 #include <folly/portability/Builtins.h>
  34 #include <glog/logging.h>
  35
  36 #if FOLLY_SSE >= 3
  37 #include <nmmintrin.h>
  38 namespace folly {
  39 namespace detail {
  40 alignas(16) extern const uint64_t groupVarintSSEMasks[];
  41 }  // namespace detail
  42 }  // namespace folly
  43 #endif
  44
  45 namespace folly {
  46 namespace detail {
  47 extern const uint8_t groupVarintLengths[];
  48 }  // namespace detail
  49 }  // namespace folly
  50
  51 namespace folly {
  52
  53 template <typename T>
  54 class GroupVarint;
  55
  56 /**
  57  * GroupVarint encoding for 32-bit values.
  58  *
  59  * Encodes 4 32-bit integers at once, each using 1-4 bytes depending on size.
  60  * There is one byte of overhead.  (The first byte contains the lengths of
  61  * the four integers encoded as two bits each; 00=1 byte .. 11=4 bytes)
  62  *
  63  * This implementation assumes little-endian and does unaligned 32-bit
  64  * accesses, so it's basically not portable outside of the x86[_64] world.
  65  */
  66 template <>
  67 class GroupVarint<uint32_t> : public detail::GroupVarintBase<uint32_t> {
  68  public:
  69
  70   /**
  71    * Return the number of bytes used to encode these four values.
  72    */
  73   static size_t size(uint32_t a, uint32_t b, uint32_t c, uint32_t d) {
  74     return kHeaderSize + kGroupSize + key(a) + key(b) + key(c) + key(d);
  75   }
  76
  77   /**
  78    * Return the number of bytes used to encode four uint32_t values stored
  79    * at consecutive positions in an array.
  80    */
  81   static size_t size(const uint32_t* p) {
  82     return size(p[0], p[1], p[2], p[3]);
  83   }
  84
  85   /**
  86    * Return the number of bytes used to encode count (<= 4) values.
  87    * If you clip a buffer after these many bytes, you can still decode
  88    * the first "count" values correctly (if the remaining size() -
  89    * partialSize() bytes are filled with garbage).
  90    */
  91   static size_t partialSize(const type* p, size_t count) {
  92     DCHECK_LE(count, kGroupSize);
  93     size_t s = kHeaderSize + count;
  94     for (; count; --count, ++p) {
  95       s += key(*p);
  96     }
  97     return s;
  98   }
  99
 100   /**
 101    * Return the number of values from *p that are valid from an encoded
 102    * buffer of size bytes.
 103    */
 104   static size_t partialCount(const char* p, size_t size) {
 105     uint8_t v = uint8_t(*p);
 106     size_t s = kHeaderSize;
 107     s += 1 + b0key(v);
 108     if (s > size) return 0;
 109     s += 1 + b1key(v);
 110     if (s > size) return 1;
 111     s += 1 + b2key(v);
 112     if (s > size) return 2;
 113     s += 1 + b3key(v);
 114     if (s > size) return 3;
 115     return 4;
 116   }
 117
 118   /**
 119    * Given a pointer to the beginning of an GroupVarint32-encoded block,
 120    * return the number of bytes used by the encoding.
 121    */
 122   static size_t encodedSize(const char* p) {
 123     return kHeaderSize + kGroupSize +
 124            b0key(uint8_t(*p)) + b1key(uint8_t(*p)) +
 125            b2key(uint8_t(*p)) + b3key(uint8_t(*p));
 126   }
 127
 128   /**
 129    * Encode four uint32_t values into the buffer pointed-to by p, and return
 130    * the next position in the buffer (that is, one character past the last
 131    * encoded byte).  p needs to have at least size()+4 bytes available.
 132    */
 133   static char* encode(char* p, uint32_t a, uint32_t b, uint32_t c, uint32_t d) {
 134     uint8_t b0key = key(a);
 135     uint8_t b1key = key(b);
 136     uint8_t b2key = key(c);
 137     uint8_t b3key = key(d);
 138     *p++ = (b3key << 6) | (b2key << 4) | (b1key << 2) | b0key;
 139     storeUnaligned(p, a);
 140     p += b0key+1;
 141     storeUnaligned(p, b);
 142     p += b1key+1;
 143     storeUnaligned(p, c);
 144     p += b2key+1;
 145     storeUnaligned(p, d);
 146     p += b3key+1;
 147     return p;
 148   }
 149
 150   /**
 151    * Encode four uint32_t values from the array pointed-to by src into the
 152    * buffer pointed-to by p, similar to encode(p,a,b,c,d) above.
 153    */
 154   static char* encode(char* p, const uint32_t* src) {
 155     return encode(p, src[0], src[1], src[2], src[3]);
 156   }
 157
 158   /**
 159    * Decode four uint32_t values from a buffer, and return the next position
 160    * in the buffer (that is, one character past the last encoded byte).
 161    * The buffer needs to have at least 3 extra bytes available (they
 162    * may be read but ignored).
 163    */
 164   static const char* decode_simple(const char* p, uint32_t* a, uint32_t* b,
 165                                    uint32_t* c, uint32_t* d) {
 166     size_t k = loadUnaligned<uint8_t>(p);
 167     const char* end = p + detail::groupVarintLengths[k];
 168     ++p;
 169     size_t k0 = b0key(k);
 170     *a = loadUnaligned<uint32_t>(p) & kMask[k0];
 171     p += k0+1;
 172     size_t k1 = b1key(k);
 173     *b = loadUnaligned<uint32_t>(p) & kMask[k1];
 174     p += k1+1;
 175     size_t k2 = b2key(k);
 176     *c = loadUnaligned<uint32_t>(p) & kMask[k2];
 177     p += k2+1;
 178     size_t k3 = b3key(k);
 179     *d = loadUnaligned<uint32_t>(p) & kMask[k3];
 180     // p += k3+1;
 181     return end;
 182   }
 183
 184   /**
 185    * Decode four uint32_t values from a buffer and store them in the array
 186    * pointed-to by dest, similar to decode(p,a,b,c,d) above.
 187    */
 188   static const char* decode_simple(const char* p, uint32_t* dest) {
 189     return decode_simple(p, dest, dest+1, dest+2, dest+3);
 190   }
 191
 192 #if FOLLY_SSE >= 3
 193   /**
 194    * Just like the non-SSSE3 decode below, but with the additional constraint
 195    * that we must be able to read at least 17 bytes from the input pointer, p.
 196    */
 197   static const char* decode(const char* p, uint32_t* dest) {
 198     uint8_t key = uint8_t(p[0]);
 199     __m128i val = _mm_loadu_si128((const __m128i*)(p+1));
 200     __m128i mask =
 201         _mm_load_si128((const __m128i*)&detail::groupVarintSSEMasks[key * 2]);
 202     __m128i r = _mm_shuffle_epi8(val, mask);
 203     _mm_storeu_si128((__m128i*)dest, r);
 204     return p + detail::groupVarintLengths[key];
 205   }
 206
 207   /**
 208    * Just like decode_simple, but with the additional constraint that
 209    * we must be able to read at least 17 bytes from the input pointer, p.
 210    */
 211   static const char* decode(const char* p, uint32_t* a, uint32_t* b,
 212                             uint32_t* c, uint32_t* d) {
 213     uint8_t key = uint8_t(p[0]);
 214     __m128i val = _mm_loadu_si128((const __m128i*)(p+1));
 215     __m128i mask =
 216         _mm_load_si128((const __m128i*)&detail::groupVarintSSEMasks[key * 2]);
 217     __m128i r = _mm_shuffle_epi8(val, mask);
 218
 219     // Extracting 32 bits at a time out of an XMM register is a SSE4 feature
 220 #if FOLLY_SSE >= 4
 221     *a = uint32_t(_mm_extract_epi32(r, 0));
 222     *b = uint32_t(_mm_extract_epi32(r, 1));
 223     *c = uint32_t(_mm_extract_epi32(r, 2));
 224     *d = uint32_t(_mm_extract_epi32(r, 3));
 225 #else  /* !__SSE4__ */
 226     *a = _mm_extract_epi16(r, 0) + (_mm_extract_epi16(r, 1) << 16);
 227     *b = _mm_extract_epi16(r, 2) + (_mm_extract_epi16(r, 3) << 16);
 228     *c = _mm_extract_epi16(r, 4) + (_mm_extract_epi16(r, 5) << 16);
 229     *d = _mm_extract_epi16(r, 6) + (_mm_extract_epi16(r, 7) << 16);
 230 #endif  /* __SSE4__ */
 231
 232     return p + detail::groupVarintLengths[key];
 233   }
 234
 235 #else  /* !__SSSE3__ */
 236   static const char* decode(const char* p, uint32_t* a, uint32_t* b,
 237                             uint32_t* c, uint32_t* d) {
 238     return decode_simple(p, a, b, c, d);
 239   }
 240
 241   static const char* decode(const char* p, uint32_t* dest) {
 242     return decode_simple(p, dest);
 243   }
 244 #endif  /* __SSSE3__ */
 245
 246  private:
 247   static uint8_t key(uint32_t x) {
 248     // __builtin_clz is undefined for the x==0 case
 249     return uint8_t(3 - (__builtin_clz(x | 1) / 8));
 250   }
 251   static size_t b0key(size_t x) { return x & 3; }
 252   static size_t b1key(size_t x) { return (x >> 2) & 3; }
 253   static size_t b2key(size_t x) { return (x >> 4) & 3; }
 254   static size_t b3key(size_t x) { return (x >> 6) & 3; }
 255
 256   static const uint32_t kMask[];
 257 };
 258
 259
 260 /**
 261  * GroupVarint encoding for 64-bit values.
 262  *
 263  * Encodes 5 64-bit integers at once, each using 1-8 bytes depending on size.
 264  * There are two bytes of overhead.  (The first two bytes contain the lengths
 265  * of the five integers encoded as three bits each; 000=1 byte .. 111 = 8 bytes)
 266  *
 267  * This implementation assumes little-endian and does unaligned 64-bit
 268  * accesses, so it's basically not portable outside of the x86[_64] world.
 269  */
 270 template <>
 271 class GroupVarint<uint64_t> : public detail::GroupVarintBase<uint64_t> {
 272  public:
 273   /**
 274    * Return the number of bytes used to encode these five values.
 275    */
 276   static size_t size(uint64_t a, uint64_t b, uint64_t c, uint64_t d,
 277                      uint64_t e) {
 278     return kHeaderSize + kGroupSize +
 279            key(a) + key(b) + key(c) + key(d) + key(e);
 280   }
 281
 282   /**
 283    * Return the number of bytes used to encode five uint64_t values stored
 284    * at consecutive positions in an array.
 285    */
 286   static size_t size(const uint64_t* p) {
 287     return size(p[0], p[1], p[2], p[3], p[4]);
 288   }
 289
 290   /**
 291    * Return the number of bytes used to encode count (<= 4) values.
 292    * If you clip a buffer after these many bytes, you can still decode
 293    * the first "count" values correctly (if the remaining size() -
 294    * partialSize() bytes are filled with garbage).
 295    */
 296   static size_t partialSize(const type* p, size_t count) {
 297     DCHECK_LE(count, kGroupSize);
 298     size_t s = kHeaderSize + count;
 299     for (; count; --count, ++p) {
 300       s += key(*p);
 301     }
 302     return s;
 303   }
 304
 305   /**
 306    * Return the number of values from *p that are valid from an encoded
 307    * buffer of size bytes.
 308    */
 309   static size_t partialCount(const char* p, size_t size) {
 310     uint16_t v = loadUnaligned<uint16_t>(p);
 311     size_t s = kHeaderSize;
 312     s += 1 + b0key(v);
 313     if (s > size) return 0;
 314     s += 1 + b1key(v);
 315     if (s > size) return 1;
 316     s += 1 + b2key(v);
 317     if (s > size) return 2;
 318     s += 1 + b3key(v);
 319     if (s > size) return 3;
 320     s += 1 + b4key(v);
 321     if (s > size) return 4;
 322     return 5;
 323   }
 324
 325   /**
 326    * Given a pointer to the beginning of an GroupVarint64-encoded block,
 327    * return the number of bytes used by the encoding.
 328    */
 329   static size_t encodedSize(const char* p) {
 330     uint16_t n = loadUnaligned<uint16_t>(p);
 331     return kHeaderSize + kGroupSize +
 332             b0key(n) + b1key(n) + b2key(n) + b3key(n) + b4key(n);
 333   }
 334
 335   /**
 336    * Encode five uint64_t values into the buffer pointed-to by p, and return
 337    * the next position in the buffer (that is, one character past the last
 338    * encoded byte).  p needs to have at least size()+8 bytes available.
 339    */
 340   static char* encode(char* p, uint64_t a, uint64_t b, uint64_t c,
 341                       uint64_t d, uint64_t e) {
 342     uint16_t b0key = key(a);
 343     uint16_t b1key = key(b);
 344     uint16_t b2key = key(c);
 345     uint16_t b3key = key(d);
 346     uint16_t b4key = key(e);
 347     storeUnaligned<uint16_t>(
 348         p,
 349         uint16_t(
 350             (b4key << 12) |
 351             (b3key << 9) |
 352             (b2key << 6) |
 353             (b1key << 3) |
 354             b0key));
 355     p += 2;
 356     storeUnaligned(p, a);
 357     p += b0key+1;
 358     storeUnaligned(p, b);
 359     p += b1key+1;
 360     storeUnaligned(p, c);
 361     p += b2key+1;
 362     storeUnaligned(p, d);
 363     p += b3key+1;
 364     storeUnaligned(p, e);
 365     p += b4key+1;
 366     return p;
 367   }
 368
 369   /**
 370    * Encode five uint64_t values from the array pointed-to by src into the
 371    * buffer pointed-to by p, similar to encode(p,a,b,c,d,e) above.
 372    */
 373   static char* encode(char* p, const uint64_t* src) {
 374     return encode(p, src[0], src[1], src[2], src[3], src[4]);
 375   }
 376
 377   /**
 378    * Decode five uint64_t values from a buffer, and return the next position
 379    * in the buffer (that is, one character past the last encoded byte).
 380    * The buffer needs to have at least 7 bytes available (they may be read
 381    * but ignored).
 382    */
 383   static const char* decode(const char* p, uint64_t* a, uint64_t* b,
 384                             uint64_t* c, uint64_t* d, uint64_t* e) {
 385     uint16_t k = loadUnaligned<uint16_t>(p);
 386     p += 2;
 387     uint8_t k0 = b0key(k);
 388     *a = loadUnaligned<uint64_t>(p) & kMask[k0];
 389     p += k0+1;
 390     uint8_t k1 = b1key(k);
 391     *b = loadUnaligned<uint64_t>(p) & kMask[k1];
 392     p += k1+1;
 393     uint8_t k2 = b2key(k);
 394     *c = loadUnaligned<uint64_t>(p) & kMask[k2];
 395     p += k2+1;
 396     uint8_t k3 = b3key(k);
 397     *d = loadUnaligned<uint64_t>(p) & kMask[k3];
 398     p += k3+1;
 399     uint8_t k4 = b4key(k);
 400     *e = loadUnaligned<uint64_t>(p) & kMask[k4];
 401     p += k4+1;
 402     return p;
 403   }
 404
 405   /**
 406    * Decode five uint64_t values from a buffer and store them in the array
 407    * pointed-to by dest, similar to decode(p,a,b,c,d,e) above.
 408    */
 409   static const char* decode(const char* p, uint64_t* dest) {
 410     return decode(p, dest, dest+1, dest+2, dest+3, dest+4);
 411   }
 412
 413  private:
 414   enum { kHeaderBytes = 2 };
 415
 416   static uint8_t key(uint64_t x) {
 417     // __builtin_clzll is undefined for the x==0 case
 418     return uint8_t(7 - (__builtin_clzll(x | 1) / 8));
 419   }
 420
 421   static uint8_t b0key(uint16_t x) { return x & 7u; }
 422   static uint8_t b1key(uint16_t x) { return (x >> 3) & 7u; }
 423   static uint8_t b2key(uint16_t x) { return (x >> 6) & 7u; }
 424   static uint8_t b3key(uint16_t x) { return (x >> 9) & 7u; }
 425   static uint8_t b4key(uint16_t x) { return (x >> 12) & 7u; }
 426
 427   static const uint64_t kMask[];
 428 };
 429
 430 typedef GroupVarint<uint32_t> GroupVarint32;
 431 typedef GroupVarint<uint64_t> GroupVarint64;
 432
 433 /**
 434  * Simplify use of GroupVarint* for the case where data is available one
 435  * entry at a time (instead of one group at a time).  Handles buffering
 436  * and an incomplete last chunk.
 437  *
 438  * Output is a function object that accepts character ranges:
 439  * out(StringPiece) appends the given character range to the output.
 440  */
 441 template <class T, class Output>
 442 class GroupVarintEncoder {
 443  public:
 444   typedef GroupVarint<T> Base;
 445   typedef T type;
 446
 447   explicit GroupVarintEncoder(Output out)
 448     : out_(out),
 449       count_(0) {
 450   }
 451
 452   ~GroupVarintEncoder() {
 453     finish();
 454   }
 455
 456   /**
 457    * Add a value to the encoder.
 458    */
 459   void add(type val) {
 460     buf_[count_++] = val;
 461     if (count_ == Base::kGroupSize) {
 462       char* p = Base::encode(tmp_, buf_);
 463       out_(StringPiece(tmp_, p));
 464       count_ = 0;
 465     }
 466   }
 467
 468   /**
 469    * Finish encoding, flushing any buffered values if necessary.
 470    * After finish(), the encoder is immediately ready to encode more data
 471    * to the same output.
 472    */
 473   void finish() {
 474     if (count_) {
 475       // This is not strictly necessary, but it makes testing easy;
 476       // uninitialized bytes are guaranteed to be recorded as taking one byte
 477       // (not more).
 478       for (size_t i = count_; i < Base::kGroupSize; i++) {
 479         buf_[i] = 0;
 480       }
 481       Base::encode(tmp_, buf_);
 482       out_(StringPiece(tmp_, Base::partialSize(buf_, count_)));
 483       count_ = 0;
 484     }
 485   }
 486
 487   /**
 488    * Return the appender that was used.
 489    */
 490   Output& output() {
 491     return out_;
 492   }
 493   const Output& output() const {
 494     return out_;
 495   }
 496
 497   /**
 498    * Reset the encoder, disregarding any state (except what was already
 499    * flushed to the output, of course).
 500    */
 501   void clear() {
 502     count_ = 0;
 503   }
 504
 505  private:
 506   Output out_;
 507   char tmp_[Base::kMaxSize];
 508   type buf_[Base::kGroupSize];
 509   size_t count_;
 510 };
 511
 512 /**
 513  * Simplify use of GroupVarint* for the case where the last group in the
 514  * input may be incomplete (but the exact size of the input is known).
 515  * Allows for extracting values one at a time.
 516  */
 517 template <typename T>
 518 class GroupVarintDecoder {
 519  public:
 520   typedef GroupVarint<T> Base;
 521   typedef T type;
 522
 523   GroupVarintDecoder() = default;
 524
 525   explicit GroupVarintDecoder(StringPiece data,
 526                               size_t maxCount = (size_t)-1)
 527     : rrest_(data.end()),
 528       p_(data.data()),
 529       end_(data.end()),
 530       limit_(end_),
 531       pos_(0),
 532       count_(0),
 533       remaining_(maxCount) {
 534   }
 535
 536   void reset(StringPiece data, size_t maxCount = (size_t)-1) {
 537     rrest_ = data.end();
 538     p_ = data.data();
 539     end_ = data.end();
 540     limit_ = end_;
 541     pos_ = 0;
 542     count_ = 0;
 543     remaining_ = maxCount;
 544   }
 545
 546   /**
 547    * Read and return the next value.
 548    */
 549   bool next(type* val) {
 550     if (pos_ == count_) {
 551       // refill
 552       size_t rem = size_t(end_ - p_);
 553       if (rem == 0 || remaining_ == 0) {
 554         return false;
 555       }
 556       // next() attempts to read one full group at a time, and so we must have
 557       // at least enough bytes readable after its end to handle the case if the
 558       // last group is full.
 559       //
 560       // The best way to ensure this is to ensure that data has at least
 561       // Base::kMaxSize - 1 bytes readable *after* the end, otherwise we'll copy
 562       // into a temporary buffer.
 563       if (limit_ - p_ < Base::kMaxSize) {
 564         memcpy(tmp_, p_, rem);
 565         p_ = tmp_;
 566         end_ = p_ + rem;
 567         limit_ = tmp_ + sizeof(tmp_);
 568       }
 569       pos_ = 0;
 570       const char* n = Base::decode(p_, buf_);
 571       if (n <= end_) {
 572         // Full group could be decoded
 573         if (remaining_ >= Base::kGroupSize) {
 574           remaining_ -= Base::kGroupSize;
 575           count_ = Base::kGroupSize;
 576           p_ = n;
 577         } else {
 578           count_ = remaining_;
 579           remaining_ = 0;
 580           p_ += Base::partialSize(buf_, count_);
 581         }
 582       } else {
 583         // Can't decode a full group
 584         count_ = Base::partialCount(p_, size_t(end_ - p_));
 585         if (remaining_ >= count_) {
 586           remaining_ -= count_;
 587           p_ = end_;
 588         } else {
 589           count_ = remaining_;
 590           remaining_ = 0;
 591           p_ += Base::partialSize(buf_, count_);
 592         }
 593         if (count_ == 0) {
 594           return false;
 595         }
 596       }
 597     }
 598     *val = buf_[pos_++];
 599     return true;
 600   }
 601
 602   StringPiece rest() const {
 603     // This is only valid after next() returned false
 604     CHECK(pos_ == count_ && (p_ == end_ || remaining_ == 0));
 605     // p_ may point to the internal buffer (tmp_), but we want
 606     // to return subpiece of the original data
 607     size_t size = size_t(end_ - p_);
 608     return StringPiece(rrest_ - size, rrest_);
 609   }
 610
 611  private:
 612   const char* rrest_;
 613   const char* p_;
 614   const char* end_;
 615   const char* limit_;
 616   char tmp_[2 * Base::kMaxSize];
 617   type buf_[Base::kGroupSize];
 618   size_t pos_;
 619   size_t count_;
 620   size_t remaining_;
 621 };
 622
 623 typedef GroupVarintDecoder<uint32_t> GroupVarint32Decoder;
 624 typedef GroupVarintDecoder<uint64_t> GroupVarint64Decoder;
 625
 626 }  // namespace folly
 627
 628 #endif /* FOLLY_X64 || defined(__i386__) || FOLLY_PPC64 */