StitchATon2/Infra/Encoders/PngStreamEncoder.cs

using System.Buffers.Binary;
using System.IO.Compression;
using StitchATon2.Infra.Buffers;

namespace StitchATon2.Infra.Encoders;

public class PngStreamEncoder : IDisposable, IAsyncDisposable
{
    private const int BufferSize = 8 * 1024;
    private const int FlushThreshold = 1024;
    
    private readonly Stream _stream;
    private readonly MemoryStream _memoryStream;
    private readonly int _width;
    private readonly int _height;

    private readonly ZLibStream _zlibStream;
    private bool _disposed;
    private bool _shouldFlush;
    
    public PngStreamEncoder(Stream writableStream, int width, int height)
    {
        _stream = writableStream;
        _width = width;
        _height = height;
        _memoryStream = new MemoryStream(BufferSize * 2);
        _zlibStream = new ZLibStream(_memoryStream, CompressionLevel.Optimal, leaveOpen: true);
        _memoryStream.SetLength(8);
        _memoryStream.Position = 8;
    }

    ~PngStreamEncoder() => Dispose();

    public async Task WriteHeader(CancellationToken cancellationToken = default)
    {
        byte[] headerBytes = [
            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG Signature
            0x00, 0x00, 0x00, 0x0D, // Length
            
            // IHDR chunk
            0x49, 0x48, 0x44, 0x52, // IHDR
            0x00, 0x00, 0x00, 0x00, // Reserve to write Width
            0x00, 0x00, 0x00, 0x00, // Reserve to write Height
            0x08, // Bit depth
            0x02, // Color type
            0x00, // Compression method
            0x00, // Filter method
            0x00, // Interlace method
            0x00, 0x00, 0x00, 0x00, // Reserve to write CRC-32
        ];
        
        BinaryPrimitives.WriteInt32BigEndian(headerBytes.AsSpan(16), _width);
        BinaryPrimitives.WriteInt32BigEndian(headerBytes.AsSpan(20), _height);
        var crc = Crc32.Compute(headerBytes.AsSpan(12, 17));
        
        BinaryPrimitives.WriteUInt32BigEndian(headerBytes.AsSpan(29), crc);
        
        await _stream.WriteAsync(headerBytes, cancellationToken);
    }

    public async Task WriteDataAsync(IBuffer<byte> buffer, bool disposeBuffer = true, CancellationToken cancellationToken = default)
    {
        try
        {
            _zlibStream.Write([0]);

            var dataSlice = buffer.Memory;
            while (dataSlice.Length > FlushThreshold)
            {
                await _zlibStream.WriteAsync(dataSlice[..FlushThreshold], cancellationToken);
                await _zlibStream.FlushAsync(cancellationToken);
                dataSlice = dataSlice[FlushThreshold..];
                if (_memoryStream.Length >= BufferSize)
                    await FlushAsync(cancellationToken);
            }

            if (dataSlice.Length > 0)
            {
                await _zlibStream.WriteAsync(dataSlice, cancellationToken);
                await _zlibStream.FlushAsync(cancellationToken);
                _shouldFlush = true;
            }
        }
        finally
        {
            if(disposeBuffer)
                buffer.Dispose();
        }
    }

    private async Task FlushAsync(CancellationToken cancellationToken)
    {
        await _zlibStream.FlushAsync(cancellationToken);
        var dataSize = (int)(_memoryStream.Length - 8);

        _memoryStream.Write("\0\0\0\0"u8);

        _memoryStream.Position = 4;
        _memoryStream.Write("IDAT"u8);

        var buffer = _memoryStream.GetBuffer();
        BinaryPrimitives.WriteInt32BigEndian(buffer, dataSize);

        // write Crc
        var crc = Crc32.Compute(buffer.AsSpan(4, dataSize + 4));
        BinaryPrimitives.WriteUInt32BigEndian(buffer.AsSpan(dataSize + 8), crc);
            
        await _stream.WriteAsync(buffer.AsMemory(0, dataSize + 12), cancellationToken);
        _memoryStream.SetLength(8);
        _memoryStream.Position = 8;
        _shouldFlush = false;
    }

    public async ValueTask WriteEndOfFileAsync(CancellationToken cancellationToken = default)
    {
        if(_shouldFlush)
            await FlushAsync(cancellationToken);
        
        var endChunk = new byte[] {
            0x00, 0x00, 0x00, 0x00, // Length
            0x49, 0x45, 0x4E, 0x44, // IEND
            0xAE, 0x42, 0x60, 0x82, // Crc
        };
        
        await _stream.WriteAsync(endChunk, cancellationToken);
        await DisposeAsync();
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _zlibStream.Dispose();
            _memoryStream.Dispose();
            _disposed = true;
        }
        
        GC.SuppressFinalize(this);
    }

    public async ValueTask DisposeAsync()
    {
        if (!_disposed)
        {
            await _zlibStream.DisposeAsync();
            await _memoryStream.DisposeAsync();
            _disposed = true;
        }
        
        GC.SuppressFinalize(this);
    }
}