/*

 * 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>

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

NAL_unit::NAL_unit()

  : skipped_bytes(DE265_SKIPPED_BYTES_INITIAL_SIZE)

{

}

NAL_unit::~NAL_unit()

{

  free(nal_data);

}

void NAL_unit::clear()

{

  header = nal_header();

  pts = 0;

  user_data = nullptr;

  // set size to zero but keep memory

  data_size = 0;

  skipped_bytes.clear();

}

LIBDE265_CHECK_RESULT 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;

}

LIBDE265_CHECK_RESULT 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;

}

bool LIBDE265_CHECK_RESULT 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()

{

  // --- free NAL queues ---

  // empty NAL queue

  NAL_unit* nal;

  while ( (nal = pop_from_NAL_queue()) ) {

    free_NAL_unit(nal);

  }

  // free the pending input NAL

  if (pending_input_NAL != nullptr) {

    free_NAL_unit(pending_input_NAL);

  }

  // free all NALs in free-list

  for (size_t i=0;i<NAL_free_list.size();i++) {

    delete NAL_free_list[i];

  }

}

LIBDE265_CHECK_RESULT NAL_unit* NAL_Parser::alloc_NAL_unit(int size)

{

  NAL_unit* nal;

  // --- get NAL-unit object ---

  if (NAL_free_list.size() > 0) {

    nal = NAL_free_list.back();

    NAL_free_list.pop_back();

  }

  else {

    nal = new NAL_unit;

  }

  nal->clear();

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

    free_NAL_unit(nal);

    return nullptr;

  }

  return nal;

}

void NAL_Parser::free_NAL_unit(NAL_unit* nal)

{

  if (nal == nullptr) {

    // Allow calling with nullptr just like regular "free()"

    return;

  }

  if (NAL_free_list.size() < DE265_NAL_FREE_LIST_SIZE) {

    NAL_free_list.push_back(nal);

  }

  else {

    delete nal;

  }

}

NAL_unit* NAL_Parser::pop_from_NAL_queue()

{

  if (NAL_queue.empty()) {

    return nullptr;

  }

  else {

    NAL_unit* nal = NAL_queue.front();

    NAL_queue.pop();

    nBytes_in_NAL_queue -= nal->size();

    return nal;

  }

}

void NAL_Parser::push_to_NAL_queue(NAL_unit* nal)

{

  NAL_queue.push(nal);

  nBytes_in_NAL_queue += nal->size();

}

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;

  }

  NAL_unit* nal = pending_input_NAL; // 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(pending_input_NAL);

          pending_input_NAL = nullptr;

          input_push_state = 0;

          return DE265_ERROR_NAL_SIZE_EXCEEDS_SECURITY_LIMIT;

        }

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

        // push this NAL decoder queue

        push_to_NAL_queue(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;

        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(pending_input_NAL);

    pending_input_NAL = nullptr;

    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;

  }

  NAL_unit* nal = alloc_NAL_unit(len);

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

    free_NAL_unit(nal);

    return DE265_ERROR_OUT_OF_MEMORY;

  }

  nal->pts = pts;

  nal->user_data = user_data;

  nal->remove_stuffing_bytes();

  push_to_NAL_queue(nal);

  return DE265_OK;

}

de265_error NAL_Parser::flush_data()

{

  if (pending_input_NAL) {

    NAL_unit* nal = pending_input_NAL;

    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(nal);

      pending_input_NAL = nullptr;

    }

    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(pending_input_NAL);

    pending_input_NAL = nullptr;

  }

  for (;;) {

    NAL_unit* nal = pop_from_NAL_queue();

    if (nal) { free_NAL_unit(nal); }

    else break;

  }

  input_push_state = 0;

  nBytes_in_NAL_queue = 0;

}