// Copyright 2016 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/gpu/v4l2/vp9/vp9_compressed_header_parser.h"

#include <array>
#include <cstring>

#include "base/check.h"
#include "base/logging.h"

namespace media {

namespace {

// 6.3.5 Inv remap prob syntax, inv_remap_prob().
Vp9Prob InvRemapProb(uint8_t delta_prob, uint8_t prob) {
  static const std::array<uint8_t, kVp9MaxProb> inv_map_table = {
      7,   20,  33,  46,  59,  72,  85,  98,  111, 124, 137, 150, 163, 176, 189,
      202, 215, 228, 241, 254, 1,   2,   3,   4,   5,   6,   8,   9,   10,  11,
      12,  13,  14,  15,  16,  17,  18,  19,  21,  22,  23,  24,  25,  26,  27,
      28,  29,  30,  31,  32,  34,  35,  36,  37,  38,  39,  40,  41,  42,  43,
      44,  45,  47,  48,  49,  50,  51,  52,  53,  54,  55,  56,  57,  58,  60,
      61,  62,  63,  64,  65,  66,  67,  68,  69,  70,  71,  73,  74,  75,  76,
      77,  78,  79,  80,  81,  82,  83,  84,  86,  87,  88,  89,  90,  91,  92,
      93,  94,  95,  96,  97,  99,  100, 101, 102, 103, 104, 105, 106, 107, 108,
      109, 110, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 125,
      126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 138, 139, 140, 141,
      142, 143, 144, 145, 146, 147, 148, 149, 151, 152, 153, 154, 155, 156, 157,
      158, 159, 160, 161, 162, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173,
      174, 175, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 190,
      191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 203, 204, 205, 206,
      207, 208, 209, 210, 211, 212, 213, 214, 216, 217, 218, 219, 220, 221, 222,
      223, 224, 225, 226, 227, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238,
      239, 240, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 253,
  };
  uint8_t v = delta_prob;
  DCHECK_LT(v, std::size(inv_map_table));
  return inv_map_table[v];
}

}  // namespace

Vp9CompressedHeaderParser::Vp9CompressedHeaderParser() = default;

// 6.3 Compressed header syntax
bool Vp9CompressedHeaderParser::ParseNoContext(
    base::span<const uint8_t> stream,
    const Vp9FrameHeader& frame_hdr,
    Vp9V4L2CompressedParseResult* result) {
  DCHECK(result);
  *result = Vp9V4L2CompressedParseResult{};
  if (!ParseInternal(stream, frame_hdr, &result->frame_context)) {
    return false;
  }
  result->compressed_header.tx_mode = tx_mode_;
  result->compressed_header.reference_mode = reference_mode_;
  return true;
}

// 6.3.1 Tx mode syntax
void Vp9CompressedHeaderParser::ReadTxMode(const Vp9FrameHeader& fhdr) {
  int tx_mode;
  if (fhdr.quant_params.IsLossless()) {
    tx_mode = Vp9CompressedHeader::ONLY_4X4;
  } else {
    tx_mode = reader_.ReadLiteral(2);
    if (tx_mode == Vp9CompressedHeader::ALLOW_32X32) {
      tx_mode += reader_.ReadLiteral(1);
    }
  }
  tx_mode_ = static_cast<Vp9CompressedHeader::Vp9TxMode>(tx_mode);
}

// 6.3.4 Decode term subexp syntax
uint8_t Vp9CompressedHeaderParser::DecodeTermSubexp() {
  if (reader_.ReadLiteral(1) == 0) {
    return reader_.ReadLiteral(4);
  }
  if (reader_.ReadLiteral(1) == 0) {
    return reader_.ReadLiteral(4) + 16;
  }
  if (reader_.ReadLiteral(1) == 0) {
    return reader_.ReadLiteral(5) + 32;
  }
  uint8_t v = reader_.ReadLiteral(7);
  if (v < 65) {
    return v + 64;
  }
  return (v << 1) - 1 + reader_.ReadLiteral(1);
}

// 6.3.3 Diff update prob syntax
void Vp9CompressedHeaderParser::DiffUpdateProb(Vp9Prob* prob) {
  const Vp9Prob kUpdateProb = 252;
  const bool must_update_probabilities = reader_.ReadBool(kUpdateProb);

  if (!must_update_probabilities) {
    return;
  }

  uint8_t delta_prob = DecodeTermSubexp();
  *prob = InvRemapProb(delta_prob, *prob);
}

// Helper function to DiffUpdateProb an array of probs.
template <int N>
void Vp9CompressedHeaderParser::DiffUpdateProbArray(Vp9Prob (&prob_array)[N]) {
  for (auto& x : prob_array) {
    DiffUpdateProb(&x);
  }
}

// 6.3.2 Tx mode probs syntax
void Vp9CompressedHeaderParser::ReadTxModeProbs(
    Vp9FrameContext* frame_context) {
  for (auto& a : frame_context->tx_probs_8x8) {
    DiffUpdateProbArray(a);
  }
  for (auto& a : frame_context->tx_probs_16x16) {
    DiffUpdateProbArray(a);
  }
  for (auto& a : frame_context->tx_probs_32x32) {
    DiffUpdateProbArray(a);
  }
}

// 6.3.7 Coef probs syntax
void Vp9CompressedHeaderParser::ReadCoefProbs(Vp9FrameContext* frame_context) {
  const std::array<int, Vp9CompressedHeader::TX_MODES>
      tx_mode_to_biggest_tx_size = {
          0, 1, 2, 3, 3,
      };
  const int max_tx_size = tx_mode_to_biggest_tx_size[tx_mode_];
  const size_t num_tx_sizes = base::checked_cast<size_t>(max_tx_size + 1);
  auto coef_probs = base::span(frame_context->coef_probs).first(num_tx_sizes);
  for (auto& tx_size_coef_probs : coef_probs) {
    if (reader_.ReadLiteral(1) == 0) {
      continue;
    }

    for (auto& ai : tx_size_coef_probs) {
      for (auto& aj : ai) {
        for (auto& ak : aj) {
          const int max_l = (+ak == +aj[0]) ? 3 : 6;
          auto probs = base::span(ak).first(base::checked_cast<size_t>(max_l));
          for (auto& prob : probs) {
            DiffUpdateProbArray(prob);
          }
        }
      }
    }
  }
}

// 6.3.8 Skip probs syntax
void Vp9CompressedHeaderParser::ReadSkipProb(Vp9FrameContext* frame_context) {
  DiffUpdateProbArray(frame_context->skip_prob);
}

// 6.3.9 Inter mode probs syntax
void Vp9CompressedHeaderParser::ReadInterModeProbs(
    Vp9FrameContext* frame_context) {
  for (auto& a : frame_context->inter_mode_probs) {
    DiffUpdateProbArray(a);
  }
}

// 6.3.10 Interp filter probs syntax
void Vp9CompressedHeaderParser::ReadInterpFilterProbs(
    Vp9FrameContext* frame_context) {
  for (auto& a : frame_context->interp_filter_probs) {
    DiffUpdateProbArray(a);
  }
}

// 6.3.11 Intra inter probs syntax
void Vp9CompressedHeaderParser::ReadIsInterProbs(
    Vp9FrameContext* frame_context) {
  DiffUpdateProbArray(frame_context->is_inter_prob);
}

// 6.3.12 Frame reference mode syntax
void Vp9CompressedHeaderParser::ReadFrameReferenceMode(
    const Vp9FrameHeader& fhdr) {
  bool compound_reference_allowed = false;
  for (int i = VP9_FRAME_LAST + 1; i < VP9_FRAME_MAX; i++) {
    if (fhdr.ref_frame_sign_bias[i] != fhdr.ref_frame_sign_bias[1]) {
      compound_reference_allowed = true;
    }
  }

  if (compound_reference_allowed && reader_.ReadLiteral(1)) {
    reference_mode_ =
        reader_.ReadLiteral(1) ? REFERENCE_MODE_SELECT : COMPOUND_REFERENCE;
  } else {
    reference_mode_ = SINGLE_REFERENCE;
  }
}

// 6.3.13 Frame reference mode probs syntax
void Vp9CompressedHeaderParser::ReadFrameReferenceModeProbs(
    Vp9FrameContext* frame_context) {
  if (reference_mode_ == REFERENCE_MODE_SELECT) {
    DiffUpdateProbArray(frame_context->comp_mode_prob);
  }

  if (reference_mode_ != COMPOUND_REFERENCE) {
    for (auto& a : frame_context->single_ref_prob) {
      DiffUpdateProbArray(a);
    }
  }

  if (reference_mode_ != SINGLE_REFERENCE) {
    DiffUpdateProbArray(frame_context->comp_ref_prob);
  }
}

// 6.3.14 Y mode probs syntax
void Vp9CompressedHeaderParser::ReadYModeProbs(Vp9FrameContext* frame_context) {
  for (auto& a : frame_context->y_mode_probs) {
    DiffUpdateProbArray(a);
  }
}

// 6.3.15 Partition probs syntax
void Vp9CompressedHeaderParser::ReadPartitionProbs(
    Vp9FrameContext* frame_context) {
  for (auto& a : frame_context->partition_probs) {
    DiffUpdateProbArray(a);
  }
}

// 6.3.16 MV probs syntax
void Vp9CompressedHeaderParser::ReadMvProbs(bool allow_high_precision_mv,
                                            Vp9FrameContext* frame_context) {
  UpdateMvProbArray(frame_context->mv_joint_probs);

  auto mv_sign_prob = base::span(frame_context->mv_sign_prob);
  auto mv_class_probs = base::span(frame_context->mv_class_probs);
  auto mv_class0_bit_prob = base::span(frame_context->mv_class0_bit_prob);
  auto mv_bits_prob = base::span(frame_context->mv_bits_prob);
  auto mv_class0_fr_probs = base::span(frame_context->mv_class0_fr_probs);
  auto mv_fr_probs = base::span(frame_context->mv_fr_probs);
  auto mv_class0_hp_prob = base::span(frame_context->mv_class0_hp_prob);
  auto mv_hp_prob = base::span(frame_context->mv_hp_prob);

  for (size_t i = 0; i < mv_sign_prob.size(); ++i) {
    UpdateMvProb(&mv_sign_prob.at(i));
    UpdateMvProbArray(mv_class_probs.at(i));
    UpdateMvProb(&mv_class0_bit_prob.at(i));
    UpdateMvProbArray(mv_bits_prob.at(i));
  }

  for (size_t i = 0; i < mv_class0_fr_probs.size(); ++i) {
    for (auto& a : mv_class0_fr_probs.at(i)) {
      UpdateMvProbArray(a);
    }
    UpdateMvProbArray(mv_fr_probs.at(i));
  }

  if (allow_high_precision_mv) {
    for (size_t i = 0; i < mv_class0_hp_prob.size(); ++i) {
      UpdateMvProb(&mv_class0_hp_prob.at(i));
      UpdateMvProb(&mv_hp_prob.at(i));
    }
  }
}

// 6.3.17 Update mv prob syntax
void Vp9CompressedHeaderParser::UpdateMvProb(Vp9Prob* prob) {
  if (reader_.ReadBool(252)) {
    *prob = reader_.ReadLiteral(7) << 1 | 1;
  }
}

// Helper function to UpdateMvProb an array of probs.
template <int N>
void Vp9CompressedHeaderParser::UpdateMvProbArray(Vp9Prob (&prob_array)[N]) {
  for (auto& x : prob_array) {
    UpdateMvProb(&x);
  }
}

// 6.3 Compressed header syntax
bool Vp9CompressedHeaderParser::ParseInternal(base::span<const uint8_t> stream,
                                              const Vp9FrameHeader& fhdr,
                                              Vp9FrameContext* frame_context) {
  DCHECK(frame_context);

  DVLOG(2) << "Vp9CompressedHeaderParser::Parse";
  if (!reader_.Initialize(stream)) {
    return false;
  }

  ReadTxMode(fhdr);
  if (tx_mode_ == Vp9CompressedHeader::TX_MODE_SELECT) {
    ReadTxModeProbs(frame_context);
  }

  ReadCoefProbs(frame_context);
  ReadSkipProb(frame_context);

  if (!fhdr.IsIntra()) {
    ReadInterModeProbs(frame_context);
    if (fhdr.interpolation_filter == SWITCHABLE) {
      ReadInterpFilterProbs(frame_context);
    }
    ReadIsInterProbs(frame_context);
    ReadFrameReferenceMode(fhdr);
    ReadFrameReferenceModeProbs(frame_context);
    ReadYModeProbs(frame_context);
    ReadPartitionProbs(frame_context);
    ReadMvProbs(fhdr.allow_high_precision_mv, frame_context);
  }

  if (!reader_.IsValid()) {
    LOG(ERROR) << "parser reads beyond the end of buffer";
    return false;
  }
  if (!reader_.ConsumePaddingBits()) {
    LOG(ERROR) << "padding bits are not zero";
    return false;
  }
  return true;
}

}  // namespace media