libjxl

dec_bit_reader.h (12140B)

---

 // 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 <stddef.h>
 #include <stdint.h>
 #include <string.h>  // 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/span.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_ASSERT(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_ASSERT(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);
 #ifdef 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_ASSERT(bits_in_buf_ >= num_bits + overread_bytes_ * kBitsPerByte);
 #endif
     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 span of the remaining (unconsumed) bytes, e.g. for passing to
   // external decoders such as Brotli.
   Span<const uint8_t> GetSpan() const {
     JXL_DASSERT(first_byte_ != nullptr);
     JXL_ASSERT(TotalBitsConsumed() % kBitsPerByte == 0);
     const size_t offset = TotalBitsConsumed() / kBitsPerByte;  // no remainder
     JXL_ASSERT(offset <= TotalBytes());
     return Bytes(first_byte_ + offset, TotalBytes() - offset);
   }
 
   // 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) {
     JXL_DASSERT(reader_ != nullptr);
     JXL_DASSERT(status_ != nullptr);
   }
   ~BitReaderScopedCloser() {
     if (reader_ != nullptr) {
       Status close_ret = reader_->Close();
       if (!close_ret) *status_ = close_ret;
     }
   }
   void CloseAndSuppressError() {
     JXL_ASSERT(reader_ != nullptr);
     (void)reader_->Close();
     reader_ = nullptr;
   }
   BitReaderScopedCloser(const BitReaderScopedCloser&) = delete;
 
  private:
   BitReader* reader_;
   Status* status_;
 };
 
 }  // namespace jxl
 
 #endif  // LIB_JXL_DEC_BIT_READER_H_