FeatherMC/src/NetworkMessage.h

#pragma once

#include <cstdint>
#include <cstring>
#include <vector>
#include <cassert>

#define VARINT_MAX_SIZE 5

namespace Feather
{
    class NetworkMessageCounter
    {
    public:

        inline void WriteVarInt(int32_t value)
        {
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                m_size++;
            } while (value != 0);
        }

        inline void WriteString(const char* string, int32_t length)
        {
            int32_t stringSize = length + 1;
            WriteVarInt(stringSize);

            m_size += stringSize;
        }

        inline void WriteLong(int64_t value)
        {
            m_size += sizeof(int64_t);
        }

        inline int32_t GetSize() const { return m_size; }

    private:
        int32_t m_size = 0;
    };

    class NetworkMessage
    {
    public:

        NetworkMessage(int32_t size)
        {
            NetworkMessageCounter packetSizeCounter;
            packetSizeCounter.WriteVarInt(size);

            size_t initialSize = size_t(size) + size_t(packetSizeCounter.GetSize());

            m_data.reserve(initialSize);
            WriteVarInt(size);
        }

        inline void WriteVarInt(int32_t value)
        {
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                m_data.push_back(temp);
            } while (value != 0);
        }

        inline void WriteString(const char* string, int32_t length)
        {
            int32_t stringSize = length + 1;
            WriteVarInt(stringSize);

            size_t oldSize = m_data.size();
            size_t newSize = oldSize + size_t(stringSize);
            m_data.resize(newSize);
            std::memcpy(&m_data[oldSize], string, size_t(stringSize));
        }

        inline void WriteLong(int64_t value)
        {
            size_t oldSize = m_data.size();
            m_data.resize(oldSize + sizeof(int64_t));
            // TODO: do this with registers if faster (see libcore.io.Memory.pokeLong)
            std::memcpy(&m_data[oldSize], &value, sizeof(int64_t));
        }

        inline virtual const uint8_t*   GetData()       const { return m_data.data(); }
        inline virtual size_t           GetDataSize()   const { return m_data.size(); }

    protected:
        // Used by subclasses with their own constructors
        NetworkMessage() {}

        std::vector<uint8_t> m_data;
    };

    #define DEFINE_MESSAGE_WRAPPER(x)                  \
        template <typename... Args>                    \
        NetworkMessage x(Args... args)                 \
        {                                              \
            NetworkMessageCounter counter;             \
            x(counter, args...);                       \
            NetworkMessage message(counter.GetSize()); \
            x(message, args...);                       \
            return message;                            \
        }

    // This class provides an alternative to NetworkMessageCounter, instead performing
    // one count and prepending the size marker to the start of the buffer.
    // 
    // To do this, you must call Finalize() before GetDataSize() or GetData()
    class DynamicNetworkMessage : public NetworkMessage
    {
    public:
        DynamicNetworkMessage(int32_t expectedSize) : NetworkMessage()
        {
            size_t initialSize = VARINT_MAX_SIZE + expectedSize;

            m_data.reserve(initialSize);
            m_data.resize(VARINT_MAX_SIZE);
        }

        void Finalize()
        {
            PrependVarIntSize();
            m_finalized = true;
        }

        // Inlining and 'final' should ideally avert vtable lookups
        // TOOD: Make these separate functions if vtable is not averted

        inline virtual const uint8_t* GetData() const final
        {
            // You must call Finalize() first
            assert(m_finalized);
            return m_data.data() + m_startOffset;
        }

        inline virtual size_t GetDataSize() const final
        {
            // You must call Finalize() first
            assert(m_finalized);
            return m_data.size() - m_startOffset;
        }
    private:
        int m_startOffset = 0;
        bool m_finalized = false;

        void PrependVarIntSize()
        {
            uint8_t sizeBytes[VARINT_MAX_SIZE];

            // compute the size of the packet excluding the size marker
            int value = int(m_data.size() - VARINT_MAX_SIZE);
            int index = 0;

            // perform WriteVarInt into the static buffer sizeBytes
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                sizeBytes[index++] = temp;
            } while (value != 0 && index < VARINT_MAX_SIZE);

            // where to start writing the size marker
            // if we used all 5 bytes, this is zero
            // if we used 3 bytes, this would be 2
            int startIndex = VARINT_MAX_SIZE - index;

            // Copy sizeBytes into our actual buffer, but align it to the right
            for (int i = 0; i < index; i++)
            {
                m_data[startIndex + i] = sizeBytes[i];
            }

            m_startOffset = startIndex;
        }
    };
}