named-binary-tag/Nbt/Serialization/NbtWriter.cs

using System.Text;
using Nbt.Tag;

namespace Nbt.Serialization;

public interface INbtWriter : IDisposable
{
    void WriteNamedTag(NamedTag namedTag);

    void WriteTag(INbtTag tag);

    void WriteTagType(NbtTagType tagType);

    void WriteSByte(sbyte b);

    void WriteShort(short s);

    void WriteInt(int i);

    void WriteLong(long l);

    void WriteFloat(float f);

    void WriteDouble(double d);

    void WriteString(string s);
}

public interface INbtAsyncWriter : IAsyncDisposable
{
    ValueTask WriteNamedTagAsync(NamedTag namedTag, CancellationToken ct = default);

    ValueTask WriteTagAsync(INbtTag tag, CancellationToken ct = default);

    ValueTask WriteTagTypeAsync(NbtTagType tagType, CancellationToken ct = default);

    ValueTask WriteSByteAsync(sbyte b, CancellationToken ct = default);

    ValueTask WriteShortAsync(short s, CancellationToken ct = default);

    ValueTask WriteIntAsync(int i, CancellationToken ct = default);

    ValueTask WriteLongAsync(long l, CancellationToken ct = default);

    ValueTask WriteFloatAsync(float f, CancellationToken ct = default);

    ValueTask WriteDoubleAsync(double d, CancellationToken ct = default);

    ValueTask WriteStringAsync(string s, CancellationToken ct = default);
}

public sealed class NbtWriter(Stream stream, bool leaveOpen = false) : INbtWriter, INbtAsyncWriter
{
    private readonly Stream stream = stream;
    private readonly bool leaveOpen = leaveOpen;

    public Stream BaseStream => stream;

    public static NbtWriter Create(Stream stream, Compression compression = Compression.None, bool leaveOpen = false) => compression switch
    {
        Compression.GZip => new(Utils.CreateGZipDeflater(stream, leaveOpen), false),
        Compression.ZLib => new(Utils.CreateZLibDeflater(stream, leaveOpen), false),
        Compression.LZ4 => new(Utils.CreateLZ4Deflater(stream, leaveOpen), false),
        Compression.Auto => throw new ArgumentException($"{nameof(Compression.Auto)} is invalid for a writer", nameof(compression)),
        _ => new(stream, leaveOpen)
    };

    #region Sync

    private void Flush() => stream.Flush();

    private void Write(byte[] buffer, int offset, int count) => stream.Write(buffer, offset, count);

    private void Write(ReadOnlySpan<byte> buffer) => stream.Write(buffer);

    private void WriteByte(byte b) => stream.WriteByte(b);

    private void WriteEndian(Span<byte> buffer)
    {
        if (BitConverter.IsLittleEndian)
        {
            for (var i = buffer.Length - 1; i >= 0; i--)
            {
                WriteByte(buffer[i]);
            }
        }
        else
        {
            Write(buffer);
        }
    }

    public void WriteSByte(sbyte b) => WriteByte((byte)b);

    public void WriteUShort(ushort s)
    {
        Span<byte> buffer = stackalloc byte[sizeof(ushort)];
        BitConverter.TryWriteBytes(buffer, s);
        WriteEndian(buffer);
    }

    public void WriteShort(short s)
    {
        Span<byte> buffer = stackalloc byte[sizeof(short)];
        BitConverter.TryWriteBytes(buffer, s);
        WriteEndian(buffer);
    }

    public void WriteUInt(uint i)
    {
        Span<byte> buffer = stackalloc byte[sizeof(uint)];
        BitConverter.TryWriteBytes(buffer, i);
        WriteEndian(buffer);
    }

    public void WriteInt(int i)
    {
        Span<byte> buffer = stackalloc byte[sizeof(int)];
        BitConverter.TryWriteBytes(buffer, i);
        WriteEndian(buffer);
    }

    public void WriteULong(ulong l)
    {
        Span<byte> buffer = stackalloc byte[sizeof(ulong)];
        BitConverter.TryWriteBytes(buffer, l);
        WriteEndian(buffer);
    }

    public void WriteLong(long l)
    {
        Span<byte> buffer = stackalloc byte[sizeof(long)];
        BitConverter.TryWriteBytes(buffer, l);
        WriteEndian(buffer);
    }

    public void WriteFloat(float f)
    {
        Span<byte> buffer = stackalloc byte[sizeof(float)];
        BitConverter.TryWriteBytes(buffer, f);
        WriteEndian(buffer);
    }

    public void WriteDouble(double d)
    {
        Span<byte> buffer = stackalloc byte[sizeof(double)];
        BitConverter.TryWriteBytes(buffer, d);
        WriteEndian(buffer);
    }

    public void WriteString(string s)
    {
        WriteUShort((ushort)s.Length);

        var buffer = Encoding.UTF8.GetBytes(s);
        Write(buffer);
    }

    public void WriteNamedTag(NamedTag namedTag) => WriteNamedTag(namedTag.Name, namedTag.Tag);

    public void WriteNamedTag(string tagName, INbtTag tag)
    {
        var tagType = NbtUtils.GetTagType(tag.Type);
        WriteTagType(tag.Type);
        WriteString(tagName);
        tagType.Write(this, tag);
    }

    public void WriteTag(INbtTag tag)
    {
        var tagType = NbtUtils.GetTagType(tag.Type);
        WriteTagType(tag.Type);
        tagType.Write(this, tag);
    }

    public void WriteTagType(NbtTagType tagType) => WriteByte((byte)tagType);

    #endregion

    #region Async

    private ValueTask FlushAsync(CancellationToken ct) => new(stream.FlushAsync(ct));

    private ValueTask WriteAsync(byte[] buffer, int offset, int count, CancellationToken ct) => new(stream.WriteAsync(buffer, offset, count, ct));

    private ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct) => stream.WriteAsync(buffer, ct);

    private ValueTask WriteByteAsync(byte b, CancellationToken ct) => new(stream.WriteAsync([b], 0, 1, ct));

    private async ValueTask WriteEndianAsync(Memory<byte> buffer, CancellationToken ct)
    {
        if (BitConverter.IsLittleEndian)
        {
            for (var i = buffer.Length - 1; i >= 0; i--)
            {
                await WriteAsync(buffer.Slice(i, 1), ct);
            }
        }
        else
        {
            await WriteAsync(buffer, ct);
        }
    }

    public ValueTask WriteSByteAsync(sbyte b, CancellationToken ct = default) => WriteByteAsync((byte)b, ct);

    public async ValueTask WriteUShortAsync(ushort s, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(ushort)];
        BitConverter.TryWriteBytes(buffer, s);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteShortAsync(short s, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(short)];
        BitConverter.TryWriteBytes(buffer, s);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteUIntAsync(uint i, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(uint)];
        BitConverter.TryWriteBytes(buffer, i);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteIntAsync(int i, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(int)];
        BitConverter.TryWriteBytes(buffer, i);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteULongAsync(ulong l, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(ulong)];
        BitConverter.TryWriteBytes(buffer, l);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteLongAsync(long l, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(long)];
        BitConverter.TryWriteBytes(buffer, l);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteFloatAsync(float f, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(float)];
        BitConverter.TryWriteBytes(buffer, f);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteDoubleAsync(double d, CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(double)];
        BitConverter.TryWriteBytes(buffer, d);
        await WriteEndianAsync(buffer, ct);
    }

    public async ValueTask WriteStringAsync(string s, CancellationToken ct = default)
    {
        await WriteUShortAsync((ushort)s.Length, ct);

        var buffer = Encoding.UTF8.GetBytes(s);
        await WriteAsync(buffer, ct);
    }

    public ValueTask WriteNamedTagAsync(NamedTag namedTag, CancellationToken ct = default) => WriteNamedTagAsync(namedTag.Name, namedTag.Tag, ct);

    public async ValueTask WriteNamedTagAsync(string tagName, INbtTag tag, CancellationToken ct = default)
    {
        var tagType = NbtUtils.GetTagType(tag.Type);
        await WriteTagTypeAsync(tag.Type, ct);
        await WriteStringAsync(tagName, ct);
        await tagType.WriteAsync(this, tag, ct);
    }

    public async ValueTask WriteTagAsync(INbtTag tag, CancellationToken ct = default)
    {
        var tagType = NbtUtils.GetTagType(tag.Type);
        await WriteTagTypeAsync(tag.Type, ct);
        await tagType.WriteAsync(this, tag, ct);
    }

    public ValueTask WriteTagTypeAsync(NbtTagType tagType, CancellationToken ct = default) => WriteByteAsync((byte)tagType, ct);

    #endregion

    #region IDisposable

    private bool _disposedValue;

    private void Dispose(bool disposing)
    {
        if (!_disposedValue)
        {
            if (disposing)
            {
                if (!leaveOpen)
                {
                    stream.Dispose();
                }
            }

            _disposedValue = true;
        }
    }

    public void Dispose()
    {
        Dispose(disposing: true);
        GC.SuppressFinalize(this);
    }

    private async ValueTask DisposeAsync(bool disposing)
    {
        if (!_disposedValue)
        {
            if (disposing)
            {
                if (!leaveOpen)
                {
                    await stream.DisposeAsync();
                }
            }

            _disposedValue = true;
        }
    }

    public async ValueTask DisposeAsync()
    {
        await DisposeAsync(disposing: true);
        GC.SuppressFinalize(this);
    }

    #endregion
}