StitchATon2/Infra/Buffers/PooledMemoryStream.cs

using System.Buffers;

namespace StitchATon2.Infra.Buffers;

public class PooledMemoryStream : Stream
{
    private byte[] _buffer;
    private int _length;
    private int _position;
    private readonly ArrayPool<byte> _pool;
    private bool _disposed;

    public PooledMemoryStream(int initialCapacity = 1024, ArrayPool<byte>? pool = null)
    {
        _pool = pool ?? ArrayPool<byte>.Shared;
        _buffer = _pool.Rent(initialCapacity);
    }

    public override bool CanRead => !_disposed;
    public override bool CanSeek => !_disposed;
    public override bool CanWrite => !_disposed;
    public override long Length => _length;
    public override long Position
    {
        get => _position;
        set
        {
            if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
            if (value < 0 || value > int.MaxValue) throw new ArgumentOutOfRangeException();
            _position = (int)value;
        }
    }

    public byte[] GetBuffer() => _buffer;
    public ArraySegment<byte> GetWrittenSegment() => new(_buffer, 0, _length);

    public override void Flush() { /* no-op */ }

    public override int Read(byte[] buffer, int offset, int count)
    {
        if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
        int available = _length - _position;
        int toRead = Math.Min(count, available);
        Buffer.BlockCopy(_buffer, _position, buffer, offset, toRead);
        _position += toRead;
        return toRead;
    }

    public override void Write(byte[] buffer, int offset, int count)
    {
        if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
        EnsureCapacity(_position + count);
        Buffer.BlockCopy(buffer, offset, _buffer, _position, count);
        _position += count;
        _length = Math.Max(_length, _position);
    }

    public override long Seek(long offset, SeekOrigin origin)
    {
        if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
        int newPos = origin switch
        {
            SeekOrigin.Begin => (int)offset,
            SeekOrigin.Current => _position + (int)offset,
            SeekOrigin.End => _length + (int)offset,
            _ => throw new ArgumentOutOfRangeException()
        };
        if (newPos < 0) throw new IOException("Negative position");
        _position = newPos;
        return _position;
    }

    public override void SetLength(long value)
    {
        if (_disposed) throw new ObjectDisposedException(nameof(PooledMemoryStream));
        if (value < 0 || value > int.MaxValue) throw new ArgumentOutOfRangeException();
        EnsureCapacity((int)value);
        _length = (int)value;
        if (_position > _length) _position = _length;
    }

    private void EnsureCapacity(int size)
    {
        if (size <= _buffer.Length) return;
        int newSize = Math.Max(size, _buffer.Length * 2);
        byte[] newBuffer = _pool.Rent(newSize);
        Buffer.BlockCopy(_buffer, 0, newBuffer, 0, _length);
        _pool.Return(_buffer, clearArray: true);
        _buffer = newBuffer;
    }

    protected override void Dispose(bool disposing)
    {
        if (!_disposed)
        {
            _pool.Return(_buffer, clearArray: true);
            _buffer = Array.Empty<byte>();
            _disposed = true;
        }
        base.Dispose(disposing);
    }
}