/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Derived from Stagefright's ABitReader.

#include "BitReader.h"

namespace mozilla
{

BitReader::BitReader(const mozilla::MediaByteBuffer* aBuffer)
  : BitReader(aBuffer->Elements(), aBuffer->Length() * 8)
{
}

BitReader::BitReader(const mozilla::MediaByteBuffer* aBuffer, size_t aBits)
  : BitReader(aBuffer->Elements(), aBits)
{
}

BitReader::BitReader(const uint8_t* aBuffer, size_t aBits)
  : mData(aBuffer)
  , mOriginalBitSize(aBits)
  , mTotalBitsLeft(aBits)
  , mSize((aBits + 7) / 8)
  , mReservoir(0)
  , mNumBitsLeft(0)
{
}

BitReader::~BitReader() { }

uint32_t
BitReader::ReadBits(size_t aNum)
{
  MOZ_ASSERT(aNum <= 32);
  if (mTotalBitsLeft < aNum) {
    NS_ASSERTION(false, "Reading past end of buffer");
    return 0;
  }
  uint32_t result = 0;
  while (aNum > 0) {
    if (mNumBitsLeft == 0) {
      FillReservoir();
    }

    size_t m = aNum;
    if (m > mNumBitsLeft) {
      m = mNumBitsLeft;
    }

    result = (result << m) | (mReservoir >> (32 - m));
    mReservoir <<= m;
    mNumBitsLeft -= m;
    mTotalBitsLeft -= m;

    aNum -= m;
  }

  return result;
}

// Read unsigned integer Exp-Golomb-coded.
uint32_t
BitReader::ReadUE()
{
  uint32_t i = 0;

  while (ReadBit() == 0 && i < 32) {
    i++;
  }
  if (i == 32) {
    // This can happen if the data is invalid, or if it's
    // short, since ReadBit() will return 0 when it runs
    // off the end of the buffer.
    NS_WARNING("Invalid H.264 data");
    return 0;
  }
  uint32_t r = ReadBits(i);
  r += (1 << i) - 1;
  return r;
}

// Read signed integer Exp-Golomb-coded.
int32_t
BitReader::ReadSE()
{
  int32_t r = ReadUE();
  if (r & 1) {
    return (r+1) / 2;
  } else {
    return -r / 2;
  }
}

uint64_t
BitReader::ReadU64()
{
  uint64_t hi = ReadU32();
  uint32_t lo = ReadU32();
  return (hi << 32) | lo;
}

uint64_t
BitReader::ReadUTF8()
{
  int64_t val = ReadBits(8);
  uint32_t top = (val & 0x80) >> 1;

  if ((val & 0xc0) == 0x80 || val >= 0xFE) {
    // error.
    return -1;
  }
  while (val & top) {
    int tmp = ReadBits(8) - 128;
    if (tmp >> 6) {
      // error.
      return -1;
    }
    val = (val << 6) + tmp;
    top <<= 5;
  }
  val &= (top << 1) - 1;
  return val;
}

size_t
BitReader::BitCount() const
{
  return mOriginalBitSize - mTotalBitsLeft;
}

size_t
BitReader::BitsLeft() const
{
  return mTotalBitsLeft;
}

void
BitReader::FillReservoir()
{
  if (mSize == 0) {
    NS_ASSERTION(false, "Attempting to fill reservoir from past end of data");
    return;
  }

  mReservoir = 0;
  size_t i;
  for (i = 0; mSize > 0 && i < 4; i++) {
    mReservoir = (mReservoir << 8) | *mData;
    mData++;
    mSize--;
  }

  mNumBitsLeft = 8 * i;
  mReservoir <<= 32 - mNumBitsLeft;
}

/* static */ uint32_t
BitReader::GetBitLength(const mozilla::MediaByteBuffer* aNAL)
{
  size_t size = aNAL->Length();

  while (size > 0 && aNAL->ElementAt(size - 1) == 0) {
    size--;
  }

  if (!size) {
    return 0;
  }

  if (size > UINT32_MAX / 8) {
    // We can't represent it, we'll use as much as we can.
    return UINT32_MAX;
  }

  uint8_t v = aNAL->ElementAt(size - 1);
  size *= 8;

  // Remove the stop bit and following trailing zeros.
  if (v) {
    // Count the consecutive zero bits (trailing) on the right by binary search.
    // Adapted from Matt Whitlock algorithm to only bother with 8 bits integers.
    uint32_t c;
    if (v & 1) {
      // Special case for odd v (assumed to happen half of the time).
      c = 0;
    } else {
      c = 1;
      if ((v & 0xf) == 0) {
        v >>= 4;
        c += 4;
      }
      if ((v & 0x3) == 0) {
        v >>= 2;
        c += 2;
      }
      c -= v & 0x1;
    }
    size -= c + 1;
  }
  return size;
}

} // namespace mozilla