StitchATon2/Domain/ImageCreators/DangerousImageCreator.cs

using System.IO.Pipelines;
using System.Runtime.CompilerServices;
using SixLabors.ImageSharp.PixelFormats;
using StitchATon2.Infra;
using StitchATon2.Infra.Buffers;
using StitchATon2.Infra.Encoders;
using StitchATon2.Infra.Synchronization;

namespace StitchATon2.Domain.ImageCreators;

public sealed class DangerousImageCreator : IDisposable
{
    private readonly GridSection _section;
    private readonly IBuffer<Int32Pixel> _mmfReadBuffer;

    private int FullWidth => _section.TileManager.Configuration.FullWidth;
    private int FullHeight => _section.TileManager.Configuration.FullHeight;

    private int OffsetX => _section.OffsetX;
    private int OffsetY => _section.OffsetY;

    private int Width => _section.Width;
    private int Height => _section.Height;

    private int TileWidth => _section.TileManager.Configuration.Width;
    private int TileHeight => _section.TileManager.Configuration.Height;
    private Tile TileOrigin => _section.Origin;
    
    private int RightmostPixelIndex => _section.TileManager.Configuration.RightTileIndex;
    private int BottomPixelIndex => _section.TileManager.Configuration.BottomTileIndex;
    
    private TileManager TileManager => _section.TileManager;

    public DangerousImageCreator(GridSection section)
    {
        _section = section;
        _mmfReadBuffer = MemoryAllocator.Allocate<Int32Pixel>(TileWidth);
    }
    
    ~DangerousImageCreator() => Dispose();

    private Task Create(float scale, Action<IBuffer<byte>> writeRowCallback, CancellationToken cancellationToken)
    {
        var scaleFactor = MathF.ReciprocalEstimate(scale);
        var targetWidth = (int)(Width / scaleFactor);
        var targetHeight = (int)(Height / scaleFactor);
        if (targetHeight == 0 || targetWidth == 0)
            return Task.CompletedTask;

        using var xLookup = Utils.BoundsMatrix(scaleFactor, targetWidth, FullWidth, OffsetX);
        using var yLookup = Utils.BoundsMatrix(scaleFactor, targetHeight, FullHeight, OffsetY);

        using var yStartMap = MemoryAllocator.Allocate<Int32Pixel>(targetWidth + 1);
        using var yEndMap = MemoryAllocator.Allocate<Int32Pixel>(targetWidth + 1);
        
        var outputBufferSize = targetWidth * Unsafe.SizeOf<Rgb24>();
        
        // Use pixel referencing to eliminate type casting
        var pxInt32 = Int32Pixel.Zero;
        ref var px = ref pxInt32;
        ref var rChannel = ref Unsafe.As<Int32Pixel, byte>(ref px);
        ref var gChannel = ref Unsafe.Add(ref rChannel, 4);
        ref var bChannel = ref Unsafe.Add(ref rChannel, 8);

        var taskQueue = TaskHelper.SynchronizedTaskFactory.StartNew(() => { }, cancellationToken);
        
        for (var y = 1; y <= targetHeight; y++)
        {
            var yStart = yLookup[y - 1];
            var yEnd = yLookup[y];
            
            var (localRow0, localOffsetY0) = int.DivRem(yStart, TileHeight);
            MapRow(localRow0, localOffsetY0, xLookup, targetWidth + 1, yStartMap);
            
            var (localRow1, localOffsetY1) = int.DivRem(yEnd, TileHeight);
            MapRow(localRow1, localOffsetY1, xLookup, targetWidth + 1, yEndMap);
            
            // Cross row
            if (localRow0 != localRow1)
                MapRow(localRow0, BottomPixelIndex, xLookup, targetWidth + 1, yEndMap, true);
            
            var outputBuffer = MemoryAllocator.Allocate<byte>(outputBufferSize);
            ref var outputChannel = ref outputBuffer.Span[0];
            var boxHeight = yEnd - yStart;
            
            for (int x1 = 1, x0 = 0; x1 <= targetWidth; x0 = x1++)
            {
                var xStart = xLookup[x1 - 1];
                var xEnd = xLookup[x1];

                px = yEndMap[x1];
                px += yStartMap[x0];
                px -= yEndMap[x0];
                px -= yStartMap[x1];
                px /= Math.Max(1, (xEnd - xStart) * boxHeight);
                
                outputChannel = rChannel;
                outputChannel = ref Unsafe.AddByteOffset(ref outputChannel, 1);
                
                outputChannel = gChannel;
                outputChannel = ref Unsafe.AddByteOffset(ref outputChannel, 1);
                
                outputChannel = bChannel;
                outputChannel = ref Unsafe.AddByteOffset(ref outputChannel, 1);
            }

            if (cancellationToken.IsCancellationRequested)
            {
                outputBuffer.Dispose();
                return Task.FromCanceled(cancellationToken);
            }
            
            cancellationToken.Register(outputBuffer.Dispose);
            taskQueue = taskQueue
                .ContinueWith(
                    _ => writeRowCallback.Invoke(outputBuffer),
                    cancellationToken,
                    TaskContinuationOptions.OnlyOnRanToCompletion,
                    TaskScheduler.Current);
        }

        return taskQueue;
    }

    public async Task WriteToPipe(PipeWriter outputPipe, float scale, CancellationToken cancellationToken = default)
    {
        var scaleFactor = MathF.ReciprocalEstimate(scale);
        var targetWidth = (int)(Width / scaleFactor);
        var targetHeight = (int)(Height / scaleFactor);
        if (targetHeight == 0 || targetWidth == 0)
            return;
        
        var encoder = new PngPipeEncoder(outputPipe, targetWidth, targetHeight);
        encoder.WriteHeader();
        
        await Create(scale, 
            output => encoder
                .WriteDataAsync(output, cancellationToken: cancellationToken)
                .Wait(cancellationToken),
            cancellationToken);
        
        await encoder.WriteEndOfFileAsync(cancellationToken);
    }

    public async Task WriteToStream(Stream outputStream, float scale, CancellationToken cancellationToken = default)
    {
        var scaleFactor = MathF.ReciprocalEstimate(scale);
        var targetWidth = (int)(Width / scaleFactor);
        var targetHeight = (int)(Height / scaleFactor);
        if (targetHeight == 0 || targetWidth == 0)
            return;
        
        var encoder = new PngStreamEncoder(outputStream, targetWidth, targetHeight);
        await encoder.WriteHeader(cancellationToken);
        
        await Create(scale, 
            output => encoder
                .WriteDataAsync(output, cancellationToken: cancellationToken)
                .Wait(cancellationToken),
            cancellationToken);
        
        await encoder.WriteEndOfFileAsync(cancellationToken);
    }
    
    private void MapRow(
        int rowOffset,
        int yOffset,
        IBuffer<int> boundsMatrix,
        int count,
        IBuffer<Int32Pixel> destination,
        bool appendMode = false)
    {
        var currentTile = TileManager.TryGetAdjacent(TileOrigin, 0, rowOffset);
        if (currentTile == null)
        {
            if (appendMode) return;
            for (var i = 0; i < count; i++)
                destination[i] = Int32Pixel.Zero;

            return;
        }
        
        var sourceMap = boundsMatrix.Span[..count];
        var xAdder = Int32Pixel.Zero;
        var xOffset = 0;
        var written = 0;
        var destinationSpan = destination.Span;
        var readBufferSpan = _mmfReadBuffer.Span;
        
        while (true)
        {
            currentTile.Integral.Acquire(yOffset, readBufferSpan);
            int localX;

            if (appendMode)
            {
                while (written < sourceMap.Length && (localX = sourceMap[written] - xOffset) < TileWidth)
                {
                    destinationSpan[written] += readBufferSpan[localX];
                    destinationSpan[written] += xAdder;
                    written++;
                }
            }
            else
            {
                while (written < sourceMap.Length && (localX = sourceMap[written] - xOffset) < TileWidth)
                {
                    destinationSpan[written] = readBufferSpan[localX];
                    destinationSpan[written] += xAdder;
                    written++;
                }
            }
            
            if (written >= sourceMap.Length)
                break;
            
            currentTile = TileManager.TryGetAdjacent(currentTile, 1, 0);
            if (currentTile == null)
            {
                destinationSpan[written] = destinationSpan[written - 1];
                break;
            }

            xAdder += readBufferSpan[RightmostPixelIndex];
            xOffset += TileWidth;
        }
    }

    public void Dispose()
    {
        _mmfReadBuffer.Dispose();
        GC.SuppressFinalize(this);
    }
}