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

using System.Buffers;
using System.Text;
using Nbt.Tag;

namespace Nbt.Serialization;

public interface INbtReader : IDisposable
{
    INbtTag ReadTag();

    NamedTag ReadNamedTag();

    NbtTagType ReadTagType();

    sbyte ReadSByte();

    short ReadShort();

    int ReadInt();

    long ReadLong();

    float ReadFloat();

    double ReadDouble();

    string ReadString();
}

public interface INbtAsyncReader : IAsyncDisposable
{
    ValueTask<INbtTag> ReadTagAsync(CancellationToken ct = default);

    ValueTask<NamedTag> ReadNamedTagAsync(CancellationToken ct = default);

    ValueTask<NbtTagType> ReadTagTypeAsync(CancellationToken ct = default);

    ValueTask<sbyte> ReadSByteAsync(CancellationToken ct = default);

    ValueTask<short> ReadShortAsync(CancellationToken ct = default);

    ValueTask<int> ReadIntAsync(CancellationToken ct = default);

    ValueTask<long> ReadLongAsync(CancellationToken ct = default);

    ValueTask<float> ReadFloatAsync(CancellationToken ct = default);

    ValueTask<double> ReadDoubleAsync(CancellationToken ct = default);

    ValueTask<string> ReadStringAsync(CancellationToken ct = default);
}

public sealed class NbtReader(Stream stream, bool leaveOpen = false) : INbtReader, INbtAsyncReader
{
    private readonly Stream stream = stream;
    private readonly bool leaveOpen = leaveOpen;

    public Stream BaseStream => stream;

    public static NbtReader Create(Stream stream, Compression compression = Compression.Auto, bool leaveOpen = false) => compression switch
    {
        Compression.Auto => CreateInflater(stream, leaveOpen),
        Compression.GZip => new(Utils.CreateGZipInflater(stream, leaveOpen), false),
        Compression.ZLib => new(Utils.CreateZLibInflater(stream, leaveOpen), false),
        Compression.LZ4 => new(Utils.CreateLZ4Deflater(stream, leaveOpen), false),
        _ => new(stream, leaveOpen)
    };

    private static NbtReader CreateInflater(Stream stream, bool leaveOpen)
    {
        if (!stream.CanSeek)
        {
            return new(stream, leaveOpen);
        }

        var cmf = stream.ReadByte();

        if (cmf < 0)
        {
            throw new EndOfStreamException();
        }

        stream.Seek(-1, SeekOrigin.Current);

        return cmf switch
        {
            0x1F => new(Utils.CreateGZipInflater(stream, leaveOpen), false),
            0x78 => new(Utils.CreateZLibInflater(stream, leaveOpen), false),
            _ => new(stream, leaveOpen)
        };
    }

    #region Sync

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

    private int Read(Span<byte> buffer) => stream.Read(buffer);

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

    private void ReadExactly(Span<byte> buffer) => stream.ReadExactly(buffer);

    private byte ReadByte()
    {
        var b = new byte[1];
        ReadExactly(b);
        return b[0];
    }

    private void ReadEndian(Span<byte> buffer)
    {
        ReadExactly(buffer);

        if (BitConverter.IsLittleEndian)
        {
            buffer.Reverse();
        }
    }

    public sbyte ReadSByte() => (sbyte)ReadByte();

    private ushort ReadUShort()
    {
        Span<byte> buffer = stackalloc byte[sizeof(ushort)];
        ReadEndian(buffer);
        return BitConverter.ToUInt16(buffer);
    }

    public short ReadShort()
    {
        Span<byte> buffer = stackalloc byte[sizeof(short)];
        ReadEndian(buffer);
        return BitConverter.ToInt16(buffer);
    }

    public int ReadInt()
    {
        Span<byte> buffer = stackalloc byte[sizeof(int)];
        ReadEndian(buffer);
        return BitConverter.ToInt32(buffer);
    }

    public long ReadLong()
    {
        Span<byte> buffer = stackalloc byte[sizeof(long)];
        ReadEndian(buffer);
        return BitConverter.ToInt64(buffer);
    }

    public float ReadFloat()
    {
        Span<byte> buffer = stackalloc byte[sizeof(float)];
        ReadEndian(buffer);
        return BitConverter.ToSingle(buffer);
    }

    public double ReadDouble()
    {
        Span<byte> buffer = stackalloc byte[sizeof(double)];
        ReadEndian(buffer);
        return BitConverter.ToDouble(buffer);
    }

    public string ReadString()
    {
        var len = ReadUShort();

        if (len == 0)
        {
            return string.Empty;
        }

        var buffer = ArrayPool<byte>.Shared.Rent(len);

        try
        {
            ReadExactly(buffer, 0, len);
            return Encoding.UTF8.GetString(buffer, 0, len);
        }
        finally
        {
            Array.Fill(buffer, (byte)0, 0, len);
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    public NbtTagType ReadTagType() => (NbtTagType)ReadByte();

    public INbtTag ReadTag()
    {
        var tagType = NbtUtils.GetTagType(ReadTagType());
        return tagType.Read(this);
    }

    public NamedTag ReadNamedTag()
    {
        var tagType = NbtUtils.GetTagType(ReadTagType());
        var tagName = ReadString();
        var tag = tagType.Read(this);
        return new(tagName, tag);
    }

    #endregion

    #region Async

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

    private ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct) => stream.ReadAsync(buffer, ct);

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

    private ValueTask ReadExactlyAsync(Memory<byte> buffer, CancellationToken ct) => stream.ReadExactlyAsync(buffer, ct);

    private async ValueTask<int> ReadByteAsync(CancellationToken ct)
    {
        var b = new byte[1];
        await ReadExactlyAsync(b, ct);
        return b[0];
    }

    private async ValueTask ReadEndianAsync(Memory<byte> buffer, CancellationToken ct)
    {
        await ReadExactlyAsync(buffer, ct);

        if (BitConverter.IsLittleEndian)
        {
            buffer.Span.Reverse();
        }
    }

    public async ValueTask<sbyte> ReadSByteAsync(CancellationToken ct = default) => (sbyte)await ReadByteAsync(ct);

    public async ValueTask<ushort> ReadUShortAsync(CancellationToken ct)
    {
        var buffer = new byte[sizeof(ushort)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToUInt16(buffer);
    }

    public async ValueTask<short> ReadShortAsync(CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(short)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToInt16(buffer);
    }

    public async ValueTask<int> ReadIntAsync(CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(int)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToInt32(buffer);
    }

    public async ValueTask<long> ReadLongAsync(CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(long)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToInt64(buffer);
    }

    public async ValueTask<float> ReadFloatAsync(CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(float)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToSingle(buffer);
    }

    public async ValueTask<double> ReadDoubleAsync(CancellationToken ct = default)
    {
        var buffer = new byte[sizeof(double)];
        await ReadEndianAsync(buffer, ct);
        return BitConverter.ToDouble(buffer);
    }

    public async ValueTask<string> ReadStringAsync(CancellationToken ct = default)
    {
        var len = await ReadUShortAsync(ct);

        if (len == 0)
        {
            return string.Empty;
        }

        var buffer = ArrayPool<byte>.Shared.Rent(len);

        try
        {
            await ReadExactlyAsync(buffer, 0, len, ct);
            return Encoding.UTF8.GetString(buffer, 0, len);
        }
        finally
        {
            Array.Fill(buffer, (byte)0, 0, len);
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    public async ValueTask<NbtTagType> ReadTagTypeAsync(CancellationToken ct = default) => (NbtTagType)await ReadByteAsync(ct);

    public async ValueTask<INbtTag> ReadTagAsync(CancellationToken ct = default)
    {
        var tagType = await ReadTagTypeAsync(ct);
        var type = NbtUtils.GetTagType(tagType);
        return await type.ReadAsync(this, ct);
    }

    public async ValueTask<NamedTag> ReadNamedTagAsync(CancellationToken ct = default)
    {
        var tagType = await ReadTagTypeAsync(ct);
        var type = NbtUtils.GetTagType(tagType);
        var tagName = await ReadStringAsync(ct);
        var tag = await type.ReadAsync(this, ct);
        return new(tagName, tag);
    }

    #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
}