bit_stream_reader.hpp

#pragma once
 
#include "fly/types/bit_stream/detail/bit_stream.hpp"
#include "fly/types/bit_stream/detail/concepts.hpp"
#include "fly/types/bit_stream/types.hpp"
#include "fly/types/numeric/endian.hpp"
 
#include <algorithm>
#include <bit>
#include <cstdint>
#include <istream>
 
namespace fly {
 
class BitStreamReader : public detail::BitStream
{
public:
    explicit BitStreamReader(std::istream &stream) noexcept;
 
    bool read_word(word_type &word);
 
    bool read_byte(byte_type &byte);
 
    template <detail::BitStreamInteger DataType>
    byte_type read_bits(DataType &bits, byte_type size);
 
    template <detail::BitStreamInteger DataType>
    byte_type peek_bits(DataType &bits, byte_type size);
 
    void discard_bits(byte_type size);
 
    bool fully_consumed() const;
 
    byte_type header() const;
 
private:
    void refill_buffer();
 
    template <detail::BitStreamInteger DataType>
    byte_type fill(DataType &buffer, byte_type bytes);
 
    std::istream &m_stream;
 
    byte_type m_header {0};
    byte_type m_remainder {0};
};
 
//==================================================================================================
template <detail::BitStreamInteger DataType>
byte_type BitStreamReader::read_bits(DataType &bits, byte_type size)
{
    if constexpr (detail::BitStreamBuffer<DataType>)
    {
        // See peek_bits for why buffer_type reads must be split.
        const byte_type size_high = size / 2;
        const byte_type size_low = size - size_high;
        std::uint32_t bits_high, bits_low;
 
        const byte_type bits_read_high = peek_bits(bits_high, size_high);
        discard_bits(bits_read_high);
 
        const byte_type bits_read_low = peek_bits(bits_low, size_low);
        discard_bits(bits_read_low);
 
        bits = (static_cast<DataType>(bits_high) << size_low) | bits_low;
        return bits_read_high + bits_read_low;
    }
    else
    {
        const byte_type bits_read = peek_bits(bits, size);
        discard_bits(bits_read);
 
        return bits_read;
    }
}
 
//==================================================================================================
template <detail::BitStreamInteger DataType>
byte_type BitStreamReader::peek_bits(DataType &bits, byte_type size)
{
    // Peek operations the site of the byte buffer are not supported because the byte buffer could
    // be in a state where it cannot be refilled.
    //
    // For example, consider a 64-bit byte buffer, and reading 6 bits and then 64 bits. After the
    // 6-bit read, there will be 58 bits left in the buffer; not enough for the next 64 bit read.
    // The buffer must then be refilled, but it cannot because there is less than 1 byte free in the
    // byte buffer.
    //
    // Ideally, given the split peek operations below, the given bits could be filled with the 58
    // available, then the byte buffer entirely refilled. But the caller then cannot discard more
    // than 6 bits, which invalidates the whole peek_bits/discard_bits semantic. BitStreamReader
    // would need to support putting bits back onto the stream.
    static_assert(
        !detail::BitStreamBuffer<DataType>,
        "peek_bits only supports types smaller than buffer_type");
 
    byte_type peeked = 0, lshift = 0;
    bits = 0;
 
    // If there are more bits to peek than are available in the byte buffer, break the peek into two
    // peeks.
    if (size > m_position)
    {
        peeked = m_position;
 
        const DataType buffer = static_cast<DataType>(m_buffer);
        lshift = size - m_position;
 
        // Fill the input bits with the remainder of byte buffer.
        bits = buffer & bit_mask<DataType>(m_position);
        bits <<= lshift - 1;
        bits <<= 1;
 
        // Refill the buffer from the stream.
        refill_buffer();
 
        // Then update the input to only peek any remaining bits next.
        size = std::min(lshift, static_cast<byte_type>(m_position - peeked));
        lshift -= size;
    }
 
    const byte_type rshift = m_position - peeked - size;
    const DataType buffer = static_cast<DataType>(m_buffer >> rshift);
 
    bits |= (buffer & bit_mask<DataType>(size)) << lshift;
    peeked += size;
 
    return peeked;
}
 
//==================================================================================================
inline void BitStreamReader::discard_bits(byte_type size)
{
    m_position -= size;
}
 
//==================================================================================================
template <detail::BitStreamInteger DataType>
byte_type BitStreamReader::fill(DataType &buffer, byte_type bytes)
{
    if (m_stream)
    {
        const std::streamsize bytes_read = m_stream_buffer->sgetn(
            reinterpret_cast<std::ios::char_type *>(&buffer),
            static_cast<std::streamsize>(bytes));
 
        buffer = endian_swap_if_non_native<std::endian::big>(buffer);
        return static_cast<byte_type>(bytes_read);
    }
 
    return 0;
}
 
} // namespace fly