Elias-Fano micro-optimizations
[folly.git] / folly / experimental / EliasFanoCoding.h
1 /*
2  * Copyright 2014 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 /**
18  * @author Philip Pronin (philipp@fb.com)
19  *
20  * Based on the paper by Sebastiano Vigna,
21  * "Quasi-succinct indices" (arxiv:1206.4300).
22  */
23
24 #ifndef FOLLY_EXPERIMENTAL_ELIAS_FANO_CODING_H
25 #define FOLLY_EXPERIMENTAL_ELIAS_FANO_CODING_H
26
27 #ifndef __GNUC__
28 #error EliasFanoCoding.h requires GCC
29 #endif
30
31 #if !FOLLY_X64
32 #error EliasFanoCoding.h requires x86_64
33 #endif
34
35 #include <cstdlib>
36 #include <limits>
37 #include <type_traits>
38 #include <boost/noncopyable.hpp>
39 #include <glog/logging.h>
40
41 #include <folly/Bits.h>
42 #include <folly/CpuId.h>
43 #include <folly/Likely.h>
44 #include <folly/Range.h>
45 #include <folly/experimental/Select64.h>
46
47 #if __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
48 #error EliasFanoCoding.h requires little endianness
49 #endif
50
51 namespace folly { namespace compression {
52
53 struct EliasFanoCompressedList {
54   EliasFanoCompressedList() { }
55
56   void free() {
57     ::free(const_cast<unsigned char*>(lower.data()));
58     ::free(const_cast<unsigned char*>(upper.data()));
59     ::free(const_cast<unsigned char*>(skipPointers.data()));
60     ::free(const_cast<unsigned char*>(forwardPointers.data()));
61   }
62
63   size_t size = 0;
64   uint8_t numLowerBits = 0;
65
66   // WARNING: EliasFanoCompressedList has no ownership of
67   // lower, upper, skipPointers and forwardPointers.
68   // The 7 bytes following the last byte of lower and upper
69   // sequences should be readable.
70   folly::ByteRange lower;
71   folly::ByteRange upper;
72
73   folly::ByteRange skipPointers;
74   folly::ByteRange forwardPointers;
75 };
76
77 // Version history:
78 // In version 1 skip / forward pointers encoding has been changed,
79 // so SkipValue = uint32_t is able to address up to ~4B elements,
80 // instead of only ~2B.
81 template <class Value,
82           class SkipValue = size_t,
83           size_t kSkipQuantum = 0,     // 0 = disabled
84           size_t kForwardQuantum = 0,  // 0 = disabled
85           size_t kVersion = 0>
86 struct EliasFanoEncoder {
87   static_assert(std::is_integral<Value>::value &&
88                 std::is_unsigned<Value>::value,
89                 "Value should be unsigned integral");
90
91   typedef EliasFanoCompressedList CompressedList;
92
93   typedef Value ValueType;
94   typedef SkipValue SkipValueType;
95
96   static constexpr size_t skipQuantum = kSkipQuantum;
97   static constexpr size_t forwardQuantum = kForwardQuantum;
98   static constexpr size_t version = kVersion;
99
100   static uint8_t defaultNumLowerBits(size_t upperBound, size_t size) {
101     if (size == 0 || upperBound < size) {
102       return 0;
103     }
104     // floor(log(upperBound / size));
105     return folly::findLastSet(upperBound / size) - 1;
106   }
107
108   // Requires: input range (begin, end) is sorted (encoding
109   // crashes if it's not).
110   // WARNING: encode() mallocates lower, upper, skipPointers
111   // and forwardPointers. As EliasFanoCompressedList has
112   // no ownership of them, you need to call free() explicitly.
113   template <class RandomAccessIterator>
114   static EliasFanoCompressedList encode(RandomAccessIterator begin,
115                                         RandomAccessIterator end) {
116     if (begin == end) {
117       return EliasFanoCompressedList();
118     }
119     EliasFanoEncoder encoder(end - begin, *(end - 1));
120     for (; begin != end; ++begin) {
121       encoder.add(*begin);
122     }
123     return encoder.finish();
124   }
125
126   EliasFanoEncoder(size_t size, ValueType upperBound) {
127     if (size == 0) {
128       return;
129     }
130
131     uint8_t numLowerBits = defaultNumLowerBits(upperBound, size);
132
133     // This is detail::writeBits56 limitation.
134     numLowerBits = std::min<uint8_t>(numLowerBits, 56);
135     CHECK_LT(numLowerBits, 8 * sizeof(Value));  // As we shift by numLowerBits.
136
137     // WARNING: Current read/write logic assumes that the 7 bytes
138     // following the last byte of lower and upper sequences are
139     // readable (stored value doesn't matter and won't be changed),
140     // so we allocate additional 7B, but do not include them in size
141     // of returned value.
142
143     // *** Lower bits.
144     const size_t lowerSize = (numLowerBits * size + 7) / 8;
145     if (lowerSize > 0) {  // numLowerBits != 0
146       lower_ = static_cast<unsigned char*>(calloc(lowerSize + 7, 1));
147     }
148
149     // *** Upper bits.
150     // Upper bits are stored using unary delta encoding.
151     // For example, (3 5 5 9) will be encoded as 1000011001000_2.
152     const size_t upperSizeBits =
153       (upperBound >> numLowerBits) +  // Number of 0-bits to be stored.
154       size;                           // 1-bits.
155     const size_t upperSize = (upperSizeBits + 7) / 8;
156     upper_ = static_cast<unsigned char*>(calloc(upperSize + 7, 1));
157
158     // *** Skip pointers.
159     // Store (1-indexed) position of every skipQuantum-th
160     // 0-bit in upper bits sequence.
161     size_t numSkipPointers = 0;
162     /* static */ if (skipQuantum != 0) {
163       /* static */ if (kVersion > 0) {
164         CHECK_LT(size, std::numeric_limits<SkipValueType>::max());
165       } else {
166         CHECK_LT(upperSizeBits, std::numeric_limits<SkipValueType>::max());
167       }
168       // 8 * upperSize is used here instead of upperSizeBits, as that is
169       // more serialization-friendly way (upperSizeBits isn't known outside of
170       // this function, unlike upperSize; thus numSkipPointers could easily be
171       // deduced from upperSize).
172       numSkipPointers = (8 * upperSize - size) / (skipQuantum ?: 1);
173       skipPointers_ = static_cast<SkipValueType*>(
174           numSkipPointers == 0
175             ? nullptr
176             : calloc(numSkipPointers, sizeof(SkipValueType)));
177     }
178
179     // *** Forward pointers.
180     // Store (1-indexed) position of every forwardQuantum-th
181     // 1-bit in upper bits sequence.
182     size_t numForwardPointers = 0;
183     /* static */ if (forwardQuantum != 0) {
184       /* static */ if (kVersion > 0) {
185         CHECK_LT(upperBound >> numLowerBits,
186                  std::numeric_limits<SkipValueType>::max());
187       } else {
188         CHECK_LT(upperSizeBits, std::numeric_limits<SkipValueType>::max());
189       }
190
191       // '?: 1' is a workaround for false 'division by zero' compile-time error.
192       numForwardPointers = size / (forwardQuantum ?: 1);
193       forwardPointers_ = static_cast<SkipValueType*>(
194         numForwardPointers == 0
195           ? nullptr
196           : malloc(numForwardPointers * sizeof(SkipValueType)));
197     }
198
199     // *** Result.
200     result_.size = size;
201     result_.numLowerBits = numLowerBits;
202     result_.lower.reset(lower_, lowerSize);
203     result_.upper.reset(upper_, upperSize);
204     result_.skipPointers.reset(
205         reinterpret_cast<unsigned char*>(skipPointers_),
206         numSkipPointers * sizeof(SkipValueType));
207     result_.forwardPointers.reset(
208         reinterpret_cast<unsigned char*>(forwardPointers_),
209         numForwardPointers * sizeof(SkipValueType));
210   }
211
212   void add(ValueType value) {
213     CHECK_GE(value, lastValue_);
214
215     const auto numLowerBits = result_.numLowerBits;
216     const ValueType upperBits = value >> numLowerBits;
217
218     // Upper sequence consists of upperBits 0-bits and (size_ + 1) 1-bits.
219     const size_t pos = upperBits + size_;
220     upper_[pos / 8] |= 1U << (pos % 8);
221     // Append numLowerBits bits to lower sequence.
222     if (numLowerBits != 0) {
223       const ValueType lowerBits = value & ((ValueType(1) << numLowerBits) - 1);
224       writeBits56(lower_, size_ * numLowerBits, numLowerBits, lowerBits);
225     }
226
227     /* static */ if (skipQuantum != 0) {
228       while ((skipPointersSize_ + 1) * skipQuantum <= upperBits) {
229         /* static */ if (kVersion > 0) {
230           // Since version 1, just the number of preceding 1-bits is stored.
231           skipPointers_[skipPointersSize_] = size_;
232         } else {
233           skipPointers_[skipPointersSize_] =
234             size_ + (skipPointersSize_ + 1) * skipQuantum;
235         }
236         ++skipPointersSize_;
237       }
238     }
239
240     /* static */ if (forwardQuantum != 0) {
241       if ((size_ + 1) % forwardQuantum == 0) {
242         const auto pos = size_ / forwardQuantum;
243         /* static */ if (kVersion > 0) {
244           // Since version 1, just the number of preceding 0-bits is stored.
245           forwardPointers_[pos] = upperBits;
246         } else {
247           forwardPointers_[pos] = upperBits + size_ + 1;
248         }
249       }
250     }
251
252     lastValue_ = value;
253     ++size_;
254   }
255
256   const EliasFanoCompressedList& finish() const {
257     CHECK_EQ(size_, result_.size);
258     return result_;
259   }
260
261  private:
262   // Writes value (with len up to 56 bits) to data starting at pos-th bit.
263   static void writeBits56(unsigned char* data, size_t pos,
264                           uint8_t len, uint64_t value) {
265     DCHECK_LE(uint32_t(len), 56);
266     DCHECK_EQ(0, value & ~((uint64_t(1) << len) - 1));
267     unsigned char* const ptr = data + (pos / 8);
268     uint64_t ptrv = folly::loadUnaligned<uint64_t>(ptr);
269     ptrv |= value << (pos % 8);
270     folly::storeUnaligned<uint64_t>(ptr, ptrv);
271   }
272
273   unsigned char* lower_ = nullptr;
274   unsigned char* upper_ = nullptr;
275   SkipValueType* skipPointers_ = nullptr;
276   SkipValueType* forwardPointers_ = nullptr;
277
278   ValueType lastValue_ = 0;
279   size_t size_ = 0;
280   size_t skipPointersSize_ = 0;
281
282   EliasFanoCompressedList result_;
283 };
284
285 // NOTE: It's recommended to compile EF coding with -msse4.2, starting
286 // with Nehalem, Intel CPUs support POPCNT instruction and gcc will emit
287 // it for __builtin_popcountll intrinsic.
288 // But we provide an alternative way for the client code: it can switch to
289 // the appropriate version of EliasFanoReader<> in realtime (client should
290 // implement this switching logic itself) by specifying instruction set to
291 // use explicitly.
292 namespace instructions {
293
294 struct Default {
295   static bool supported(const folly::CpuId& cpuId = {}) {
296     return true;
297   }
298   static inline uint64_t popcount(uint64_t value) {
299     return __builtin_popcountll(value);
300   }
301   static inline int ctz(uint64_t value) {
302     DCHECK_GT(value, 0);
303     return __builtin_ctzll(value);
304   }
305   static inline uint64_t blsr(uint64_t value) {
306     return value & (value - 1);
307   }
308 };
309
310 struct Nehalem : public Default {
311   static bool supported(const folly::CpuId& cpuId = {}) {
312     return cpuId.popcnt();
313   }
314   static inline uint64_t popcount(uint64_t value) {
315     // POPCNT is supported starting with Intel Nehalem, AMD K10.
316     uint64_t result;
317     asm ("popcntq %1, %0" : "=r" (result) : "r" (value));
318     return result;
319   }
320 };
321
322 struct Haswell : public Nehalem {
323   static bool supported(const folly::CpuId& cpuId = {}) {
324     return Nehalem::supported(cpuId) && cpuId.bmi1();
325   }
326   static inline uint64_t blsr(uint64_t value) {
327     // BMI1 is supported starting with Intel Haswell, AMD Piledriver.
328     // BLSR combines two instuctions into one and reduces register pressure.
329     uint64_t result;
330     asm ("blsrq %1, %0" : "=r" (result) : "r" (value));
331     return result;
332   }
333 };
334
335 }  // namespace instructions
336
337 namespace detail {
338
339 template <class Encoder, class Instructions>
340 class UpperBitsReader {
341   typedef typename Encoder::SkipValueType SkipValueType;
342  public:
343   typedef typename Encoder::ValueType ValueType;
344
345   explicit UpperBitsReader(const EliasFanoCompressedList& list)
346     : forwardPointers_(list.forwardPointers.data()),
347       skipPointers_(list.skipPointers.data()),
348       start_(list.upper.data()) {
349     reset();
350   }
351
352   void reset() {
353     block_ = start_ != nullptr ? folly::loadUnaligned<block_t>(start_) : 0;
354     outer_ = 0;
355     inner_ = -1;
356     position_ = -1;
357     value_ = 0;
358   }
359
360   size_t position() const { return position_; }
361   ValueType value() const { return value_; }
362
363   ValueType next() {
364     // Skip to the first non-zero block.
365     while (block_ == 0) {
366       outer_ += sizeof(block_t);
367       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
368     }
369
370     ++position_;
371     inner_ = Instructions::ctz(block_);
372     block_ = Instructions::blsr(block_);
373
374     return setValue();
375   }
376
377   ValueType skip(size_t n) {
378     DCHECK_GT(n, 0);
379
380     position_ += n;  // n 1-bits will be read.
381
382     // Use forward pointer.
383     if (Encoder::forwardQuantum > 0 && n > Encoder::forwardQuantum) {
384       // Workaround to avoid 'division by zero' compile-time error.
385       constexpr size_t q = Encoder::forwardQuantum ?: 1;
386
387       const size_t steps = position_ / q;
388       const size_t dest =
389         folly::loadUnaligned<SkipValueType>(
390             forwardPointers_ + (steps - 1) * sizeof(SkipValueType));
391
392       /* static */ if (Encoder::version > 0) {
393         reposition(dest + steps * q);
394       } else {
395         reposition(dest);
396       }
397       n = position_ + 1 - steps * q;  // n is > 0.
398       // correct inner_ will be set at the end.
399     }
400
401     size_t cnt;
402     // Find necessary block.
403     while ((cnt = Instructions::popcount(block_)) < n) {
404       n -= cnt;
405       outer_ += sizeof(block_t);
406       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
407     }
408
409     // Skip to the n-th one in the block.
410     DCHECK_GT(n, 0);
411     inner_ = select64<Instructions>(block_, n - 1);
412     block_ &= (block_t(-1) << inner_) << 1;
413
414     return setValue();
415   }
416
417   // Skip to the first element that is >= v and located *after* the current
418   // one (so even if current value equals v, position will be increased by 1).
419   ValueType skipToNext(ValueType v) {
420     DCHECK_GE(v, value_);
421
422     // Use skip pointer.
423     if (Encoder::skipQuantum > 0 && v >= value_ + Encoder::skipQuantum) {
424       // Workaround to avoid 'division by zero' compile-time error.
425       constexpr size_t q = Encoder::skipQuantum ?: 1;
426
427       const size_t steps = v / q;
428       const size_t dest =
429         folly::loadUnaligned<SkipValueType>(
430             skipPointers_ + (steps - 1) * sizeof(SkipValueType));
431
432       /* static */ if (Encoder::version > 0) {
433         reposition(dest + q * steps);
434         position_ = dest - 1;
435       } else {
436         reposition(dest);
437         position_ = dest - q * steps - 1;
438       }
439       // Correct inner_ and value_ will be set during the next()
440       // call at the end.
441
442       // NOTE: Corresponding block of lower bits sequence may be
443       // prefetched here (via __builtin_prefetch), but experiments
444       // didn't show any significant improvements.
445     }
446
447     // Skip by blocks.
448     size_t cnt;
449     size_t skip = v - (8 * outer_ - position_ - 1);
450
451     constexpr size_t kBitsPerBlock = 8 * sizeof(block_t);
452     while ((cnt = Instructions::popcount(~block_)) < skip) {
453       skip -= cnt;
454       position_ += kBitsPerBlock - cnt;
455       outer_ += sizeof(block_t);
456       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
457     }
458
459     if (LIKELY(skip)) {
460       auto inner = select64<Instructions>(~block_, skip - 1);
461       position_ += inner - skip + 1;
462       block_ &= block_t(-1) << inner;
463     }
464
465     next();
466     return value_;
467   }
468
469   ValueType jump(size_t n) {
470     if (Encoder::forwardQuantum == 0 || n <= Encoder::forwardQuantum) {
471       reset();
472     } else {
473       position_ = -1;  // Avoid reading the head, skip() will reposition.
474     }
475     return skip(n);
476   }
477
478   ValueType jumpToNext(ValueType v) {
479     if (Encoder::skipQuantum == 0 || v < Encoder::skipQuantum) {
480       reset();
481     } else {
482       value_ = 0;  // Avoid reading the head, skipToNext() will reposition.
483     }
484     return skipToNext(v);
485   }
486
487  private:
488   ValueType setValue() {
489     value_ = static_cast<ValueType>(8 * outer_ + inner_ - position_);
490     return value_;
491   }
492
493   void reposition(size_t dest) {
494     outer_ = dest / 8;
495     block_ = folly::loadUnaligned<block_t>(start_ + outer_);
496     block_ &= ~((block_t(1) << (dest % 8)) - 1);
497   }
498
499   typedef unsigned long long block_t;
500   const unsigned char* const forwardPointers_;
501   const unsigned char* const skipPointers_;
502   const unsigned char* const start_;
503   block_t block_;
504   size_t outer_;  // Outer offset: number of consumed bytes in upper.
505   size_t inner_;  // Inner offset: (bit) position in current block.
506   size_t position_;  // Index of current value (= #reads - 1).
507   ValueType value_;
508 };
509
510 }  // namespace detail
511
512 template <class Encoder,
513           class Instructions = instructions::Default>
514 class EliasFanoReader : private boost::noncopyable {
515  public:
516   typedef Encoder EncoderType;
517   typedef typename Encoder::ValueType ValueType;
518
519   explicit EliasFanoReader(const EliasFanoCompressedList& list)
520     : list_(list),
521       lowerMask_((ValueType(1) << list_.numLowerBits) - 1),
522       upper_(list_) {
523     DCHECK(Instructions::supported());
524     // To avoid extra branching during skipTo() while reading
525     // upper sequence we need to know the last element.
526     if (UNLIKELY(list_.size == 0)) {
527       lastValue_ = 0;
528       return;
529     }
530     ValueType lastUpperValue = 8 * list_.upper.size() - list_.size;
531     auto it = list_.upper.end() - 1;
532     DCHECK_NE(*it, 0);
533     lastUpperValue -= 8 - folly::findLastSet(*it);
534     lastValue_ = readLowerPart(list_.size - 1) |
535                  (lastUpperValue << list_.numLowerBits);
536   }
537
538   void reset() {
539     upper_.reset();
540     progress_ = 0;
541     value_ = 0;
542   }
543
544   bool next() {
545     if (UNLIKELY(progress_ == list_.size)) {
546       value_ = std::numeric_limits<ValueType>::max();
547       return false;
548     }
549     value_ = readLowerPart(progress_) |
550              (upper_.next() << list_.numLowerBits);
551     ++progress_;
552     return true;
553   }
554
555   bool skip(size_t n) {
556     CHECK_GT(n, 0);
557
558     progress_ += n;
559     if (LIKELY(progress_ <= list_.size)) {
560       if (LIKELY(n < kLinearScanThreshold)) {
561         for (size_t i = 0; i < n; ++i) upper_.next();
562       } else {
563         upper_.skip(n);
564       }
565       value_ = readLowerPart(progress_ - 1) |
566         (upper_.value() << list_.numLowerBits);
567       return true;
568     }
569
570     progress_ = list_.size;
571     value_ = std::numeric_limits<ValueType>::max();
572     return false;
573   }
574
575   bool skipTo(ValueType value) {
576     DCHECK_GE(value, value_);
577     if (value <= value_) {
578       return true;
579     } else if (value > lastValue_) {
580       progress_ = list_.size;
581       value_ = std::numeric_limits<ValueType>::max();
582       return false;
583     }
584
585     size_t upperValue = (value >> list_.numLowerBits);
586     size_t upperSkip = upperValue - upper_.value();
587     // The average density of ones in upper bits is 1/2.
588     // LIKELY here seems to make things worse, even for small skips.
589     if (upperSkip < 2 * kLinearScanThreshold) {
590       do {
591         upper_.next();
592       } while (UNLIKELY(upper_.value() < upperValue));
593     } else {
594       upper_.skipToNext(upperValue);
595     }
596
597     iterateTo(value);
598     return true;
599   }
600
601   bool jump(size_t n) {
602     if (LIKELY(n - 1 < list_.size)) {  // n > 0 && n <= list_.size
603       progress_ = n;
604       value_ = readLowerPart(n - 1) | (upper_.jump(n) << list_.numLowerBits);
605       return true;
606     } else if (n == 0) {
607       reset();
608       return true;
609     }
610     progress_ = list_.size;
611     value_ = std::numeric_limits<ValueType>::max();
612     return false;
613   }
614
615   ValueType jumpTo(ValueType value) {
616     if (value <= 0) {
617       reset();
618       return true;
619     } else if (value > lastValue_) {
620       progress_ = list_.size;
621       value_ = std::numeric_limits<ValueType>::max();
622       return false;
623     }
624
625     upper_.jumpToNext(value >> list_.numLowerBits);
626     iterateTo(value);
627     return true;
628   }
629
630   size_t size() const { return list_.size; }
631
632   size_t position() const { return progress_ - 1; }
633   ValueType value() const { return value_; }
634
635  private:
636   ValueType readLowerPart(size_t i) const {
637     DCHECK_LT(i, list_.size);
638     const size_t pos = i * list_.numLowerBits;
639     const unsigned char* ptr = list_.lower.data() + (pos / 8);
640     const uint64_t ptrv = folly::loadUnaligned<uint64_t>(ptr);
641     return lowerMask_ & (ptrv >> (pos % 8));
642   }
643
644   void iterateTo(ValueType value) {
645     while (true) {
646       value_ = readLowerPart(upper_.position()) |
647         (upper_.value() << list_.numLowerBits);
648       if (LIKELY(value_ >= value)) break;
649       upper_.next();
650     }
651     progress_ = upper_.position() + 1;
652   }
653
654   constexpr static size_t kLinearScanThreshold = 8;
655
656   const EliasFanoCompressedList list_;
657   const ValueType lowerMask_;
658   detail::UpperBitsReader<Encoder, Instructions> upper_;
659   size_t progress_ = 0;
660   ValueType value_ = 0;
661   ValueType lastValue_;
662 };
663
664 }}  // namespaces
665
666 #endif  // FOLLY_EXPERIMENTAL_ELIAS_FANO_CODING_H