bit_writer.h

Go to the documentation of this file.

#pragma once
#include "../utility/assert.h"
#include "../utility/crc.h"
#include "../utility/endian.h"
#include "../utility/meta.h"
 
#include "byte_buffer.h"
#include "serialize_traits.h"
#include "stream_traits.h"
 
#include <cstdint>
#include <cstring>
#include <memory>
#include <type_traits>
 
namespace bitstream
{
    template<typename Policy>
    class bit_writer
    {
    public:
        static constexpr bool writing = true;
        static constexpr bool reading = false;
 
        template<typename... Ts,
            typename = std::enable_if_t<std::is_constructible_v<Policy, Ts...>>>
        bit_writer(Ts&&... args)
            noexcept(std::is_nothrow_constructible_v<Policy, Ts...>) :
            m_Policy(std::forward<Ts>(args) ...),
            m_Scratch(0),
            m_ScratchBits(0),
            m_WordIndex(0) {}
 
        bit_writer(const bit_writer&) = delete;
        
        bit_writer(bit_writer&& other) noexcept :
            m_Policy(std::move(other.m_Policy)),
            m_Scratch(other.m_Scratch),
            m_ScratchBits(other.m_ScratchBits),
            m_WordIndex(other.m_WordIndex)
        {
            other.m_Scratch = 0;
            other.m_ScratchBits = 0;
            other.m_WordIndex = 0;
        }
 
        bit_writer& operator=(const bit_writer&) = delete;
 
        bit_writer& operator=(bit_writer&& rhs) noexcept
        {
            m_Policy = std::move(rhs.m_Policy);
            m_Scratch = rhs.m_Scratch;
            m_ScratchBits = rhs.m_ScratchBits;
            m_WordIndex = rhs.m_WordIndex;
 
            rhs.m_Scratch = 0;
            rhs.m_ScratchBits = 0;
            rhs.m_WordIndex = 0;
            
            return *this;
        }
 
        [[nodiscard]] uint8_t* get_buffer() const noexcept { return reinterpret_cast<uint8_t*>(m_Policy.get_buffer()); }
 
        [[nodiscard]] uint32_t get_num_bits_serialized() const noexcept { return m_Policy.get_num_bits_serialized(); }
 
        [[nodiscard]] uint32_t get_num_bytes_serialized() const noexcept { return get_num_bits_serialized() > 0U ? ((get_num_bits_serialized() - 1U) / 8U + 1U) : 0U; }
 
        [[nodiscard]] bool can_serialize_bits(uint32_t num_bits) const noexcept { return m_Policy.can_serialize_bits(num_bits); }
 
        [[nodiscard]] uint32_t get_remaining_bits() const noexcept { return get_total_bits() - get_num_bits_serialized(); }
 
        [[nodiscard]] uint32_t get_total_bits() const noexcept { return m_Policy.get_total_bits(); }
        
        uint32_t flush() noexcept
        {
            if (m_ScratchBits > 0U)
            {
                uint32_t* ptr = m_Policy.get_buffer() + m_WordIndex;
                uint32_t ptr_value = static_cast<uint32_t>(m_Scratch >> 32U);
                *ptr = utility::to_big_endian32(ptr_value);
 
                m_Scratch = 0U;
                m_ScratchBits = 0U;
                m_WordIndex++;
            }
 
            return get_num_bits_serialized();
        }
 
        [[nodiscard]] bool pad_to_size(uint32_t num_bytes) noexcept
        {
            uint32_t num_bits_written = get_num_bits_serialized();
 
            BS_ASSERT(num_bytes * 8U >= num_bits_written);
 
            BS_ASSERT(can_serialize_bits(num_bytes * 8U - num_bits_written));
 
            if (num_bits_written == 0)
            {
                BS_ASSERT(m_Policy.extend(num_bytes * 8U - num_bits_written));
 
                std::memset(m_Policy.get_buffer(), 0, num_bytes);
                
                m_Scratch = 0;
                m_ScratchBits = 0;
                m_WordIndex = num_bytes / 4;
                
                return true;
            }
            
            uint32_t remainder = (num_bytes * 8U - num_bits_written) % 32U;
            uint32_t zero = 0;
 
            // Align to byte
            if (remainder != 0U)
                BS_ASSERT(serialize_bits(zero, remainder));
 
            uint32_t offset = get_num_bits_serialized() / 32;
            uint32_t max = num_bytes / 4;
 
            // Serialize words
            for (uint32_t i = offset; i < max; i++)
                BS_ASSERT(serialize_bits(zero, 32));
 
            return true;
        }
 
        [[nodiscard]] bool pad(uint32_t num_bytes) noexcept
        {
            return pad_to_size(get_num_bytes_serialized() + num_bytes);
        }
 
        [[nodiscard]] bool align() noexcept
        {
            uint32_t remainder = m_ScratchBits % 8U;
            if (remainder != 0U)
            {
                uint32_t zero = 0U;
                bool status = serialize_bits(zero, 8U - remainder);
 
                BS_ASSERT(status && get_num_bits_serialized() % 8U == 0U);
            }
            return true;
        }
 
        [[nodiscard]] bool serialize_bits(uint32_t value, uint32_t num_bits) noexcept
        {
            BS_ASSERT(num_bits > 0U && num_bits <= 32U);
 
            BS_ASSERT(m_Policy.extend(num_bits));
 
            // This is actually slower
            // Possibly due to unlikely branching
            /*if (num_bits == 32U && m_ScratchBits == 0U)
            {
                uint32_t* ptr = m_Policy.get_buffer() + m_WordIndex;
 
                *ptr = utility::to_big_endian32(value);
 
                m_WordIndex++;
 
                return true;
            }*/
 
            uint32_t offset = 64U - num_bits - m_ScratchBits;
            uint64_t ls_value = static_cast<uint64_t>(value) << offset;
 
            m_Scratch |= ls_value;
            m_ScratchBits += num_bits;
 
            if (m_ScratchBits >= 32U)
            {
                uint32_t* ptr = m_Policy.get_buffer() + m_WordIndex;
                uint32_t ptr_value = static_cast<uint32_t>(m_Scratch >> 32U);
                *ptr = utility::to_big_endian32(ptr_value);
                m_Scratch <<= 32ULL;
                m_ScratchBits -= 32U;
                m_WordIndex++;
            }
 
            return true;
        }
 
        [[nodiscard]] bool serialize_bytes(const uint8_t* bytes, uint32_t num_bits) noexcept
        {
            BS_ASSERT(num_bits > 0U);
            
            BS_ASSERT(can_serialize_bits(num_bits));
            
            // Write the byte array as words
            const uint32_t* word_buffer = reinterpret_cast<const uint32_t*>(bytes);
            uint32_t num_words = num_bits / 32U;
            
            if (m_ScratchBits % 32U == 0U && num_words > 0U)
            {
                BS_ASSERT(m_Policy.extend(num_words * 32U));
 
                // If the written buffer is word-aligned, just memcpy it
                std::memcpy(m_Policy.get_buffer() + m_WordIndex, word_buffer, num_words * 4U);
                
                m_WordIndex += num_words;
            }
            else
            {
                // If the buffer is not word-aligned, serialize a word at a time
                for (uint32_t i = 0U; i < num_words; i++)
                {
                    // Casting a byte-array to an int is wrong on little-endian systems
                    // We have to swap the bytes around
                    uint32_t value = utility::to_big_endian32(word_buffer[i]);
                    BS_ASSERT(serialize_bits(value, 32U));
                }
            }
            
            // Early exit if the word-count matches
            if (num_bits % 32U == 0U)
                return true;
            
            uint32_t remaining_bits = num_bits - num_words * 32U;
            
            uint32_t num_bytes = (remaining_bits - 1U) / 8U + 1U;
            for (uint32_t i = 0U; i < num_bytes; i++)
            {
                uint32_t value = static_cast<uint32_t>(bytes[num_words * 4U + i]);
                BS_ASSERT(serialize_bits(value, (std::min)(remaining_bits - i * 8U, 8U)));
            }
 
            return true;
        }
 
        template<typename T>
        [[nodiscard]] bool serialize_into(bit_writer<T>& writer) const noexcept
        {
            uint8_t* buffer = reinterpret_cast<uint8_t*>(m_Policy.get_buffer());
            uint32_t num_bits = get_num_bits_serialized();
            uint32_t remainder_bits = num_bits % 8U;
 
            BS_ASSERT(writer.serialize_bytes(buffer, num_bits - remainder_bits));
 
            if (remainder_bits > 0U)
            {
                uint32_t byte_value = buffer[num_bits / 8U] >> (8U - remainder_bits);
                BS_ASSERT(writer.serialize_bits(byte_value, remainder_bits));
            }
 
            return true;
        }
 
        template<typename Trait, typename... Args, typename = utility::has_serialize_t<Trait, bit_writer, Args...>>
        [[nodiscard]] bool serialize(Args&&... args) noexcept(utility::is_serialize_noexcept_v<Trait, bit_writer, Args...>)
        {
            return serialize_traits<Trait>::serialize(*this, std::forward<Args>(args)...);
        }
 
        template<typename... Args, typename Trait, typename = utility::has_deduce_serialize_t<Trait, bit_writer, Args...>>
        [[nodiscard]] bool serialize(Trait&& arg, Args&&... args) noexcept(utility::is_deduce_serialize_noexcept_v<Trait, bit_writer, Args...>)
        {
            return serialize_traits<utility::deduce_trait_t<Trait, bit_writer, Args...>>::serialize(*this, std::forward<Trait>(arg), std::forward<Args>(args)...);
        }
 
    private:
        Policy m_Policy;
 
        uint64_t m_Scratch;
        int m_ScratchBits;
        size_t m_WordIndex;
    };
 
    using fixed_bit_writer = bit_writer<fixed_policy>;
 
    template<typename T>
    using growing_bit_writer = bit_writer<growing_policy<T>>;
}