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

#define LIB_JXL_ENC_BIT_WRITER_H_

// BitWriter class: unbuffered writes using unaligned 64-bit stores.

#include <jxl/memory_manager.h>

#include <cstddef>

#include <cstdint>

#include <functional>

#include <memory>

#include <utility>

#include <vector>

#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"

#include "lib/jxl/padded_bytes.h"

namespace jxl {

struct AuxOut;

enum class LayerType : uint8_t;

struct BitWriter {

  // Upper bound on `n_bits` in each call to Write. We shift a 64-bit word by

  // 7 bits (max already valid bits in the last byte) and at least 1 bit is

  // needed to zero-initialize the bit-stream ahead (i.e. if 7 bits are valid

  // and we write 57 bits, then the next write will access a byte that was not

  // yet zero-initialized).

  static constexpr size_t kMaxBitsPerCall = 56;

  explicit BitWriter(JxlMemoryManager* memory_manager)

      : bits_written_(0), storage_(memory_manager) {}

  // Disallow copying - may lead to bugs.

  BitWriter(const BitWriter&) = delete;

  BitWriter& operator=(const BitWriter&) = delete;

  BitWriter(BitWriter&&) = default;

  BitWriter& operator=(BitWriter&&) = default;

  size_t BitsWritten() const { return bits_written_; }

  JxlMemoryManager* memory_manager() const { return storage_.memory_manager(); }

  Span<const uint8_t> GetSpan() const {

    // Callers must ensure byte alignment to avoid uninitialized bits.

    JXL_DASSERT(bits_written_ % kBitsPerByte == 0);

    return Bytes(storage_.data(), DivCeil(bits_written_, kBitsPerByte));

  }

  // Example usage: bytes = std::move(writer).TakeBytes(); Useful for the

  // top-level encoder which returns PaddedBytes, not a BitWriter.

  // *this must be an rvalue reference and is invalid afterwards.

  PaddedBytes&& TakeBytes() && {

    // Callers must ensure byte alignment to avoid uninitialized bits.

    JXL_DASSERT(bits_written_ % kBitsPerByte == 0);

    Status status = storage_.resize(DivCeil(bits_written_, kBitsPerByte));

    JXL_DASSERT(status);

    // Can never fail, because we are resizing to a lower size.

    (void)status;

    return std::move(storage_);

  }

  Status Shrink() { return storage_.shrink_to_fit(); }

  // Must be byte-aligned before calling.

  Status AppendByteAligned(const Span<const uint8_t>& span);

  // NOTE: no allotment needed, the other BitWriters have already been charged.

  Status AppendByteAligned(

      const std::vector<std::unique_ptr<BitWriter>>& others);

  Status AppendUnaligned(const BitWriter& other);

  // Writes bits into bytes in increasing addresses, and within a byte

  // least-significant-bit first.

  //

  // The function can write up to 56 bits in one go.

  void Write(size_t n_bits, uint64_t bits);

  // This should only rarely be used - e.g. when the current location will be

  // referenced via byte offset (TOCs point to groups), or byte-aligned reading

  // is required for speed.

  void ZeroPadToByte() {

    const size_t remainder_bits =

        RoundUpBitsToByteMultiple(bits_written_) - bits_written_;

    if (remainder_bits == 0) return;

    Write(remainder_bits, 0);

    JXL_DASSERT(bits_written_ % kBitsPerByte == 0);

  }

  Status WithMaxBits(size_t max_bits, LayerType layer,

                     AuxOut* JXL_RESTRICT aux_out,

                     const std::function<Status()>& function,

                     bool finished_histogram = false);

 private:

  class Allotment {

   public:

    explicit Allotment(size_t max_bits);

    ~Allotment();

    Allotment(const Allotment& other) = delete;

    Allotment(Allotment&& other) = delete;

    Allotment& operator=(const Allotment&) = delete;

    Allotment& operator=(Allotment&&) = delete;

    // Call after writing a histogram, but before ReclaimUnused.

    Status FinishedHistogram(BitWriter* JXL_RESTRICT writer);

    size_t HistogramBits() const {

      JXL_DASSERT(called_);

      return histogram_bits_;

    }

    Status ReclaimAndCharge(BitWriter* JXL_RESTRICT writer, LayerType layer,

                            AuxOut* JXL_RESTRICT aux_out);

   private:

    friend struct BitWriter;

    // Expands a BitWriter's storage. Must happen before calling Write or

    // ZeroPadToByte. Must call ReclaimUnused after writing to reclaim the

    // unused storage so that BitWriter memory use remains tightly bounded.

    Status Init(BitWriter* JXL_RESTRICT writer);

    Status PrivateReclaim(BitWriter* JXL_RESTRICT writer,

                          size_t* JXL_RESTRICT used_bits,

                          size_t* JXL_RESTRICT unused_bits);

    size_t prev_bits_written_;

    const size_t max_bits_;

    size_t histogram_bits_ = 0;

    bool called_ = false;

    Allotment* parent_;

  };

  size_t bits_written_;

  PaddedBytes storage_;

  Allotment* current_allotment_ = nullptr;

};

}  // namespace jxl

#endif  // LIB_JXL_ENC_BIT_WRITER_H_