banan-os/userspace/libraries/LibDEFLATE/include/LibDEFLATE/BitStream.h

119 lines
2.3 KiB
C++

#pragma once
#include <BAN/Vector.h>
#include <BAN/ByteSpan.h>
namespace LibDEFLATE
{
class BitInputStream
{
public:
BitInputStream(BAN::ConstByteSpan data)
: m_data(data)
{ }
BAN::ErrorOr<uint16_t> peek_bits(size_t count)
{
ASSERT(count <= 16);
while (m_bit_buffer_len < count)
{
if (m_data.empty())
return BAN::Error::from_errno(ENOBUFS);
m_bit_buffer |= m_data[0] << m_bit_buffer_len;
m_bit_buffer_len += 8;
m_data = m_data.slice(1);
}
return m_bit_buffer & ((1 << count) - 1);
}
BAN::ErrorOr<uint16_t> take_bits(size_t count)
{
const uint16_t result = TRY(peek_bits(count));
m_bit_buffer >>= count;
m_bit_buffer_len -= count;
return result;
}
BAN::ErrorOr<void> take_byte_aligned(uint8_t* output, size_t bytes)
{
ASSERT(m_bit_buffer % 8 == 0);
while (m_bit_buffer_len && bytes)
{
*output++ = m_bit_buffer;
m_bit_buffer >>= 8;
m_bit_buffer_len -= 8;
bytes--;
}
if (bytes > m_data.size())
return BAN::Error::from_errno(EINVAL);
memcpy(output, m_data.data(), bytes);
m_data = m_data.slice(bytes);
return {};
}
void skip_to_byte_boundary()
{
const size_t bits_to_remove = m_bit_buffer_len % 8;
m_bit_buffer >>= bits_to_remove;
m_bit_buffer_len -= bits_to_remove;
}
private:
BAN::ConstByteSpan m_data;
uint32_t m_bit_buffer { 0 };
uint8_t m_bit_buffer_len { 0 };
};
class BitOutputStream
{
public:
BAN::ErrorOr<void> write_bits(uint16_t value, size_t count)
{
ASSERT(m_bit_buffer_len < 8);
ASSERT(count <= 16);
const uint16_t mask = (1 << count) - 1;
m_bit_buffer |= (value & mask) << m_bit_buffer_len;
m_bit_buffer_len += count;
while (m_bit_buffer_len >= 8)
{
TRY(m_data.push_back(m_bit_buffer));
m_bit_buffer >>= 8;
m_bit_buffer_len -= 8;
}
return {};
}
BAN::ErrorOr<void> pad_to_byte_boundary()
{
ASSERT(m_bit_buffer_len < 8);
if (m_bit_buffer_len == 0)
return {};
TRY(m_data.push_back(m_bit_buffer));
m_bit_buffer = 0;
m_bit_buffer_len = 0;
return {};
}
BAN::Vector<uint8_t> take_buffer()
{
ASSERT(m_bit_buffer_len == 0);
return BAN::move(m_data);
}
private:
BAN::Vector<uint8_t> m_data;
uint32_t m_bit_buffer { 0 };
uint8_t m_bit_buffer_len { 0 };
};
}