Eliminate need for reading upper bound at construction.

Summary:
# Note: This is a resubmission after reverting this change.
The failing `DCHECK` has been fixed with an explicit cast to `size_t` to accommodate long lists.

# Summary

This diff introduces the following changes to the EliasFanoReader implementation.

- `EliasFanoReader` no longer requires knowledge of the last value in the list at construction time. This removes the need to access the last byte of the encoded list.
- Operations in `UpperBitsReader` are now responsible for ensuring validity.
- Removes constraint that the last set bit in upperBits must be in the last byte of the encoded list.

In addition, the diff also extends the unit-tests for Elias Fano coding to add arbitrary extensions to the upperBound at construction to ensure that additional 0-blocks at the end of the list do not affect the behavior of the reader.

Reviewed By: ot, luciang

Differential Revision: D22211304

fbshipit-source-id: 0dbe904c9fd8cd9568a480355e5e6a4525922966

folly/experimental/EliasFanoCoding.h

@@ -24,9 +24,11 @@
 #pragma once
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdlib>
 #include <limits>
 #include <type_traits>
+#include <utility>
 
 #include <folly/Likely.h>
 #include <folly/Portability.h>
@@ -354,18 +356,24 @@ struct EliasFanoEncoderV2<
 
 namespace detail {
 
-template <class Encoder, class Instructions, class SizeType>
+template <
+    class Encoder,
+    class Instructions,
+    class SizeType,
+    bool kUnchecked = false>
 class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
                         SkipPointers<Encoder::skipQuantum> {
-  typedef typename Encoder::SkipValueType SkipValueType;
+  using SkipValueType = typename Encoder::SkipValueType;
 
  public:
-  typedef typename Encoder::ValueType ValueType;
+  using ValueType = typename Encoder::ValueType;
 
   explicit UpperBitsReader(const typename Encoder::CompressedList& list)
       : ForwardPointers<Encoder::forwardQuantum>(list.forwardPointers),
         SkipPointers<Encoder::skipQuantum>(list.skipPointers),
-        start_(list.upper) {
+        start_(list.upper),
+        size_(list.size),
+        upperBound_(8 * list.upperSizeBytes - size_) {
     reset();
   }
 
@@ -380,11 +388,23 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
   FOLLY_ALWAYS_INLINE SizeType position() const {
     return position_;
   }
+
   FOLLY_ALWAYS_INLINE ValueType value() const {
     return value_;
   }
 
-  FOLLY_ALWAYS_INLINE ValueType previous() {
+  FOLLY_ALWAYS_INLINE bool valid() const {
+    return position() < size();
+  }
+
+  FOLLY_ALWAYS_INLINE SizeType size() const {
+    return size_;
+  }
+
+  FOLLY_ALWAYS_INLINE bool previous() {
+    if (!kUnchecked && UNLIKELY(position() == 0)) {
+      return false;
+    }
     size_t inner;
     block_t block;
     getPreviousInfo(block, inner, outer_);
@@ -394,7 +414,12 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
     return setValue(inner);
   }
 
-  FOLLY_ALWAYS_INLINE ValueType next() {
+  // Returns `true` if the next position is within bounds.
+  FOLLY_ALWAYS_INLINE bool next() {
+    if (!kUnchecked && UNLIKELY(position() + 1 >= size())) {
+      return setDone();
+    }
+
     // Skip to the first non-zero block.
     while (block_ == 0) {
       outer_ += sizeof(block_t);
@@ -404,12 +429,14 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
     ++position_;
     size_t inner = Instructions::ctz(block_);
     block_ = Instructions::blsr(block_);
-
     return setValue(inner);
   }
 
-  FOLLY_ALWAYS_INLINE ValueType skip(SizeType n) {
+  FOLLY_ALWAYS_INLINE bool skip(SizeType n) {
     DCHECK_GT(n, 0);
+    if (!kUnchecked && UNLIKELY(position_ + n >= size())) {
+      return setDone();
+    }
 
     position_ += n; // n 1-bits will be read.
 
@@ -435,15 +462,18 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
     DCHECK_GT(n, 0);
     size_t inner = select64<Instructions>(block_, n - 1);
     block_ &= (block_t(-1) << inner) << 1;
-
     return setValue(inner);
   }
 
   // Skip to the first element that is >= v and located *after* the current
   // one (so even if current value equals v, position will be increased by 1).
-  FOLLY_ALWAYS_INLINE ValueType skipToNext(ValueType v) {
+  FOLLY_ALWAYS_INLINE bool skipToNext(ValueType v) {
     DCHECK_GE(v, value_);
 
+    if (!kUnchecked && UNLIKELY(v > upperBound_)) {
+      return setDone();
+    }
+
     // Use skip pointer.
     if (Encoder::skipQuantum > 0 &&
         UNLIKELY(v >= value_ + Encoder::skipQuantum)) {
@@ -451,6 +481,12 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
       const size_t dest = folly::loadUnaligned<SkipValueType>(
           this->skipPointers_ + (steps - 1) * sizeof(SkipValueType));
 
+      DCHECK_LE(dest, size());
+      // Check if we've reached the end.
+      if (!kUnchecked && UNLIKELY(dest == size())) {
+        return setDone();
+      }
+
       reposition(dest + Encoder::skipQuantum * steps);
       position_ = dest - 1;
 
@@ -470,6 +506,7 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
       skip -= cnt;
       position_ += kBitsPerBlock - cnt;
       outer_ += sizeof(block_t);
+      DCHECK_LT(outer_, (static_cast<size_t>(upperBound_) + size()) / 8);
       block_ = folly::loadUnaligned<block_t>(start_ + outer_);
     }
 
@@ -479,17 +516,24 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
       block_ &= block_t(-1) << inner;
     }
 
-    next();
-    return value_;
+    DCHECK_LT(position() + 1, size() + 1);
+    return next();
   }
 
   /**
-   * Prepare to skip to `value`. This is a constant-time operation that will
-   * prefetch memory required for a `skipTo(value)` call.
+   * Try to prepare to skip to value. This is a constant-time operation
+   * that will attempt to prefetch memory required for a subsequent
+   * skipTo(value) call if the value to skip to is within this list.
    *
-   * @return position of reader
+   * Returns:
+   *   {true, position of the reader} if the skip is valid,
+   *   {false, size()} otherwise.
    */
-  FOLLY_ALWAYS_INLINE SizeType prepareSkipTo(ValueType v) const {
+  FOLLY_ALWAYS_INLINE std::pair<bool, SizeType> prepareSkipTo(
+      ValueType v) const {
+    if (!kUnchecked && UNLIKELY(v > upperBound_)) {
+      return std::make_pair(false, size());
+    }
     auto position = position_;
 
     if (Encoder::skipQuantum > 0 && v >= value_ + Encoder::skipQuantum) {
@@ -498,6 +542,12 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
       const size_t dest = folly::loadUnaligned<SkipValueType>(
           this->skipPointers_ + (steps - 1) * sizeof(SkipValueType));
 
+      DCHECK_LE(dest, size());
+      // Check if we're going out of bounds.
+      if (!kUnchecked && UNLIKELY(dest == size())) {
+        return std::make_pair(false, size());
+      }
+
       position = dest - 1;
       outer = (dest + Encoder::skipQuantum * steps) / 8;
 
@@ -512,7 +562,7 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
       __builtin_prefetch(addr + kCacheLineSize);
     }
 
-    return position;
+    return std::make_pair(true, position);
   }
 
   FOLLY_ALWAYS_INLINE ValueType previousValue() const {
@@ -534,14 +584,15 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
     return (start_[bitPos / 8] & (1 << (bitPos % 8))) == 0;
   }
 
-  FOLLY_ALWAYS_INLINE void setDone(SizeType endPos) {
-    position_ = endPos;
+  FOLLY_ALWAYS_INLINE bool setDone() {
+    position_ = size_;
+    return false;
   }
 
  private:
-  FOLLY_ALWAYS_INLINE ValueType setValue(size_t inner) {
+  FOLLY_ALWAYS_INLINE bool setValue(size_t inner) {
     value_ = static_cast<ValueType>(8 * outer_ + inner - position_);
-    return value_;
+    return true;
   }
 
   FOLLY_ALWAYS_INLINE void reposition(SizeType dest) {
@@ -573,6 +624,8 @@ class UpperBitsReader : ForwardPointers<Encoder::forwardQuantum>,
   }
 
   const unsigned char* const start_;
+  const SizeType size_; // Size of the list.
+  const ValueType upperBound_; // Upper bound of values in this list.
   block_t block_;
   SizeType position_; // Index of current value (= #reads - 1).
   OuterType outer_; // Outer offset: number of consumed bytes in upper.
@@ -598,24 +651,8 @@ class EliasFanoReader {
   typedef typename Encoder::ValueType ValueType;
 
   explicit EliasFanoReader(const typename Encoder::CompressedList& list)
-      : upper_(list),
-        lower_(list.lower),
-        size_(list.size),
-        numLowerBits_(list.numLowerBits) {
+      : upper_(list), lower_(list.lower), numLowerBits_(list.numLowerBits) {
     DCHECK(Instructions::supported());
-    // To avoid extra branching during skipTo() while reading
-    // upper sequence we need to know the last element.
-    // If kUnchecked == true, we do not check that skipTo() is called
-    // within the bounds, so we can avoid initializing lastValue_.
-    if (kUnchecked || UNLIKELY(list.size == 0)) {
-      lastValue_ = 0;
-      return;
-    }
-    ValueType lastUpperValue = ValueType(8 * list.upperSizeBytes - size_);
-    auto it = list.upper + list.upperSizeBytes - 1;
-    DCHECK_NE(*it, 0);
-    lastUpperValue -= 8 - folly::findLastSet(*it);
-    lastValue_ = readLowerPart(size_ - 1) | (lastUpperValue << numLowerBits_);
   }
 
   void reset() {
@@ -624,25 +661,21 @@ class EliasFanoReader {
   }
 
   bool previous() {
-    if (!kUnchecked && UNLIKELY(position() == 0)) {
+    if (LIKELY(upper_.previous())) {
+      setCurrentValue();
+      return true;
+    }
     reset();
     return false;
   }
-    upper_.previous();
-    value_ =
-        readLowerPart(upper_.position()) | (upper_.value() << numLowerBits_);
-    return true;
-  }
 
   bool next() {
-    if (!kUnchecked && UNLIKELY(position() + 1 >= size_)) {
-      return setDone();
-    }
-    upper_.next();
-    value_ =
-        readLowerPart(upper_.position()) | (upper_.value() << numLowerBits_);
+    if (LIKELY(upper_.next())) {
+      setCurrentValue();
       return true;
     }
+    return setDone();
+  }
 
   /**
    * Advances by n elements. n = 0 is allowed and has no effect. Returns false
@@ -653,15 +686,12 @@ class EliasFanoReader {
       return valid();
     }
 
-    if (kUnchecked || LIKELY(position() + n < size_)) {
-      upper_.skip(n);
-      value_ =
-          readLowerPart(upper_.position()) | (upper_.value() << numLowerBits_);
-      return true;
-    }
-
+    if (!upper_.skip(n)) {
       return setDone();
     }
+    setCurrentValue();
+    return true;
+  }
 
   /**
    * Skips to the first element >= value whose position is greater or equal to
@@ -671,19 +701,26 @@ class EliasFanoReader {
   bool skipTo(ValueType value) {
     if (value != kInvalidValue) {
       DCHECK_GE(value + 1, value_ + 1);
+      if (UNLIKELY(value == value_)) {
+        return true;
       }
+    }
+
+    ValueType upperValue = value >> numLowerBits_;
 
-    if (!kUnchecked && UNLIKELY(value > lastValue_)) {
+    if (UNLIKELY(!upper_.skipToNext(upperValue))) {
+      // The value is beyond the last value in the list.
       return setDone();
-    } else if (UNLIKELY(value == value_)) {
-      return true;
     }
 
-    ValueType upperValue = value >> numLowerBits_;
-    upper_.skipToNext(upperValue);
-    iterateTo(value);
+    do {
+      setCurrentValue();
+      if (LIKELY(value_ >= value)) {
         return true;
       }
+    } while (LIKELY(upper_.next()));
+    return setDone();
+  }
 
   /**
    * Prepare to skip to `value` by prefetching appropriate memory in both the
@@ -693,31 +730,33 @@ class EliasFanoReader {
     // Also works when value_ == kInvalidValue.
     if (value != kInvalidValue) {
       DCHECK_GE(value + 1, value_ + 1);
-    }
-
-    if ((!kUnchecked && value > lastValue_) || (value == value_)) {
+      if (UNLIKELY(value == value_)) {
         return;
       }
+    }
 
     // Do minimal computation required to prefetch address used in
     // `readLowerPart()`.
     ValueType upperValue = (value >> numLowerBits_);
-    const auto upperPosition = upper_.prepareSkipTo(upperValue);
-    const auto addr = lower_ + (upperPosition * numLowerBits_ / 8);
+    const auto skipDetails = upper_.prepareSkipTo(upperValue);
+    // If we have a valid skip position, prefetch the lower bits for it.
+    if (LIKELY(skipDetails.first)) {
+      const auto addr = lower_ + (skipDetails.second * numLowerBits_ / 8);
       __builtin_prefetch(addr);
       __builtin_prefetch(addr + kCacheLineSize);
     }
+  }
 
   /**
    * Jumps to the element at position n. The reader can be in any state. Returns
    * false if n >= size().
    */
   bool jump(SizeType n) {
-    if (n + 1 < upper_.position() + 1) { // Also works if position() == -1.
+    if (n + 1 < position() + 1) { // Also works if position() == -1.
       reset();
       n += 1; // Initial position is -1.
     } else {
-      n -= upper_.position();
+      n -= position();
     }
     return skip(n);
   }
@@ -738,7 +777,7 @@ class EliasFanoReader {
       // We might be in the middle of a run, iterate backwards to the beginning.
       auto valueLower = Instructions::bzhi(value_, numLowerBits_);
       while (!upper_.isAtBeginningOfRun() &&
-             readLowerPart(upper_.position() - 1) == valueLower) {
+             readLowerPart(position() - 1) == valueLower) {
         upper_.previous();
       }
       return true;
@@ -753,20 +792,15 @@ class EliasFanoReader {
     return skipTo(value);
   }
 
-  ValueType lastValue() const {
-    CHECK(!kUnchecked);
-    return lastValue_;
-  }
-
   ValueType previousValue() const {
     DCHECK_GT(position(), 0);
     DCHECK_LT(position(), size());
-    return readLowerPart(upper_.position() - 1) |
+    return readLowerPart(position() - 1) |
         (upper_.previousValue() << numLowerBits_);
   }
 
   SizeType size() const {
-    return size_;
+    return upper_.size();
   }
 
   bool valid() const {
@@ -787,37 +821,27 @@ class EliasFanoReader {
 
   FOLLY_ALWAYS_INLINE bool setDone() {
     value_ = kInvalidValue;
-    upper_.setDone(size_);
-    return false;
+    return upper_.setDone();
+  }
+
+  FOLLY_ALWAYS_INLINE void setCurrentValue() {
+    value_ = readLowerPart(position()) | (upper_.value() << numLowerBits_);
   }
 
   FOLLY_ALWAYS_INLINE ValueType readLowerPart(SizeType i) const {
-    DCHECK_LT(i, size_);
+    DCHECK_LT(i, size());
     const size_t pos = i * numLowerBits_;
     const unsigned char* ptr = lower_ + (pos / 8);
     const uint64_t ptrv = folly::loadUnaligned<uint64_t>(ptr);
     // This removes the branch in the fallback implementation of
     // bzhi. The condition is verified at encoding time.
-    assume(numLowerBits_ < sizeof(ValueType) * 8);
+    folly::assume(numLowerBits_ < sizeof(ValueType) * 8);
     return Instructions::bzhi(ptrv >> (pos % 8), numLowerBits_);
   }
 
-  FOLLY_ALWAYS_INLINE void iterateTo(ValueType value) {
-    while (true) {
-      value_ =
-          readLowerPart(upper_.position()) | (upper_.value() << numLowerBits_);
-      if (LIKELY(value_ >= value)) {
-        break;
-      }
-      upper_.next();
-    }
-  }
-
-  detail::UpperBitsReader<Encoder, Instructions, SizeType> upper_;
+  detail::UpperBitsReader<Encoder, Instructions, SizeType, kUnchecked> upper_;
   const uint8_t* lower_;
-  SizeType size_;
   ValueType value_ = kInvalidValue;
-  ValueType lastValue_;
   uint8_t numLowerBits_;
 };
 

folly/experimental/test/CodingTestUtils.h

@@ -35,6 +35,11 @@
 namespace folly {
 namespace compression {
 
+template <typename ValueType, class List>
+folly::Optional<std::size_t> getUniverseUpperBound(const List& /* list */) {
+  return folly::none;
+}
+
 template <class URNG>
 std::vector<uint64_t> generateRandomList(
     size_t n,
@@ -268,6 +273,16 @@ void testSkipTo(const std::vector<uint64_t>& data, const List& list) {
     EXPECT_EQ(reader.position(), reader.size());
     EXPECT_FALSE(reader.next());
   }
+  // Skip past the last element and before the upperBound.
+  using ValueType = typename Reader::ValueType;
+  if (const auto upperBound = getUniverseUpperBound<ValueType>(list);
+      upperBound && *upperBound != data.back()) {
+    Reader reader(list);
+    EXPECT_FALSE(reader.skipTo(*upperBound));
+    EXPECT_FALSE(reader.valid());
+    EXPECT_EQ(reader.position(), reader.size());
+    EXPECT_FALSE(reader.next());
+  }
 }
 
 template <class Reader, class List>
@@ -362,9 +377,19 @@ void testEmpty() {
   }
 }
 
+// `upperBoundExtension` is required to inject additional 0-blocks
+// at the end of the list. This allows us to test lists with a large gap between
+// last element and universe upper bound, to exercise bounds-checking when
+// skipping past the last element
 template <class Reader, class Encoder>
-void testAll(const std::vector<uint64_t>& data) {
-  auto list = Encoder::encode(data.begin(), data.end());
+void testAll(
+    const std::vector<uint64_t>& data,
+    uint64_t upperBoundExtension = 0) {
+  Encoder encoder(data.size(), data.back() + upperBoundExtension);
+  for (const auto value : data) {
+    encoder.add(value);
+  }
+  auto list = encoder.finish();
   testNext<Reader>(data, list);
   testSkip<Reader>(data, list);
   testSkipTo<Reader>(data, list);

folly/experimental/test/EliasFanoCodingTest.cpp

@@ -21,11 +21,32 @@
 #include <vector>
 
 #include <folly/Benchmark.h>
+#include <folly/Optional.h>
+#include <folly/Random.h>
 #include <folly/experimental/EliasFanoCoding.h>
 #include <folly/experimental/Select64.h>
 #include <folly/experimental/test/CodingTestUtils.h>
 #include <folly/init/Init.h>
 
+namespace folly {
+namespace compression {
+
+// Overload to help CodingTestUtils retrieve the universe upperbound
+// of the list for certain test cases.
+template <typename ValueType, typename T>
+folly::Optional<std::size_t> getUniverseUpperBound(
+    const EliasFanoCompressedListBase<T>& list) {
+  constexpr ValueType maxUpperValue = std::numeric_limits<ValueType>::max();
+  const ValueType maxUpperBits = maxUpperValue >> list.numLowerBits;
+  const ValueType upperBitsUniverse = std::min(
+      static_cast<ValueType>(8 * list.upperSizeBytes - list.size),
+      maxUpperBits);
+  return (upperBitsUniverse << list.numLowerBits) |
+      ((1 << list.numLowerBits) - 1);
+}
+} // namespace compression
+} // namespace folly
+
 using namespace folly::compression;
 
 namespace {
@@ -114,6 +135,13 @@ class EliasFanoCodingTest : public ::testing::Test {
     // max() cannot be read, as it is assumed an invalid value.
     // TODO(ott): It should be possible to lift this constraint.
     testAll<Reader, Encoder>({0, 1, std::numeric_limits<uint32_t>::max() - 1});
+    // Test data with additional trailing 0s in the upperBits by extending
+    // the upper bound.
+    constexpr uint64_t minUpperBoundExtension = 2;
+    constexpr uint64_t maxUpperBoundExtension = 1024;
+    testAll<Reader, Encoder>(
+        generateRandomList(100 * 1000, 10 * 1000 * 1000),
+        folly::Random::rand32(minUpperBoundExtension, maxUpperBoundExtension));
   }
 };