// 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_