folly/GroupVarint.h

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