fix heap-buffer-overflow (asan) in EliasFanoCoding
[folly.git] / folly / experimental / EliasFanoCoding.h
1 /*
2  * Copyright 2013 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 !defined(__x86_64__)
32 #error EliasFanoCoding.h requires x86_64
33 #endif
34
35 #include <cstdlib>
36 #include <algorithm>
37 #include <limits>
38 #include <type_traits>
39 #include <boost/noncopyable.hpp>
40 #include <glog/logging.h>
41 #include "folly/Bits.h"
42 #include "folly/CpuId.h"
43 #include "folly/Likely.h"
44 #include "folly/Range.h"
45
46 #if __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
47 #error EliasFanoCoding.h requires little endianness
48 #endif
49
50 namespace folly { namespace compression {
51
52 template <class Value,
53           class SkipValue = size_t,
54           size_t kSkipQuantum = 0,     // 0 = disabled
55           size_t kForwardQuantum = 0>  // 0 = disabled
56 struct EliasFanoCompressedList {
57   static_assert(std::is_integral<Value>::value &&
58                 std::is_unsigned<Value>::value,
59                 "Value should be unsigned integral");
60
61   typedef Value ValueType;
62   typedef SkipValue SkipValueType;
63
64   EliasFanoCompressedList()
65     : size(0), numLowerBits(0) { }
66
67   static constexpr size_t skipQuantum = kSkipQuantum;
68   static constexpr size_t forwardQuantum = kForwardQuantum;
69
70   size_t size;
71   uint8_t numLowerBits;
72
73   // WARNING: EliasFanoCompressedList has no ownership of
74   // lower, upper, skipPointers and forwardPointers.
75   // The 7 bytes following the last byte of lower and upper
76   // sequences should be readable.
77   folly::ByteRange lower;
78   folly::ByteRange upper;
79
80   folly::ByteRange skipPointers;
81   folly::ByteRange forwardPointers;
82
83   void free() {
84     ::free(const_cast<unsigned char*>(lower.data()));
85     ::free(const_cast<unsigned char*>(upper.data()));
86     ::free(const_cast<unsigned char*>(skipPointers.data()));
87     ::free(const_cast<unsigned char*>(forwardPointers.data()));
88   }
89
90   static uint8_t defaultNumLowerBits(size_t upperBound, size_t size) {
91     if (size == 0 || upperBound < size) {
92       return 0;
93     }
94     // floor(log(upperBound / size));
95     return folly::findLastSet(upperBound / size) - 1;
96   }
97
98   // WARNING: encode() mallocates lower, upper, skipPointers
99   // and forwardPointers. As EliasFanoCompressedList has
100   // no ownership of them, you need to call free() explicitly.
101   static void encode(const ValueType* list, size_t size,
102                      EliasFanoCompressedList& result) {
103     encode(list, list + size, result);
104   }
105
106   template <class RandomAccessIterator>
107   static void encode(RandomAccessIterator begin,
108                      RandomAccessIterator end,
109                      EliasFanoCompressedList& result) {
110     auto list = begin;
111     const size_t size = end - begin;
112
113     if (size == 0) {
114       result = EliasFanoCompressedList();
115       return;
116     }
117
118     DCHECK(std::is_sorted(list, list + size));
119
120     const ValueType upperBound = list[size - 1];
121     uint8_t numLowerBits = defaultNumLowerBits(upperBound, size);
122
123     // This is detail::writeBits56 limitation.
124     numLowerBits = std::min<uint8_t>(numLowerBits, 56);
125     CHECK_LT(numLowerBits, 8 * sizeof(Value));  // As we shift by numLowerBits.
126
127     // WARNING: Current read/write logic assumes that the 7 bytes
128     // following the last byte of lower and upper sequences are
129     // readable (stored value doesn't matter and won't be changed),
130     // so we allocate additional 7B, but do not include them in size
131     // of returned value.
132
133     // *** Lower bits.
134     const size_t lowerSize = (numLowerBits * size + 7) / 8;
135     unsigned char* lower = nullptr;
136     if (lowerSize > 0) {  // numLowerBits != 0
137       lower = static_cast<unsigned char*>(calloc(lowerSize + 7, 1));
138       const ValueType lowerMask = (ValueType(1) << numLowerBits) - 1;
139       for (size_t i = 0; i < size; ++i) {
140         const ValueType lowerBits = list[i] & lowerMask;
141         writeBits56(lower, i * numLowerBits, numLowerBits, lowerBits);
142       }
143     }
144
145     // *** Upper bits.
146     // Upper bits are stored using unary delta encoding.
147     // For example, (3 5 5 9) will be encoded as 1000011001000_2.
148     const size_t upperSizeBits =
149       (upperBound >> numLowerBits) +  // Number of 0-bits to be stored.
150       size;                           // 1-bits.
151     const size_t upperSize = (upperSizeBits + 7) / 8;
152     unsigned char* const upper =
153       static_cast<unsigned char*>(calloc(upperSize + 7, 1));
154     for (size_t i = 0; i < size; ++i) {
155       const ValueType upperBits = list[i] >> numLowerBits;
156       const size_t pos = upperBits + i;  // upperBits 0-bits and (i + 1) 1-bits.
157       upper[pos / 8] |= 1U << (pos % 8);
158     }
159
160     // *** Skip pointers.
161     // Store (1-indexed) position of every skipQuantum-th
162     // 0-bit in upper bits sequence.
163     SkipValueType* skipPointers = nullptr;
164     size_t numSkipPointers = 0;
165     /* static */ if (skipQuantum != 0) {
166       // Workaround to avoid 'division by zero' compile-time error.
167       constexpr size_t q = skipQuantum ?: 1;
168       CHECK_LT(upperSizeBits, std::numeric_limits<SkipValueType>::max());
169       // 8 * upperSize is used here instead of upperSizeBits, as that is
170       // more serialization-friendly way.
171       numSkipPointers = (8 * upperSize - size) / q;
172       skipPointers = static_cast<SkipValueType*>(
173           numSkipPointers == 0
174             ? nullptr
175             : calloc(numSkipPointers, sizeof(SkipValueType)));
176
177       for (size_t i = 0, pos = 0; i < size; ++i) {
178         const ValueType upperBits = list[i] >> numLowerBits;
179         for (; (pos + 1) * q <= upperBits; ++pos) {
180           skipPointers[pos] = i + (pos + 1) * q;
181         }
182       }
183     }
184
185     // *** Forward pointers.
186     // Store (1-indexed) position of every forwardQuantum-th
187     // 1-bit in upper bits sequence.
188     SkipValueType* forwardPointers = nullptr;
189     size_t numForwardPointers = 0;
190     /* static */ if (forwardQuantum != 0) {
191       // Workaround to avoid 'division by zero' compile-time error.
192       constexpr size_t q = forwardQuantum ?: 1;
193       CHECK_LT(upperSizeBits, std::numeric_limits<SkipValueType>::max());
194
195       numForwardPointers = size / q;
196       forwardPointers = static_cast<SkipValueType*>(
197         numForwardPointers == 0
198           ? nullptr
199           : malloc(numForwardPointers * sizeof(SkipValueType)));
200
201       for (size_t i = q - 1, pos = 0; i < size; i += q, ++pos) {
202         const ValueType upperBits = list[i] >> numLowerBits;
203         forwardPointers[pos] = upperBits + i + 1;
204       }
205     }
206
207     // *** Result.
208     result.size = size;
209     result.numLowerBits = numLowerBits;
210     result.lower.reset(lower, lowerSize);
211     result.upper.reset(upper, upperSize);
212     result.skipPointers.reset(
213         reinterpret_cast<unsigned char*>(skipPointers),
214         numSkipPointers * sizeof(SkipValueType));
215     result.forwardPointers.reset(
216         reinterpret_cast<unsigned char*>(forwardPointers),
217         numForwardPointers * sizeof(SkipValueType));
218   }
219
220  private:
221   // Writes value (with len up to 56 bits) to data starting at pos-th bit.
222   static void writeBits56(unsigned char* data, size_t pos,
223                           uint8_t len, uint64_t value) {
224     DCHECK_LE(uint32_t(len), 56);
225     DCHECK_EQ(0, value & ~((uint64_t(1) << len) - 1));
226     unsigned char* const ptr = data + (pos / 8);
227     uint64_t ptrv = folly::loadUnaligned<uint64_t>(ptr);
228     ptrv |=  value << (pos % 8);
229     folly::storeUnaligned<uint64_t>(ptr, ptrv);
230   }
231 };
232
233 // NOTE: It's recommended to compile EF coding with -msse4.2, starting
234 // with Nehalem, Intel CPUs support POPCNT instruction and gcc will emit
235 // it for __builtin_popcountll intrinsic.
236 // But we provide an alternative way for the client code: it can switch to
237 // the appropriate version of EliasFanoReader<> in realtime (client should
238 // implement this switching logic itself) by specifying instruction set to
239 // use explicitly.
240 namespace instructions {
241
242 struct Default {
243   static bool supported() {
244     return true;
245   }
246   static inline uint64_t popcount(uint64_t value) {
247     return __builtin_popcountll(value);
248   }
249   static inline int ctz(uint64_t value) {
250     DCHECK_GT(value, 0);
251     return __builtin_ctzll(value);
252   }
253 };
254
255 struct Fast : public Default {
256   static bool supported() {
257     folly::CpuId cpuId;
258     return cpuId.popcnt();
259   }
260   static inline uint64_t popcount(uint64_t value) {
261     uint64_t result;
262     asm ("popcntq %1, %0" : "=r" (result) : "r" (value));
263     return result;
264   }
265 };
266
267 }  // namespace instructions
268
269 namespace detail {
270
271 template <class CompressedList, class Instructions>
272 class UpperBitsReader {
273   typedef typename CompressedList::SkipValueType SkipValueType;
274  public:
275   typedef typename CompressedList::ValueType ValueType;
276
277   explicit UpperBitsReader(const CompressedList& list)
278     : forwardPointers_(list.forwardPointers.data()),
279       skipPointers_(list.skipPointers.data()),
280       start_(list.upper.data()),
281       block_(start_ != nullptr ? folly::loadUnaligned<block_t>(start_) : 0),
282       outer_(0),  // outer offset: number of consumed bytes in upper.
283       inner_(-1),  // inner offset: (bit) position in current block.
284       position_(-1),  // index of current value (= #reads - 1).
285       value_(0) { }
286
287   size_t position() const { return position_; }
288   ValueType value() const { return value_; }
289
290   ValueType next() {
291     // Skip to the first non-zero block.
292     while (block_ == 0) {
293       outer_ += sizeof(block_t);
294       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
295     }
296
297     ++position_;
298     inner_ = Instructions::ctz(block_);
299     block_ &= block_ - 1;
300
301     return setValue();
302   }
303
304   ValueType skip(size_t n) {
305     DCHECK_GT(n, 0);
306
307     position_ += n;  // n 1-bits will be read.
308
309     // Use forward pointer.
310     if (CompressedList::forwardQuantum > 0 &&
311         n > CompressedList::forwardQuantum) {
312       // Workaround to avoid 'division by zero' compile-time error.
313       constexpr size_t q = CompressedList::forwardQuantum ?: 1;
314
315       const size_t steps = position_ / q;
316       const size_t dest =
317         folly::loadUnaligned<SkipValueType>(
318             forwardPointers_ + (steps - 1) * sizeof(SkipValueType));
319
320       reposition(dest);
321       n = position_ + 1 - steps * q;  // n is > 0.
322       // correct inner_ will be set at the end.
323     }
324
325     size_t cnt;
326     // Find necessary block.
327     while ((cnt = Instructions::popcount(block_)) < n) {
328       n -= cnt;
329       outer_ += sizeof(block_t);
330       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
331     }
332
333     // NOTE: Trying to skip half-block here didn't show any
334     // performance improvements.
335
336     DCHECK_GT(n, 0);
337
338     // Kill n - 1 least significant 1-bits.
339     for (size_t i = 0; i < n - 1; ++i) {
340       block_ &= block_ - 1;
341     }
342
343     inner_ = Instructions::ctz(block_);
344     block_ &= block_ - 1;
345
346     return setValue();
347   }
348
349   // Skip to the first element that is >= v and located *after* the current
350   // one (so even if current value equals v, position will be increased by 1).
351   ValueType skipToNext(ValueType v) {
352     DCHECK_GE(v, value_);
353
354     // Use skip pointer.
355     if (CompressedList::skipQuantum > 0 &&
356         v >= value_ + CompressedList::skipQuantum) {
357       // Workaround to avoid 'division by zero' compile-time error.
358       constexpr size_t q = CompressedList::skipQuantum ?: 1;
359
360       const size_t steps = v / q;
361       const size_t dest =
362         folly::loadUnaligned<SkipValueType>(
363             skipPointers_ + (steps - 1) * sizeof(SkipValueType));
364
365       reposition(dest);
366       position_ = dest - q * steps - 1;
367       // Correct inner_ and value_ will be set during the next()
368       // call at the end.
369
370       // NOTE: Corresponding block of lower bits sequence may be
371       // prefetched here (via __builtin_prefetch), but experiments
372       // didn't show any significant improvements.
373     }
374
375     // Skip by blocks.
376     size_t cnt;
377     size_t skip = v - (8 * outer_ - position_ - 1);
378
379     constexpr size_t kBitsPerBlock = 8 * sizeof(block_t);
380     while ((cnt = Instructions::popcount(~block_)) < skip) {
381       skip -= cnt;
382       position_ += kBitsPerBlock - cnt;
383       outer_ += sizeof(block_t);
384       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
385     }
386
387     // Try to skip half-block.
388     constexpr size_t kBitsPerHalfBlock = 4 * sizeof(block_t);
389     constexpr block_t halfBlockMask = (block_t(1) << kBitsPerHalfBlock) - 1;
390     if ((cnt = Instructions::popcount(~block_ & halfBlockMask)) < skip) {
391       position_ += kBitsPerHalfBlock - cnt;
392       block_ &= ~halfBlockMask;
393     }
394
395     // Just skip until we see expected value.
396     while (next() < v) { }
397     return value_;
398   }
399
400  private:
401   ValueType setValue() {
402     value_ = static_cast<ValueType>(8 * outer_ + inner_ - position_);
403     return value_;
404   }
405
406   void reposition(size_t dest) {
407     outer_ = dest / 8;
408     block_ = folly::loadUnaligned<block_t>(start_ + outer_);
409     block_ &= ~((block_t(1) << (dest % 8)) - 1);
410   }
411
412   typedef unsigned long long block_t;
413   const unsigned char* const forwardPointers_;
414   const unsigned char* const skipPointers_;
415   const unsigned char* const start_;
416   block_t block_;
417   size_t outer_;
418   size_t inner_;
419   size_t position_;
420   ValueType value_;
421 };
422
423 }  // namespace detail
424
425 template <class CompressedList,
426           class Instructions = instructions::Default>
427 class EliasFanoReader : private boost::noncopyable {
428  public:
429   typedef typename CompressedList::ValueType ValueType;
430
431   explicit EliasFanoReader(const CompressedList& list)
432     : list_(list),
433       lowerMask_((ValueType(1) << list_.numLowerBits) - 1),
434       upper_(list),
435       progress_(0),
436       value_(0) {
437     DCHECK(Instructions::supported());
438     // To avoid extra branching during skipTo() while reading
439     // upper sequence we need to know the last element.
440     if (UNLIKELY(list_.size == 0)) {
441       lastValue_ = 0;
442       return;
443     }
444     ValueType lastUpperValue = 8 * list_.upper.size() - list_.size;
445     auto it = list_.upper.end() - 1;
446     DCHECK_NE(*it, 0);
447     lastUpperValue -= 8 - folly::findLastSet(*it);
448     lastValue_ = readLowerPart(list_.size - 1) |
449                  (lastUpperValue << list_.numLowerBits);
450   }
451
452   size_t size() const { return list_.size; }
453
454   size_t position() const { return progress_ - 1; }
455   ValueType value() const { return value_; }
456
457   bool next() {
458     if (UNLIKELY(progress_ == list_.size)) {
459       value_ = std::numeric_limits<ValueType>::max();
460       return false;
461     }
462     value_ = readLowerPart(progress_) |
463              (upper_.next() << list_.numLowerBits);
464     ++progress_;
465     return true;
466   }
467
468   bool skip(size_t n) {
469     CHECK_GT(n, 0);
470
471     progress_ += n - 1;
472     if (LIKELY(progress_ < list_.size)) {
473       value_ = readLowerPart(progress_) |
474                (upper_.skip(n) << list_.numLowerBits);
475       ++progress_;
476       return true;
477     }
478
479     progress_ = list_.size;
480     value_ = std::numeric_limits<ValueType>::max();
481     return false;
482   }
483
484   bool skipTo(ValueType value) {
485     DCHECK_GE(value, value_);
486     if (value <= value_) {
487       return true;
488     }
489     if (value > lastValue_) {
490       progress_ = list_.size;
491       value_ = std::numeric_limits<ValueType>::max();
492       return false;
493     }
494
495     upper_.skipToNext(value >> list_.numLowerBits);
496     progress_ = upper_.position();
497     value_ = readLowerPart(progress_) |
498              (upper_.value() << list_.numLowerBits);
499     ++progress_;
500     while (value_ < value) {
501       value_ = readLowerPart(progress_) |
502                (upper_.next() << list_.numLowerBits);
503       ++progress_;
504     }
505
506     return true;
507   }
508
509  private:
510   ValueType readLowerPart(size_t i) const {
511     const size_t pos = i * list_.numLowerBits;
512     const unsigned char* ptr = list_.lower.data() + (pos / 8);
513     const uint64_t ptrv = folly::loadUnaligned<uint64_t>(ptr);
514     return lowerMask_ & (ptrv >> (pos % 8));
515   }
516
517   const CompressedList list_;
518   const ValueType lowerMask_;
519   detail::UpperBitsReader<CompressedList, Instructions> upper_;
520   size_t progress_;
521   ValueType value_;
522   ValueType lastValue_;
523 };
524
525 }}  // namespaces
526
527 #endif  // FOLLY_EXPERIMENTAL_ELIAS_FANO_CODING_H