folly/GroupVarint.h

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