tor-browser

dec_bit_reader.h (11484B)

---

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#ifndef LIB_JXL_DEC_BIT_READER_H_
#define LIB_JXL_DEC_BIT_READER_H_

// Bounds-checked bit reader; 64-bit buffer with support for deferred refills
// and switching to reading byte-aligned words.

#include <cstddef>
#include <cstdint>
#include <cstring>  // memcpy

#ifdef __BMI2__
#include <immintrin.h>
#endif

#include "lib/jxl/base/byte_order.h"
#include "lib/jxl/base/common.h"
#include "lib/jxl/base/compiler_specific.h"
#include "lib/jxl/base/status.h"

namespace jxl {

// Reads bits previously written to memory by BitWriter. Uses unaligned 8-byte
// little-endian loads.
class BitReader {
public:
 static constexpr size_t kMaxBitsPerCall = 56;

 // Constructs an invalid BitReader, to be overwritten before usage.
 BitReader()
     : buf_(0),
       bits_in_buf_(0),
       next_byte_{nullptr},
       end_minus_8_{nullptr},
       first_byte_(nullptr) {}
 BitReader(const BitReader&) = delete;

 // bytes need not be aligned nor padded!
 template <class ArrayLike>
 explicit BitReader(const ArrayLike& bytes)
     : buf_(0),
       bits_in_buf_(0),
       next_byte_(bytes.data()),
       // Assumes first_byte_ >= 8.
       end_minus_8_(bytes.data() - 8 + bytes.size()),
       first_byte_(bytes.data()) {
   Refill();
 }
 ~BitReader() {
   // Close() must be called before destroying an initialized bit reader.
   // Invalid bit readers will have a nullptr in first_byte_.
   JXL_DASSERT(close_called_ || !first_byte_);
 }

 // Move operator needs to invalidate the other BitReader such that it is
 // irrelevant if we call Close() on it or not.
 BitReader& operator=(BitReader&& other) noexcept {
   // Ensure the current instance was already closed, before we overwrite it
   // with other.
   JXL_DASSERT(close_called_ || !first_byte_);

   JXL_DASSERT(!other.close_called_);
   buf_ = other.buf_;
   bits_in_buf_ = other.bits_in_buf_;
   next_byte_ = other.next_byte_;
   end_minus_8_ = other.end_minus_8_;
   first_byte_ = other.first_byte_;
   overread_bytes_ = other.overread_bytes_;
   close_called_ = other.close_called_;

   other.first_byte_ = nullptr;
   other.next_byte_ = nullptr;
   return *this;
 }
 BitReader& operator=(const BitReader& other) = delete;

 // For time-critical reads, refills can be shared by multiple reads.
 // Based on variant 4 (plus bounds-checking), see
 // fgiesen.wordpress.com/2018/02/20/reading-bits-in-far-too-many-ways-part-2/
 JXL_INLINE void Refill() {
   if (JXL_UNLIKELY(next_byte_ > end_minus_8_)) {
     BoundsCheckedRefill();
   } else {
     // It's safe to load 64 bits; insert valid (possibly nonzero) bits above
     // bits_in_buf_. The shift requires bits_in_buf_ < 64.
     buf_ |= LoadLE64(next_byte_) << bits_in_buf_;

     // Advance by bytes fully absorbed into the buffer.
     next_byte_ += (63 - bits_in_buf_) >> 3;

     // We absorbed a multiple of 8 bits, so the lower 3 bits of bits_in_buf_
     // must remain unchanged, otherwise the next refill's shifted bits will
     // not align with buf_. Set the three upper bits so the result >= 56.
     bits_in_buf_ |= 56;
     JXL_DASSERT(56 <= bits_in_buf_ && bits_in_buf_ < 64);
   }
 }

 // Returns the bits that would be returned by Read without calling Advance().
 // It is legal to PEEK at more bits than present in the bitstream (required
 // by Huffman), and those bits will be zero.
 template <size_t N>
 JXL_INLINE uint64_t PeekFixedBits() const {
   static_assert(N <= kMaxBitsPerCall, "Reading too many bits in one call.");
   JXL_DASSERT(!close_called_);
   return buf_ & ((1ULL << N) - 1);
 }

 JXL_INLINE uint64_t PeekBits(size_t nbits) const {
   JXL_DASSERT(nbits <= kMaxBitsPerCall);
   JXL_DASSERT(!close_called_);

   // Slightly faster but requires BMI2. It is infeasible to make the many
   // callers reside between begin/end_target, especially because only the
   // callers in dec_ans are time-critical. Therefore only enabled if the
   // entire binary is compiled for (and thus requires) BMI2.
#if defined(__BMI2__) && defined(__x86_64__)
   return _bzhi_u64(buf_, nbits);
#else
   const uint64_t mask = (1ULL << nbits) - 1;
   return buf_ & mask;
#endif
 }

 // Removes bits from the buffer. Need not match the previous Peek size, but
 // the buffer must contain at least num_bits (this prevents consuming more
 // than the total number of bits).
 JXL_INLINE void Consume(size_t num_bits) {
   JXL_DASSERT(!close_called_);
   JXL_DASSERT(bits_in_buf_ >= num_bits);
   if (JXL_CRASH_ON_ERROR) {
     // When JXL_CRASH_ON_ERROR is defined, it is a fatal error to read more
     // bits than available in the stream. A non-zero overread_bytes_ implies
     // that next_byte_ is already at the end of the stream, so we don't need
     // to check that.
     JXL_DASSERT(bits_in_buf_ >= num_bits + overread_bytes_ * kBitsPerByte);
   }
   bits_in_buf_ -= num_bits;
   buf_ >>= num_bits;
 }

 JXL_INLINE uint64_t ReadBits(size_t nbits) {
   JXL_DASSERT(!close_called_);
   Refill();
   const uint64_t bits = PeekBits(nbits);
   Consume(nbits);
   return bits;
 }

 template <size_t N>
 JXL_INLINE uint64_t ReadFixedBits() {
   JXL_DASSERT(!close_called_);
   Refill();
   const uint64_t bits = PeekFixedBits<N>();
   Consume(N);
   return bits;
 }

 // Equivalent to calling ReadFixedBits(1) `skip` times, but much faster.
 // `skip` is typically large.
 void SkipBits(size_t skip) {
   JXL_DASSERT(!close_called_);
   // Buffer is large enough - don't zero buf_ below.
   if (JXL_UNLIKELY(skip <= bits_in_buf_)) {
     Consume(skip);
     return;
   }

   // First deduct what we can satisfy from the buffer
   skip -= bits_in_buf_;
   bits_in_buf_ = 0;
   // Not enough to call Advance - that may leave some bits in the buffer
   // which were previously ABOVE bits_in_buf.
   buf_ = 0;

   // Skip whole bytes
   const size_t whole_bytes = skip / kBitsPerByte;
   skip %= kBitsPerByte;
   if (JXL_UNLIKELY(whole_bytes >
                    static_cast<size_t>(end_minus_8_ + 8 - next_byte_))) {
     // This is already an overflow condition (skipping past the end of the bit
     // stream). However if we increase next_byte_ too much we risk overflowing
     // that value and potentially making it valid again (next_byte_ < end).
     // This will set next_byte_ to the end of the stream and still consume
     // some bits in overread_bytes_, however the TotalBitsConsumed() will be
     // incorrect (still larger than the TotalBytes()).
     next_byte_ = end_minus_8_ + 8;
     skip += kBitsPerByte;
   } else {
     next_byte_ += whole_bytes;
   }

   Refill();
   Consume(skip);
 }

 size_t TotalBitsConsumed() const {
   const size_t bytes_read = static_cast<size_t>(next_byte_ - first_byte_);
   return (bytes_read + overread_bytes_) * kBitsPerByte - bits_in_buf_;
 }

 Status JumpToByteBoundary() {
   const size_t remainder = TotalBitsConsumed() % kBitsPerByte;
   if (remainder == 0) return true;
   if (JXL_UNLIKELY(ReadBits(kBitsPerByte - remainder) != 0)) {
     return JXL_FAILURE("Non-zero padding bits");
   }
   return true;
 }

 // For interoperability with other bitreaders (for resuming at
 // non-byte-aligned positions).
 const uint8_t* FirstByte() const { return first_byte_; }
 size_t TotalBytes() const {
   return static_cast<size_t>(end_minus_8_ + 8 - first_byte_);
 }

 // Returns whether all the bits read so far have been within the input bounds.
 // When reading past the EOF, the Read*() and Consume() functions return zeros
 // but flag a failure when calling Close() without checking this function.
 Status AllReadsWithinBounds() {
   // Mark up to which point the user checked the out of bounds condition. If
   // the user handles the condition at higher level (e.g. fetch more bytes
   // from network, return a custom JXL_FAILURE, ...), Close() should not
   // output a debug error (which would break tests with JXL_CRASH_ON_ERROR
   // even when legitimately handling the situation at higher level). This is
   // used by Bundle::CanRead.
   checked_out_of_bounds_bits_ = TotalBitsConsumed();
   if (TotalBitsConsumed() > TotalBytes() * kBitsPerByte) {
     return false;
   }
   return true;
 }

 // Close the bit reader and return whether all the previous reads were
 // successful. Close must be called once.
 Status Close() {
   JXL_DASSERT(!close_called_);
   close_called_ = true;
   if (!first_byte_) return true;
   if (TotalBitsConsumed() > checked_out_of_bounds_bits_ &&
       TotalBitsConsumed() > TotalBytes() * kBitsPerByte) {
     return JXL_FAILURE("Read more bits than available in the bit_reader");
   }
   return true;
 }

private:
 // Separate function avoids inlining this relatively cold code into callers.
 JXL_NOINLINE void BoundsCheckedRefill() {
   const uint8_t* end = end_minus_8_ + 8;

   // Read whole bytes until we have [56, 64) bits (same as LoadLE64)
   for (; bits_in_buf_ < 64 - kBitsPerByte; bits_in_buf_ += kBitsPerByte) {
     if (next_byte_ >= end) break;
     buf_ |= static_cast<uint64_t>(*next_byte_++) << bits_in_buf_;
   }
   JXL_DASSERT(bits_in_buf_ < 64);

   // Add extra bytes as 0 at the end of the stream in the bit_buffer_. If
   // these bits are read, Close() will return a failure.
   size_t extra_bytes = (63 - bits_in_buf_) / kBitsPerByte;
   overread_bytes_ += extra_bytes;
   bits_in_buf_ += extra_bytes * kBitsPerByte;

   JXL_DASSERT(bits_in_buf_ < 64);
   JXL_DASSERT(bits_in_buf_ >= 56);
 }

 JXL_NOINLINE uint32_t BoundsCheckedReadByteAlignedWord() {
   if (next_byte_ + 1 < end_minus_8_ + 8) {
     uint32_t ret = LoadLE16(next_byte_);
     next_byte_ += 2;
     return ret;
   }
   overread_bytes_ += 2;
   return 0;
 }

 uint64_t buf_;
 size_t bits_in_buf_;  // [0, 64)
 const uint8_t* JXL_RESTRICT next_byte_;
 const uint8_t* end_minus_8_;  // for refill bounds check
 const uint8_t* first_byte_;   // for GetSpan

 // Number of bytes past the end that were loaded into the buf_. These bytes
 // are not read from memory, but instead assumed 0. It is an error (likely due
 // to an invalid stream) to Consume() more bits than specified in the range
 // passed to the constructor.
 uint64_t overread_bytes_{0};
 bool close_called_{false};

 uint64_t checked_out_of_bounds_bits_{0};
};

// Closes a BitReader when the BitReaderScopedCloser goes out of scope. When
// closing the bit reader, if the status result was failure it sets this failure
// to the passed variable pointer. Typical usage.
//
// Status ret = true;
// {
//   BitReader reader(...);
//   BitReaderScopedCloser reader_closer(&reader, &ret);
//
//   // ... code that can return errors here ...
// }
// // ... more code that doesn't use the BitReader.
// return ret;

class BitReaderScopedCloser {
public:
 BitReaderScopedCloser(BitReader& reader, Status& status)
     : reader_(&reader), status_(&status) {}
 ~BitReaderScopedCloser() {
   if (reader_ != nullptr) {
     Status close_ret = reader_->Close();
     if (!close_ret) *status_ = close_ret;
   }
 }
 BitReaderScopedCloser(const BitReaderScopedCloser&) = delete;

private:
 BitReader* reader_;
 Status* status_;
};

}  // namespace jxl

#endif  // LIB_JXL_DEC_BIT_READER_H_