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 sealed class NbtReader(Stream stream, bool leaveOpen = false) : INbtReader
{
    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)
    {
        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),
            _ => throw new UnknownCompressionSchemeException(cmf)
        };
    }

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

    private int ReadBytes(byte[] buffer) => ReadBytes(buffer, 0, buffer.Length);

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

    private void ReadAllBytes(byte[] buffer, int offset, int count)
    {
        if (count == 0)
        {
            return;
        }

        var bytesLeft = count;

        do
        {
            var bytesRead = ReadBytes(buffer, offset + count - bytesLeft, bytesLeft);

            if (bytesRead == 0)
            {
                throw new EndOfStreamException();
            }

            bytesLeft -= bytesRead;
        } while (bytesLeft > 0);
    }

    private void ReadAllBytes(byte[] buffer) => ReadAllBytes(buffer, 0, buffer.Length);

    private void ReadAllBytes(Span<byte> buffer)
    {
        var count = buffer.Length;

        if (count == 0)
        {
            return;
        }

        var bytesLeft = count;

        do
        {
            var bytesRead = ReadBytes(buffer.Slice(count - bytesLeft, bytesLeft));

            if (bytesRead == 0)
            {
                throw new EndOfStreamException();
            }

            bytesLeft -= bytesRead;
        } while (bytesLeft > 0);
    }

    private int Read() => stream.ReadByte();

    private int ReadAndThrowIfEndOfStream()
    {
        var b = Read();
        return b < 0 ? throw new EndOfStreamException() : b;
    }

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

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

    private byte ReadByte() => (byte)ReadAndThrowIfEndOfStream();

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

    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
        {
            ReadAllBytes(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);
    }

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

    #endregion
}