StitchATon2/Infra/Encoders/UnsafePngEncoder.cs

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

namespace StitchATon2.Infra.Encoders;

public class UnsafePngEncoder : IDisposable
{
    private const int BufferSize = 8 * 1024;
    private const int FlushThreshold = 1024;
    private const int PipeChunkThreshold = 16 * 1024;
    
    private readonly PipeWriter _outputPipe;
    private readonly int _width;
    private readonly int _height;

    private MemoryHandle? _memoryHandle;
    private readonly RawPointerStream _memoryStream;// = new RawPointerStream();
    private ZLibStream? _zlibStream;// = new ZLibStream(_memoryStream, CompressionLevel.Optimal, leaveOpen: true);
    private bool _disposed;
    private bool _shouldFlush;
    
    public UnsafePngEncoder(PipeWriter outputPipe, int width, int height)
    {
        _outputPipe = outputPipe;
        _width = width;
        _height = height;

        _memoryStream = new RawPointerStream();
    }

    ~UnsafePngEncoder() => Dispose();

    public void WriteHeader()
    {
        Span<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[16..], _width);
        BinaryPrimitives.WriteInt32BigEndian(headerBytes[20..], _height);
        var crc = Crc32.Compute(headerBytes.Slice(12, 17));
        
        BinaryPrimitives.WriteUInt32BigEndian(headerBytes[29..], crc);
        
        _outputPipe.Write(headerBytes);
    }

    private unsafe void Initialize()
    {
        if (_memoryHandle == null)
        {
            var memory = _outputPipe.GetMemory(PipeChunkThreshold);
            var handle = memory.Pin();
            _memoryStream.Initialize((byte*)handle.Pointer, 0, memory.Length);
            _memoryHandle = handle;
            
            _memoryStream.SetLength(8);
            _memoryStream.Position = 8;
            _zlibStream = new ZLibStream(_memoryStream, CompressionLevel.Optimal, true);
        }
    }

    public async Task WriteDataAsync(IBuffer<byte> buffer, bool disposeBuffer = true, CancellationToken cancellationToken = default)
    {
        Initialize();
        _zlibStream!.Write([0]);
        
        var offset = 0;
        while (buffer.Length - offset > FlushThreshold)
        {
            _zlibStream.Write(buffer.Span.Slice(offset, FlushThreshold));
            await _zlibStream.FlushAsync(cancellationToken);
            offset += FlushThreshold;
            if(_memoryStream.Length >= BufferSize)
                await FlushAsync(cancellationToken);
        }
        
        if (buffer.Length > offset)
        {
            _zlibStream.Write(buffer.Span[offset..]);
            await _zlibStream.FlushAsync(cancellationToken);
            _shouldFlush = true;
        }
        
        if(disposeBuffer) buffer.Dispose();
    }

    private async Task FlushAsync(CancellationToken cancellationToken)
    {
        await _zlibStream!.FlushAsync(cancellationToken);
        var dataSize = (int)(_memoryStream.Length - 8);
        
        _memoryStream.Position = 0;
        Span<byte> buffer = stackalloc byte[4];
        BinaryPrimitives.WriteInt32BigEndian(buffer, dataSize);
        _memoryStream.Write(buffer);
        _memoryStream.Write("IDAT"u8);

        _memoryStream.Position = 4;

        // write Crc
        var crc = Crc32.Compute(_memoryStream, dataSize + 4);
        BinaryPrimitives.WriteUInt32BigEndian(buffer, crc);
        _memoryStream.Write(buffer);

        _outputPipe.Advance((int)_memoryStream.Length);
        
        await _memoryStream.DisposeAsync();
        _memoryHandle!.Value.Dispose();
        _memoryHandle = null;
        
        _shouldFlush = false;
    }

    public async Task WriteEndOfFileAsync(CancellationToken cancellationToken = default)
    {
        if(_shouldFlush)
            await FlushAsync(cancellationToken);
        
        Span<byte> endChunk = [
            0x00, 0x00, 0x00, 0x00, // Length
            0x49, 0x45, 0x4E, 0x44, // IEND
            0xAE, 0x42, 0x60, 0x82, // Crc
        ];
        
        _outputPipe.Write(endChunk);
        Dispose();
    }

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

    private unsafe class RawPointerStream : UnmanagedMemoryStream
    {
        public void Initialize(byte* pointer, int length, int capacity)
        {
            Initialize(pointer, length, capacity, FileAccess.ReadWrite);
        }
    }
}