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

#define LIB_JXL_FIELDS_H_

// Forward/backward-compatible 'bundles' with auto-serialized 'fields'.

#include <cstddef>

#include <cstdint>

#include <cstdlib>

#include <cstring>

#include "lib/jxl/base/compiler_specific.h"

#include "lib/jxl/base/status.h"

#include "lib/jxl/dec_bit_reader.h"

#include "lib/jxl/field_encodings.h"

namespace jxl {

struct AuxOut;

struct BitWriter;

// Integer coders: BitsCoder (raw), U32Coder (table), U64Coder (varint).

// Reads/writes a given (fixed) number of bits <= 32.

namespace BitsCoder {

size_t MaxEncodedBits(size_t bits);

Status CanEncode(size_t bits, uint32_t value,

                 size_t* JXL_RESTRICT encoded_bits);

uint32_t Read(size_t bits, BitReader* JXL_RESTRICT reader);

// Returns false if the value is too large to encode.

Status Write(size_t bits, uint32_t value, BitWriter* JXL_RESTRICT writer);

}  // namespace BitsCoder

// Encodes u32 using a lookup table and/or extra bits, governed by a per-field

// encoding `enc` which consists of four distributions `d` chosen via a 2-bit

// selector (least significant = 0). Each d may have two modes:

// - direct: if d.IsDirect(), the value is d.Direct();

// - offset: the value is derived from d.ExtraBits() extra bits plus d.Offset();

// This encoding is denser than Exp-Golomb or Gamma codes when both small and

// large values occur.

//

// Examples:

// Direct: U32Enc(Val(8), Val(16), Val(32), Bits(6)), value 32 => 10b.

// Offset: U32Enc(Val(0), BitsOffset(1, 1), BitsOffset(2, 3), BitsOffset(8, 8))

//   defines the following prefix code:

//   00 -> 0

//   01x -> 1..2

//   10xx -> 3..7

//   11xxxxxxxx -> 8..263

namespace U32Coder {

size_t MaxEncodedBits(U32Enc enc);

Status CanEncode(U32Enc enc, uint32_t value, size_t* JXL_RESTRICT encoded_bits);

uint32_t Read(U32Enc enc, BitReader* JXL_RESTRICT reader);

// Returns false if the value is too large to encode.

Status Write(U32Enc enc, uint32_t value, BitWriter* JXL_RESTRICT writer);

// "private"

Status ChooseSelector(U32Enc enc, uint32_t value,

                      uint32_t* JXL_RESTRICT selector,

                      size_t* JXL_RESTRICT total_bits);

}  // namespace U32Coder

// Encodes 64-bit unsigned integers with a fixed distribution, taking 2 bits

// to encode 0, 6 bits to encode 1 to 16, 10 bits to encode 17 to 272, 15 bits

// to encode up to 4095, and on the order of log2(value) * 1.125 bits for

// larger values.

namespace U64Coder {

constexpr size_t MaxEncodedBits() { return 2 + 12 + 6 * (8 + 1) + (4 + 1); }

uint64_t Read(BitReader* JXL_RESTRICT reader);

// Returns false if the value is too large to encode.

Status Write(uint64_t value, BitWriter* JXL_RESTRICT writer);

// Can always encode, but useful because it also returns bit size.

Status CanEncode(uint64_t value, size_t* JXL_RESTRICT encoded_bits);

}  // namespace U64Coder

// IEEE 754 half-precision (binary16). Refuses to read/write NaN/Inf.

namespace F16Coder {

constexpr size_t MaxEncodedBits() { return 16; }

// Returns false if the bit representation is NaN or infinity

Status Read(BitReader* JXL_RESTRICT reader, float* JXL_RESTRICT value);

// Returns false if the value is too large to encode.

Status Write(float value, BitWriter* JXL_RESTRICT writer);

Status CanEncode(float value, size_t* JXL_RESTRICT encoded_bits);

}  // namespace F16Coder

// A "bundle" is a forward- and backward compatible collection of fields.

// They are used for SizeHeader/FrameHeader/GroupHeader. Bundles can be

// extended by appending(!) fields. Optional fields may be omitted from the

// bitstream by conditionally visiting them. When reading new bitstreams with

// old code, we skip unknown fields at the end of the bundle. This requires

// storing the amount of extra appended bits, and that fields are visited in

// chronological order of being added to the format, because old decoders

// cannot skip some future fields and resume reading old fields. Similarly,

// new readers query bits in an "extensions" field to skip (groups of) fields

// not present in old bitstreams. Note that each bundle must include an

// "extensions" field prior to freezing the format, otherwise it cannot be

// extended.

//

// To ensure interoperability, there will be no opaque fields.

//

// HOWTO:

// - basic usage: define a struct with member variables ("fields") and a

//   VisitFields(v) member function that calls v->U32/Bool etc. for each

//   field, specifying their default values. The ctor must call

//   Bundle::Init(this).

//

// - print a trace of visitors: ensure each bundle has a static Name() member

//   function, and change Bundle::Print* to return true.

//

// - optional fields: in VisitFields, add if (v->Conditional(your_condition))

//   { v->Bool(default, &field); }. This prevents reading/writing field

//   if !your_condition, which is typically computed from a prior field.

//   WARNING: to ensure all fields are initialized, do not add an else branch;

//   instead add another if (v->Conditional(!your_condition)).

//

// - repeated fields: for dynamic sizes, use e.g. std::vector and in

//   VisitFields, if (v->IsReading()) field.resize(size) before accessing field.

//   For static or bounded sizes, use an array or std::array. In all cases,

//   simply visit each array element as if it were a normal field.

//

// - nested bundles: add a bundle as a normal field and in VisitFields call

//   JXL_RETURN_IF_ERROR(v->VisitNested(&nested));

//

// - allow future extensions: define a "uint64_t extensions" field and call

//   v->BeginExtensions(&extensions) after visiting all non-extension fields,

//   and `return v->EndExtensions();` after the last extension field.

//

// - encode an entire bundle in one bit if ALL its fields equal their default

//   values: add a "mutable bool all_default" field and as the first visitor:

//   if (v->AllDefault(*this, &all_default)) {

//     // Overwrite all serialized fields, but not any nonserialized_*.

//     v->SetDefault(this);

//     return true;

//   }

//   Note: if extensions are present, AllDefault() == false.

namespace Bundle {

constexpr size_t kMaxExtensions = 64;  // bits in u64

// Initializes fields to the default values. It is not recursive to nested

// fields, this function is intended to be called in the constructors so

// each nested field will already Init itself.

void Init(Fields* JXL_RESTRICT fields);

// Similar to Init, but recursive to nested fields.

void SetDefault(Fields* JXL_RESTRICT fields);

// Returns whether ALL fields (including `extensions`, if present) are equal

// to their default value.

bool AllDefault(const Fields& fields);

// Returns max number of bits required to encode a T.

size_t MaxBits(const Fields& fields);

// Returns whether a header's fields can all be encoded, i.e. they have a

// valid representation. If so, "*total_bits" is the exact number of bits

// required. Called by Write.

Status CanEncode(const Fields& fields, size_t* JXL_RESTRICT extension_bits,

                 size_t* JXL_RESTRICT total_bits);

Status Read(BitReader* reader, Fields* JXL_RESTRICT fields);

// Returns whether enough bits are available to fully read this bundle using

// Read. Also returns true in case of a codestream error (other than not being

// large enough): that means enough bits are available to determine there's an

// error, use Read to get such error status.

// NOTE: this advances the BitReader, a different one pointing back at the

// original bit position in the codestream must be created to use Read after

// this.

bool CanRead(BitReader* reader, Fields* JXL_RESTRICT fields);

Status Write(const Fields& fields, BitWriter* JXL_RESTRICT writer, size_t layer,

             AuxOut* aux_out);

}  // namespace Bundle

// Different subclasses of Visitor are passed to implementations of Fields

// throughout their lifetime. Templates used to be used for this but dynamic

// polymorphism produces more compact executables than template reification did.

class Visitor {

 public:

  virtual ~Visitor() = default;

  virtual Status Visit(Fields* fields) = 0;

  virtual Status Bool(bool default_value, bool* JXL_RESTRICT value) = 0;

  virtual Status U32(U32Enc, uint32_t, uint32_t*) = 0;

  // Helper to construct U32Enc from U32Distr.

  Status U32(const U32Distr d0, const U32Distr d1, const U32Distr d2,

             const U32Distr d3, const uint32_t default_value,

             uint32_t* JXL_RESTRICT value) {

    return U32(U32Enc(d0, d1, d2, d3), default_value, value);

  }

  template <typename EnumT>

  Status Enum(const EnumT default_value, EnumT* JXL_RESTRICT value) {

    uint32_t u32 = static_cast<uint32_t>(*value);

    // 00 -> 0

    // 01 -> 1

    // 10xxxx -> 2..17

    // 11yyyyyy -> 18..81

    JXL_RETURN_IF_ERROR(U32(Val(0), Val(1), BitsOffset(4, 2), BitsOffset(6, 18),

                            static_cast<uint32_t>(default_value), &u32));

    *value = static_cast<EnumT>(u32);

    return EnumValid(*value);

  }

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::cms::TransferFunction>(jxl::cms::TransferFunction, jxl::cms::TransferFunction*)

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::cms::ColorSpace>(jxl::cms::ColorSpace, jxl::cms::ColorSpace*)

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::cms::WhitePoint>(jxl::cms::WhitePoint, jxl::cms::WhitePoint*)

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::cms::Primaries>(jxl::cms::Primaries, jxl::cms::Primaries*)

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::cms::RenderingIntent>(jxl::cms::RenderingIntent, jxl::cms::RenderingIntent*)

Unexecuted instantiation: jxl::Status jxl::Visitor::Enum<jxl::ExtraChannel>(jxl::ExtraChannel, jxl::ExtraChannel*)

  virtual Status Bits(size_t bits, uint32_t default_value,

                      uint32_t* JXL_RESTRICT value) = 0;

  virtual Status U64(uint64_t default_value, uint64_t* JXL_RESTRICT value) = 0;

  virtual Status F16(float default_value, float* JXL_RESTRICT value) = 0;

  // Returns whether VisitFields should visit some subsequent fields.

  // "condition" is typically from prior fields, e.g. flags.

  // Overridden by InitVisitor and MaxBitsVisitor.

  virtual Status Conditional(bool condition) { return condition; }

  // Overridden by InitVisitor, AllDefaultVisitor and CanEncodeVisitor.

  virtual Status AllDefault(const Fields& /*fields*/,

                            bool* JXL_RESTRICT all_default) {

    JXL_RETURN_IF_ERROR(Bool(true, all_default));

    return *all_default;

  }

  virtual void SetDefault(Fields* /*fields*/) {

    // Do nothing by default, this is overridden by ReadVisitor.

  }

  // Returns the result of visiting a nested Bundle.

  // Overridden by InitVisitor.

  virtual Status VisitNested(Fields* fields) { return Visit(fields); }

  // Overridden by ReadVisitor. Enables dynamically-sized fields.

  virtual bool IsReading() const { return false; }

  virtual Status BeginExtensions(uint64_t* JXL_RESTRICT extensions) = 0;

  virtual Status EndExtensions() = 0;

};

namespace fields_internal {

// A bundle can be in one of three states concerning extensions: not-begun,

// active, ended. Bundles may be nested, so we need a stack of states.

class ExtensionStates {

 public:

  void Push() {

    // Initial state = not-begun.

    begun_ <<= 1;

    ended_ <<= 1;

  }

  // Clears current state; caller must check IsEnded beforehand.

  void Pop() {

    begun_ >>= 1;

    ended_ >>= 1;

  }

  // Returns true if state == active || state == ended.

  Status IsBegun() const { return (begun_ & 1) != 0; }

  // Returns true if state != not-begun && state != active.

  Status IsEnded() const { return (ended_ & 1) != 0; }

  void Begin() {

    JXL_ASSERT(!IsBegun());

    JXL_ASSERT(!IsEnded());

    begun_ += 1;

  }

  void End() {

    JXL_ASSERT(IsBegun());

    JXL_ASSERT(!IsEnded());

    ended_ += 1;

  }

 private:

  // Current state := least-significant bit of begun_ and ended_.

  uint64_t begun_ = 0;

  uint64_t ended_ = 0;

};

// Visitors generate Init/AllDefault/Read/Write logic for all fields. Each

// bundle's VisitFields member function calls visitor->U32 etc. We do not

// overload operator() because a function name is easier to search for.

class VisitorBase : public Visitor {

 public:

  explicit VisitorBase() = default;

  ~VisitorBase() override { JXL_ASSERT(depth_ == 0); }

  // This is the only call site of Fields::VisitFields.

  // Ensures EndExtensions was called.

  Status Visit(Fields* fields) override {

    depth_ += 1;

    JXL_ASSERT(depth_ <= Bundle::kMaxExtensions);

    extension_states_.Push();

    const Status ok = fields->VisitFields(this);

    if (ok) {

      // If VisitFields called BeginExtensions, must also call

      // EndExtensions.

      JXL_ASSERT(!extension_states_.IsBegun() || extension_states_.IsEnded());

    } else {

      // Failed, undefined state: don't care whether EndExtensions was

      // called.

    }

    extension_states_.Pop();

    JXL_ASSERT(depth_ != 0);

    depth_ -= 1;

    return ok;

  }

  // For visitors accepting a const Visitor, need to const-cast so we can call

  // the non-const Visitor::VisitFields. NOTE: C is not modified except the

  // `all_default` field by CanEncodeVisitor.

  Status VisitConst(const Fields& t) { return Visit(const_cast<Fields*>(&t)); }

  // Derived types (overridden by InitVisitor because it is unsafe to read

  // from *value there)

  Status Bool(bool default_value, bool* JXL_RESTRICT value) override {

    uint32_t bits = *value ? 1 : 0;

    JXL_RETURN_IF_ERROR(Bits(1, static_cast<uint32_t>(default_value), &bits));

    JXL_DASSERT(bits <= 1);

    *value = bits == 1;

    return true;

  }

  // Overridden by ReadVisitor and WriteVisitor.

  // Called before any conditional visit based on "extensions".

  // Overridden by ReadVisitor, CanEncodeVisitor and WriteVisitor.

  Status BeginExtensions(uint64_t* JXL_RESTRICT extensions) override {

    JXL_RETURN_IF_ERROR(U64(0, extensions));

    extension_states_.Begin();

    return true;

  }

  // Called after all extension fields (if any). Although non-extension

  // fields could be visited afterward, we prefer the convention that

  // extension fields are always the last to be visited. Overridden by

  // ReadVisitor.

  Status EndExtensions() override {

    extension_states_.End();

    return true;

  }

 private:

  size_t depth_ = 0;  // to check nesting

  ExtensionStates extension_states_;

};

}  // namespace fields_internal

Status CheckHasEnoughBits(Visitor* visitor, size_t bits);

}  // namespace jxl

#endif  // LIB_JXL_FIELDS_H_