/*

 * H.265 video codec.

 * Copyright (c) 2013-2014 struktur AG, Dirk Farin <farin@struktur.de>

 *

 * This file is part of libde265.

 *

 * libde265 is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libde265 is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libde265.  If not, see <http://www.gnu.org/licenses/>.

 */

#include "nal-parser.h"

#include <string.h>

#include <assert.h>

#include <stdlib.h>

#include <stdio.h>

#include <stdint.h>

#include <limits.h>

#include <utility>

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

NAL_unit::~NAL_unit()

{

  free(nal_data);

}

[[nodiscard]] bool NAL_unit::resize(int new_size)

{

  if (capacity < new_size) {

    // Grow the buffer geometrically (1.5x) rather than to the exact requested

    // size. NAL_Parser::push_data() appends to the pending NAL one input chunk

    // at a time, increasing the request by a roughly constant amount each call.

    // With exact-size allocation every chunk would reallocate and copy the

    // whole accumulated buffer (O(n^2) for a single oversized NAL); spare

    // capacity amortizes the total copying to O(n). Here new_size > capacity >= 0,

    // so the 1.5x term is computed in 64 bits and only used when it both exceeds

    // the request and still fits in 'int'.

    int alloc_size = new_size;

    int64_t grow = static_cast<int64_t>(capacity) + capacity / 2;

    if (grow > new_size && grow <= INT_MAX) {

      alloc_size = static_cast<int>(grow);

    }

    unsigned char* newbuffer = static_cast<unsigned char*>(malloc(alloc_size));

    if (newbuffer == nullptr) {

      return false;

    }

    if (nal_data != nullptr) {

      memcpy(newbuffer, nal_data, data_size);

      free(nal_data);

    }

    nal_data = newbuffer;

    capacity = alloc_size;

  }

  return true;

}

[[nodiscard]] bool NAL_unit::append(const unsigned char* in_data, int n)

{

  if (!resize(data_size + n)) {

    return false;

  }

  if (n > 0) {

    memcpy(nal_data + data_size, in_data, n);

  }

  data_size += n;

  return true;

}

[[nodiscard]] bool NAL_unit::set_data(const unsigned char* in_data, int n)

{

  if (!resize(n)) {

    return false;

  }

  if (n > 0) {

    memcpy(nal_data, in_data, n);

  }

  data_size = n;

  return true;

}

void NAL_unit::insert_skipped_byte(uint32_t pos)

{

  skipped_bytes.push_back(pos);

}

uint32_t NAL_unit::num_skipped_bytes_before(uint32_t byte_position, uint32_t headerLength) const

{

  if (skipped_bytes.empty()) {

    return 0;

  }

  for (int k=skipped_bytes.size()-1;k>=0;k--)

    if (skipped_bytes[k] >= headerLength &&

        skipped_bytes[k]-headerLength <= byte_position) {

      return k+1;

    }

  return 0;

}

void NAL_unit::remove_stuffing_bytes()

{

  // Remove emulation-prevention bytes: every 0x03 that immediately follows two

  // 0x00 bytes is dropped (and the zero-run reset, so 00 00 03 03 keeps the

  // trailing 03). This is done in a single in-place forward-compaction pass in

  // O(n) time. A previous implementation called memmove() on the remaining tail

  // for each removed byte, which is O(n^2) and can be abused by a payload that

  // is densely packed with 00 00 03 triplets.

  uint8_t* d = data();

  const int n = size();

  int w = 0;       // write position == length of the compacted output so far

  int zeros = 0;   // number of consecutive 0x00 bytes already written to output

  for (int r=0; r<n; r++) {

    uint8_t b = d[r];

    if (zeros >= 2 && b == 3) {

      // 'r' is the position of this byte in the original (uncompacted) NAL,

      // which equals (compacted position) + num_skipped_bytes() — the value the

      // previous memmove-based code recorded here.

      insert_skipped_byte(r);

      zeros = 0;

      continue;

    }

    d[w++] = b;

    zeros = (b == 0) ? zeros + 1 : 0;

  }

  set_size(w);

}

NAL_Parser::NAL_Parser() = default;

NAL_Parser::~NAL_Parser()

{

  // The NAL queue and the pending input NAL hold owning unique_ptrs, so their

  // contents are released automatically. Nothing to do.

}

[[nodiscard]] std::unique_ptr<NAL_unit> NAL_Parser::alloc_NAL_unit(int size)

{

  // A freshly constructed NAL_unit is already in the cleared state (empty

  // buffer, empty skipped-byte list), so no clear() is needed here.

  auto nal = std::make_unique<NAL_unit>();

  if (!nal->resize(size)) {

    return nullptr;   // 'nal' is deleted as it goes out of scope

  }

  return nal;

}

void NAL_Parser::free_NAL_unit(std::unique_ptr<NAL_unit> /*nal*/)

{

  // Releasing a NAL is just destroying it: ownership is moved in by value, so the

  // NAL is deleted when the argument goes out of scope here (a moved-from / null

  // argument is a harmless no-op). Kept as a named operation so call sites read as

  // an explicit release, and because it makes a double release impossible to express.

}

std::unique_ptr<NAL_unit> NAL_Parser::pop_from_NAL_queue()

{

  if (NAL_queue.empty()) {

    return nullptr;

  }

  else {

    std::unique_ptr<NAL_unit> nal = std::move(NAL_queue.front());

    NAL_queue.pop();

    nBytes_in_NAL_queue -= nal->size();

    return nal;

  }

}

void NAL_Parser::push_to_NAL_queue(std::unique_ptr<NAL_unit> nal)

{

  nBytes_in_NAL_queue += nal->size();

  NAL_queue.push(std::move(nal));

}

de265_error NAL_Parser::push_data(const unsigned char* data, int len,

                                  de265_PTS pts, void* user_data)

{

  end_of_frame = false;

  if (pending_input_NAL == nullptr) {

    pending_input_NAL = alloc_NAL_unit(len+3);

    if (pending_input_NAL == nullptr) {

      return DE265_ERROR_OUT_OF_MEMORY;

    }

    pending_input_NAL->pts = pts;

    pending_input_NAL->user_data = user_data;

  }

  // Raw working pointer for byte access; ownership stays in pending_input_NAL.

  NAL_unit* nal = pending_input_NAL.get(); // shortcut

  // Resize output buffer so that complete input would fit.

  // We add 3, because in the worst case 3 extra bytes are created for an input byte.

  if (!nal->resize(nal->size() + len + 3)) {

    return DE265_ERROR_OUT_OF_MEMORY;

  }

  unsigned char* out = nal->data() + nal->size();

  for (int i=0;i<len;i++) {

    /*

    printf("state=%d input=%02x (%p) (output size: %d)\n",ctx->input_push_state, *data, data,

           out - ctx->nal_data.data);

    */

    switch (input_push_state) {

    case 0:

    case 1:

      if (*data == 0) { input_push_state++; }

      else { input_push_state=0; }

      break;

    case 2:

      if      (*data == 1) { input_push_state=3; } // nal->clear_skipped_bytes(); }

      else if (*data == 0) { } // *out++ = 0; }

      else { input_push_state=0; }

      break;

    case 3:

      *out++ = *data;

      input_push_state = 4;

      break;

    case 4:

      *out++ = *data;

      input_push_state = 5;

      break;

    case 5:

      if (*data==0) { input_push_state=6; }

      else { *out++ = *data; }

      break;

    case 6:

      if (*data==0) { input_push_state=7; }

      else {

        *out++ = 0;

        *out++ = *data;

        input_push_state=5;

      }

      break;

    case 7:

      if      (*data==0) { *out++ = 0; }

      else if (*data==3) {

        *out++ = 0; *out++ = 0; input_push_state=5;

        // remember which byte we removed

        nal->insert_skipped_byte((out - nal->data()) + nal->num_skipped_bytes());

      }

      else if (*data==1) {

#if DEBUG_INSERT_STREAM_ERRORS

        if ((rand()%100)<90 && nal_data.size>0) {

          int pos = rand()%nal_data.size;

          int bit = rand()%8;

          nal->nal_data.data[pos] ^= 1<<bit;

          //printf("inserted error...\n");

        }

#endif

        // enforce the maximum NAL size: drop an oversized NAL and resync

        if (!nal_size_within_limit(out - nal->data())) {

          free_NAL_unit(std::move(pending_input_NAL));

          input_push_state = 0;

          return DE265_ERROR_NAL_SIZE_EXCEEDS_SECURITY_LIMIT;

        }

        nal->set_size(out - nal->data());;

        // push this completed NAL onto the decoder queue (transfers ownership)

        push_to_NAL_queue(std::move(pending_input_NAL));

        // initialize new, empty NAL unit

        pending_input_NAL = alloc_NAL_unit(len+3);

        if (pending_input_NAL == nullptr) {

          return DE265_ERROR_OUT_OF_MEMORY;

        }

        pending_input_NAL->pts = pts;

        pending_input_NAL->user_data = user_data;

        nal = pending_input_NAL.get();

        out = nal->data();

        input_push_state=3;

        //nal->clear_skipped_bytes();

      }

      else {

        *out++ = 0;

        *out++ = 0;

        *out++ = *data;

        input_push_state=5;

      }

      break;

    }

    data++;

  }

  nal->set_size(out - nal->data());

  // Enforce the maximum NAL size on the still-incomplete pending NAL. This bounds

  // memory when a single NAL grows across many push_data() calls without ever

  // reaching a start code. The oversized pending NAL is dropped and the parser

  // resyncs at the next start code.

  if (!nal_size_within_limit(nal->size())) {

    free_NAL_unit(std::move(pending_input_NAL));

    input_push_state = 0;

    return DE265_ERROR_NAL_SIZE_EXCEEDS_SECURITY_LIMIT;

  }

  return DE265_OK;

}

de265_error NAL_Parser::push_NAL(const unsigned char* data, int len,

                                 de265_PTS pts, void* user_data)

{

  // Cannot use byte-stream input and NAL input at the same time.

  assert(pending_input_NAL == nullptr);

  // A NAL unit must at least contain its two-byte header. Reject anything shorter

  // (including a negative length) before touching any state: a zero-length unit

  // would otherwise end up as a memcpy() with a NULL destination and, once queued,

  // would fail header parsing in decode_NAL() and abort decoding of the stream.

  if (len < 2) {

    return DE265_ERROR_INVALID_ARGUMENT;

  }

  end_of_frame = false;

  // enforce the maximum NAL size to bound memory usage

  if (!nal_size_within_limit(len)) {

    return DE265_ERROR_NAL_SIZE_EXCEEDS_SECURITY_LIMIT;

  }

  std::unique_ptr<NAL_unit> nal = alloc_NAL_unit(len);

  if (nal == nullptr || !nal->set_data(data, len)) {

    free_NAL_unit(std::move(nal));

    return DE265_ERROR_OUT_OF_MEMORY;

  }

  nal->pts = pts;

  nal->user_data = user_data;

  nal->remove_stuffing_bytes();

  push_to_NAL_queue(std::move(nal));

  return DE265_OK;

}

de265_error NAL_Parser::flush_data()

{

  if (pending_input_NAL) {

    NAL_unit* nal = pending_input_NAL.get();

    uint8_t null[2] = { 0,0 };

    // append bytes that are implied by the push state

    if (input_push_state==6) {

      if (!nal->append(null,1)) {

        return DE265_ERROR_OUT_OF_MEMORY;

      }

    }

    if (input_push_state==7) {

      if (!nal->append(null,2)) {

        return DE265_ERROR_OUT_OF_MEMORY;

      }

    }

    // only push the NAL if it contains at least the NAL header

    if (input_push_state>=5) {

      push_to_NAL_queue(std::move(pending_input_NAL));

    }

    input_push_state = 0;

  }

  return DE265_OK;

}

void NAL_Parser::remove_pending_input_data()

{

  // --- remove pending input data ---

  if (pending_input_NAL) {

    free_NAL_unit(std::move(pending_input_NAL));

  }

  for (;;) {

    std::unique_ptr<NAL_unit> nal = pop_from_NAL_queue();

    if (nal) { free_NAL_unit(std::move(nal)); }

    else break;

  }

  input_push_state = 0;

  nBytes_in_NAL_queue = 0;

}