BitStream.h

//
// Created by Lectem on 12/11/2016.
//

#pragma once

#include <stdint.h>
#include <SystemUtils.h>
#include <algorithm>
#include <assert.h>
#include <climits>
#include <type_traits>
#include "IOBase.h"
namespace WorldStone
{

class BitStreamView : public IOBase
{
    using byte                     = unsigned char;
    size_t      size               = 0;       
    size_t      firstBitOffset     = 0;       
    size_t      currentBitPosition = 0;       
    const byte* buffer             = nullptr; 

public:
    BitStreamView() = default;
    BitStreamView(const void* inputBuffer, size_t sizeInBits, size_t firstBitOffsetInBuffer = 0)
        : size(sizeInBits),
          firstBitOffset(firstBitOffsetInBuffer),
          currentBitPosition(firstBitOffsetInBuffer),
          buffer(static_cast<const byte*>(inputBuffer))
    {
        assert(firstBitOffset + size <= bufferSizeInBits());
    }

    BitStreamView createSubView(size_t newbufferSizeInBits) const
    {
        // 0Bits is a special case as it should always have a size of 0
        if (newbufferSizeInBits == 0) return {};
        const size_t curBytesPos     = currentBitPosition / CHAR_BIT;
        const size_t bitPosInCurByte = currentBitPosition % CHAR_BIT;
        assert(tell() + newbufferSizeInBits <= size);
        assert(currentBitPosition + newbufferSizeInBits <= bufferSizeInBits());
        return {buffer + curBytesPos, newbufferSizeInBits, bitPosInCurByte};
    }

    size_t tell() const { return currentBitPosition - firstBitOffset; }
    void setPosition(size_t newPosition)
    {
        assert(newPosition >= 0_z && newPosition < size);
        currentBitPosition = newPosition + firstBitOffset;
    }
    size_t bitPositionInBuffer() const { return currentBitPosition; }
    void skip(size_t nbBits)
    {
        assert(currentBitPosition + nbBits < bufferSizeInBits());
        currentBitPosition += nbBits;
    }

    void alignToByte()
    {
        // Remove the last 3 bytes to be a multiple of 8, but add 7 before so that we round up
        currentBitPosition = (currentBitPosition + 7_z) & (~0x7_z);
    }

    size_t bufferSizeInBytes() const { return (size + firstBitOffset + 7) / CHAR_BIT; }
    size_t bufferSizeInBits() const { return bufferSizeInBytes() * CHAR_BIT; }
    size_t sizeInBits() const { return size; }

    bool readBool()
    {
        const size_t currentBytesPos     = currentBitPosition / CHAR_BIT;
        const size_t bitPosInCurrentByte = currentBitPosition % CHAR_BIT;
        const int    mask                = (1 << bitPosInCurrentByte);
        currentBitPosition++;
        return (buffer[currentBytesPos] & mask) == mask;
    }
    uint32_t readBit() { return uint32_t(readBool()); }

    template<typename RetType = uint32_t>
    RetType readUnsigned(unsigned nbBits)
    {
        static_assert(std::is_unsigned<RetType>::value, "You must return an unsigned type !");
        RetType value           = 0;
        size_t  curBytesPos     = currentBitPosition / CHAR_BIT;
        size_t  bitPosInCurByte = currentBitPosition % CHAR_BIT;
        currentBitPosition += nbBits;
        // The following variable is needed because we are reading a little endian input.
        size_t curDestBitPosition = 0;
        while (curDestBitPosition < nbBits)
        {
            // How many bits we can read in this byte ?
            const size_t bitsToReadInCurByte =
                std::min(CHAR_BIT - bitPosInCurByte, nbBits - curDestBitPosition);
            const RetType mask = RetType((1U << bitsToReadInCurByte) - 1U);
            // The bits we got from the current byte
            const RetType inBits = (RetType(buffer[curBytesPos++]) >> bitPosInCurByte) & mask;
            // Place them at the right position
            value |= inBits << curDestBitPosition;
            curDestBitPosition += bitsToReadInCurByte;
            // Next time we start at the beginning of the byte, perhaps we could optimize by
            // Treating the first byte as a special case ?
            bitPosInCurByte = 0;
        }
        return value;
    }

    uint8_t readUnsigned8OrLess(const int nbBits)
    {
        const size_t curBytesPos     = currentBitPosition / CHAR_BIT;
        const int    bitPosInCurByte = currentBitPosition % CHAR_BIT;
        currentBitPosition += size_t(nbBits);
        // Note: Could get rid of the condition by having one extra byte at the end of the stream
        const uint16_t shortFromBuffer =
            buffer[curBytesPos] |
            ((bitPosInCurByte + nbBits > 8) ? uint16_t(buffer[curBytesPos + 1] << CHAR_BIT) : 0);
        const unsigned mask  = 0xFF >> (CHAR_BIT - nbBits);
        const uint8_t  value = uint8_t((shortFromBuffer >> bitPosInCurByte) & mask);
        return value;
    }

    uint32_t read0Bits() { return 0u; }

    template<unsigned NbBits, typename std::enable_if<(NbBits > 0)>::type* = nullptr>
    int32_t           readSigned()
    {
        return Utils::signExtend<int32_t, NbBits>(readUnsigned(NbBits));
    }

    template<unsigned NbBits, typename std::enable_if<(NbBits == 0)>::type* = nullptr>
    int32_t           readSigned()
    {
        currentBitPosition += NbBits;
        return 0;
    }
};
}