/*
 * Copyright 2009, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include "BitBuffer.h"


// #pragma mark - BitReader


struct BitBuffer::BitReader {
        const uint8*    data;
        uint64                  bitSize;
        uint32                  bitOffset;

        BitReader(const uint8* data, uint64 bitSize, uint32 bitOffset)
                :
                data(data),
                bitSize(bitSize),
                bitOffset(bitOffset)
        {
        }

        uint8 ReadByte()
        {
                uint8 byte = *data;
                data++;
                bitSize -= 8;

                if (bitOffset == 0)
                        return byte;

                return (byte << bitOffset) | (*data >> (8 - bitOffset));
        }

        uint8 ReadBits(uint32 count)
        {
                uint8 byte = *data;
                bitSize -= count;
                bitOffset += count;

                if (bitOffset <= 8) {
                        if (bitOffset == 8) {
                                bitOffset = 0;
                                data++;
                                return byte & ((1 << count) - 1);
                        }

                        return (byte >> (8 - bitOffset)) & ((1 << count) - 1);
                }

                data++;
                bitOffset -= 8;
                return ((byte << bitOffset) | (*data >> (8 - bitOffset)))
                        & ((1 << count) - 1);
        }
};


// #pragma mark - BitBuffer


BitBuffer::BitBuffer()
        :
        fMissingBits(0)
{
}


BitBuffer::~BitBuffer()
{
}


bool
BitBuffer::AddBytes(const void* data, size_t size)
{
        if (size == 0)
                return true;

        if (fMissingBits == 0) {
                size_t oldSize = fBytes.Size();
                if (!fBytes.AddUninitialized(size))
                        return false;

                memcpy(fBytes.Elements() + oldSize, data, size);
                return true;
        }

        return AddBits(data, (uint64)size * 8, 0);
}


bool
BitBuffer::AddBits(const void* _data, uint64 bitSize, uint32 bitOffset)
{
        if (bitSize == 0)
                return true;

        const uint8* data = (const uint8*)_data + bitOffset / 8;
        bitOffset %= 8;

        BitReader reader(data, bitSize, bitOffset);

        // handle special case first: no more bits than missing
        size_t oldSize = fBytes.Size();
        if (fMissingBits > 0 && bitSize <= fMissingBits) {
                fMissingBits -= bitSize;
                uint8 bits = reader.ReadBits(bitSize) << fMissingBits;
                fBytes[oldSize - 1] |= bits;
                return true;
        }

        // resize the buffer
        if (!fBytes.AddUninitialized((reader.bitSize - fMissingBits + 7) / 8))
                return false;

        // fill in missing bits
        if (fMissingBits > 0) {
                fBytes[oldSize - 1] |= reader.ReadBits(fMissingBits);
                fMissingBits = 0;
        }

        // read full bytes as long as we can
        uint8* buffer = fBytes.Elements() + oldSize;
        while (reader.bitSize >= 8) {
                *buffer = reader.ReadByte();
                buffer++;
        }

        // If we have left-over bits, write a partial byte.
        if (reader.bitSize > 0) {
                fMissingBits = 8 - reader.bitSize;
                *buffer = reader.ReadBits(reader.bitSize) << fMissingBits;
        }

        return true;
}


bool
BitBuffer::AddZeroBits(uint64 bitSize)
{
        if (bitSize == 0)
                return true;

        // handle special case first: no more bits than missing
        size_t oldSize = fBytes.Size();
        if (fMissingBits > 0 && bitSize <= fMissingBits) {
                fMissingBits -= bitSize;
                return true;
        }

        // resize the buffer
        if (!fBytes.AddUninitialized((bitSize - fMissingBits + 7) / 8))
                return false;

        // fill in missing bits
        if (fMissingBits > 0) {
                bitSize -= fMissingBits;
                fMissingBits = 0;
        }

        // zero the remaining bytes, including a potentially partial last byte
        uint8* buffer = fBytes.Elements() + oldSize;
        memset(buffer, 0, (bitSize + 7) / 8);
        bitSize %= 8;

        if (bitSize > 0)
                fMissingBits = 8 - bitSize;

        return true;
}