solve edge case

Infra/Buffers/PooledMemoryStream.cs

@@ -0,0 +1,102 @@
+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);
+    }
+}