StitchATon2/Infra/ImageIntegral.cs

using System.Buffers;
using System.IO.MemoryMappedFiles;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using SixLabors.ImageSharp.Formats;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.PixelFormats;
using StitchATon2.Infra.Buffers;

namespace StitchATon2.Infra;

public class ImageIntegral : IDisposable
{
    private const int MaxProcessingQueue = 4;

    private readonly string _imagePath;
    private readonly string _outputDirectory;
    private readonly int _width;
    private readonly int _height;

    private IMemoryOwner<ManualResetEventSlim>? _rowLocks;
    private MemoryMappedFile? _memoryMappedFile;
    private readonly object _lock = new();

    private readonly ManualResetEventSlim _queueLock = new(true);
    private readonly ManualResetEventSlim _initializationLock = new(false);

    private volatile int _processedRows;
    private volatile int _queueCounter;

    public ImageIntegral(string imagePath, string outputDirectory, int width, int height)
    {
        _imagePath = imagePath;
        _outputDirectory = outputDirectory;
        _width = width;
        _height = height;
    }
    
    ~ImageIntegral() => Dispose();

    public void Acquire(int row, Int32Pixel[] buffer, CancellationToken cancellationToken = default)
    {
        Acquire(row, cancellationToken);
        ReadRow(row, buffer);
    }

    public void Acquire(int row, Span<Int32Pixel> buffer, CancellationToken cancellationToken = default)
    {
        Acquire(row, cancellationToken);
        ReadRow(row, buffer);
    }

    private void Acquire(int row, CancellationToken cancellationToken)
    {
        if (_memoryMappedFile is null)
        {
            lock (_lock)
            {
                if (_memoryMappedFile is null)
                {
                    Task.Factory.StartNew(() => Initialize(cancellationToken), cancellationToken);
                    _initializationLock.Wait(cancellationToken);
                    _initializationLock.Dispose();
                }
            }
        }

        _rowLocks?.Memory.Span[row].Wait(cancellationToken);
    }

    private void Initialize(CancellationToken cancellationToken)
    {
        var fileName = Path.GetFileNameWithoutExtension(_imagePath);
        var path = Path.Combine(_outputDirectory, $"{fileName}.mmf");

        var backedFileStream = InitializeBackedFile(path, out var header);
        _processedRows = header.ProcessedRows;

        if (header.ProcessedRows >= _height)
        {
            // When statement above is true
            // then it is guaranteed that backed file is valid and fully processed
            _memoryMappedFile = MemoryMappedFile.CreateFromFile(path, FileMode.Open);
            _initializationLock.Set();
            return;
        }
        
        var taskQueue = backedFileStream == null
            ? Task.CompletedTask
            : AllocateBackedFile(backedFileStream, header);

        taskQueue = taskQueue.ContinueWith(
            _ =>
            {
                _memoryMappedFile = MemoryMappedFile.CreateFromFile(path, FileMode.Open);
                _initializationLock.Set();
            },
            cancellationToken);

        // initialize resource gating, all rows is expected to be locked
        // if the backed file require to allocate, it should be safe to do this
        // asynchronously
        var rowLocks = MemoryAllocator.AllocateManaged<ManualResetEventSlim>(_height);
        var rowLocksSpan = rowLocks.Memory.Span;
        for (int i = 0; i < _height; i++)
        {
            var isOpen = i < header.ProcessedRows;
            rowLocksSpan[i] = new ManualResetEventSlim(isOpen);
        }

        _rowLocks = rowLocks;
        ProcessIntegral(taskQueue, cancellationToken);
    }

    private void ProcessIntegral(Task taskQueue, CancellationToken cancellationToken)
    {
        PngDecoderOptions decoderOptions = new()
        {
            PngCrcChunkHandling = PngCrcChunkHandling.IgnoreAll,
            GeneralOptions = new DecoderOptions
            {
                MaxFrames = 1,
                SkipMetadata = true,
            }
        };
        
        using var fileStream = File.OpenRead(_imagePath);
        using var image = PngDecoder.Instance.Decode<Rgb24>(decoderOptions, fileStream);
        var imageBuffer = image.Frames.RootFrame.PixelBuffer;
        
        var accumulator = Int32Pixel.Zero;
        var buffer = MemoryAllocator.AllocateArray<Int32Pixel>(_width);
        var processedRows = _processedRows;
        Interlocked.Exchange(ref _queueCounter, 0);
        
        // First row
        if (processedRows == 0)
        {
            var sourceRow = imageBuffer.DangerousGetRowSpan(0);
            for (var x = 0; x < sourceRow.Length; x++)
            {
                accumulator.Accumulate(sourceRow[x]);
                buffer[x] = accumulator;
            }
            
            taskQueue = QueueWriterTask(taskQueue, 0, buffer.Clone(_width), cancellationToken);
            processedRows++;
        }
        else
        {
            ReadRow(processedRows - 1, buffer);
        }
        
        if(cancellationToken.IsCancellationRequested)
            return;

        var prevBuffer = buffer;
        buffer = MemoryAllocator.AllocateArray<Int32Pixel>(_width);
        
        for (int y = processedRows; y < image.Height; y++)
        {
            var sourceRow = imageBuffer.DangerousGetRowSpan(y);
            accumulator = (Int32Pixel)sourceRow[0];
            buffer[0] = accumulator + prevBuffer[0];

            // Process all other columns
            for (var x = 1; x < sourceRow.Length; x++)
            {
                accumulator.Accumulate(sourceRow[x]);
                buffer[x] = accumulator + prevBuffer[x];
            }

            if (_queueCounter >= MaxProcessingQueue)
            {
                _queueLock.Reset();
                _queueLock.Wait(cancellationToken);
            }

            if(cancellationToken.IsCancellationRequested)
                break;

            var writeBuffer = prevBuffer;
            Array.Copy(buffer.Array, writeBuffer.Array, image.Width);
            taskQueue = QueueWriterTask(taskQueue, y, writeBuffer, cancellationToken);
            prevBuffer = buffer;
            buffer = MemoryAllocator.AllocateArray<Int32Pixel>(_width);
        }

        buffer.Dispose();
        if(cancellationToken.IsCancellationRequested)
            return;

        taskQueue = taskQueue.ContinueWith(task =>
        {
            if (task.IsCompletedSuccessfully)
            {
                DisposeRowLocks();
                _queueLock.Dispose();
            }
        }, cancellationToken);

        taskQueue.Wait(cancellationToken);
    }

    private Task QueueWriterTask(
        Task taskQueue,
        int row,
        ArrayOwner<Int32Pixel> writeBuffer,
        CancellationToken cancellationToken)
    {
        Interlocked.Increment(ref _queueCounter);
        cancellationToken.Register(writeBuffer.Dispose);
        return taskQueue.ContinueWith(_ =>
        {
            using (var view = AcquireView(row, MemoryMappedFileAccess.Write))
                view.WriteArray(0, writeBuffer.Array, 0, _width);
                
            writeBuffer.Dispose();
            _rowLocks!.Memory.Span[row].Set();
            Interlocked.Increment(ref _processedRows);

            using (var view = AcquireHeaderView(MemoryMappedFileAccess.Write))
                view.Write(16, _processedRows);
                
            Interlocked.Decrement(ref _queueCounter);
            _queueLock.Set();
        }, cancellationToken);
    }

    private MemoryMappedViewAccessor AcquireHeaderView(MemoryMappedFileAccess access)
        => _memoryMappedFile!.CreateViewAccessor(0, Header.Size, access);

    private MemoryMappedViewAccessor AcquireView(int row, MemoryMappedFileAccess access)
    {
        var size = _width * Int32Pixel.Size;
        var offset = row * size + Header.Size;
        return _memoryMappedFile!.CreateViewAccessor(offset, size, access);
    }

    private void ReadRow(int row, Int32Pixel[] readBuffer)
    {
        using var view = AcquireView(row, MemoryMappedFileAccess.Read);
        view.ReadArray(0, readBuffer, 0, _width);
    }

    private void ReadRow(int row, Span<Int32Pixel> buffer)
    {
        using var view = AcquireView(row, MemoryMappedFileAccess.Read);
        view.DangerousReadSpan(0, buffer, 0, _width);
    }

    private void ReadRow(int row, IBuffer<Int32Pixel> buffer)
    {
        using var view = AcquireView(row, MemoryMappedFileAccess.Read);
        view.DangerousReadSpan(0, buffer.Span, 0, _width);
    }

    private FileStream? InitializeBackedFile(string path, out Header header)
    {
        var expectedHeader = Header.CreateInitial(_width, _height);
        
        // Expectation when file exists:
        // - throws IOException when it's being processed, handle it if possible
        // - returns null if file is valid
        // - delete the existing file if it's not valid (modified from external)
        FileStream fs;
        if (File.Exists(path))
        {
            fs = File.OpenRead(path);
            if (fs.Length < expectedHeader.Length + Header.Size)
            {
                fs.Dispose();
                File.Delete(path);
            }
            
            Span<byte> headerBytes = stackalloc byte[Header.Size];
            fs.ReadExactly(headerBytes);
            
            header = MemoryMarshal.Cast<byte, Header>(headerBytes)[0];
            var isValid = expectedHeader.Identifier == header.Identifier
                          && expectedHeader.Width == header.Width
                          && expectedHeader.Height == header.Height
                          && expectedHeader.Length == header.Length;
            
            fs.Dispose();
            Console.WriteLine($"Image integral file found: {path}");
            
            if (!isValid) File.Delete(path);
            else return null;
        }
        
        // Expected process:
        // - Initialize file creation
        // - Write the header
        // - Allocate initial content by Header.Length
        // - Action above should be done asynchronously and return the task handle.
        var fsOptions = new FileStreamOptions
        {
            Access = FileAccess.Write,
            Share = FileShare.None,
            Mode = FileMode.CreateNew,
            PreallocationSize = Header.Size,
        };
        
        Console.WriteLine($"Create image integral file: {path}");
        
        Directory.CreateDirectory(_outputDirectory);
        fs = File.Open(path, fsOptions);
        fs.Write(MemoryMarshal.AsBytes([expectedHeader]));
        header = expectedHeader;
        return fs;
    }

    private static async Task AllocateBackedFile(FileStream fileStream, Header header)
    {
        // The input filestream is expected to be empty with
        // initial cursor at the beginning of the file and the content
        // is pre-allocated for at least Header.Length bytes
        // No other process should be accessed the file while being
        // allocated.
        // Allocated bytes is not necessary to be zeroed.

        // const int writeBufferSize = 4 * 1024;
        // using var writeBuffer = MemoryPool<byte>.Shared.Rent(writeBufferSize);
        //
        // var written = 0;
        // while (written + writeBufferSize < header.Length)
        // {
        //     await fileStream.WriteAsync(writeBuffer.Memory, cancellationToken);
        //     written += writeBufferSize;
        // }
        //
        // if (written < header.Length)
        // {
        //     await fileStream.WriteAsync(writeBuffer.Memory[..(header.Length - written)], cancellationToken);
        // }
        
        fileStream.SetLength(header.Length + Header.Size);
        
        await fileStream.DisposeAsync();
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct Header
    {
        private const int Signature = 0x47544e49; // INTG

        public static int Size => Unsafe.SizeOf<Header>();
        
        public uint Identifier;
        public int Width;
        public int Height;
        public int Length;
        public int ProcessedRows;

        public static Header CreateInitial(int width, int height) => new()
        {
            Identifier = Signature,
            Width = width,
            Height = height,
            Length = width * height * Int32Pixel.Size,
            ProcessedRows = 0,
        };
    }

    private void DisposeRowLocks()
    {
        if (_rowLocks is { } locks)
        {
            _rowLocks = null;
            var lockSpan = locks.Memory.Span;
            for(int i = 0; i < _height; i++)
                lockSpan[i].Dispose();

            locks.Dispose();
        }
    }

    public void Dispose()
    {
        DisposeRowLocks();
        _memoryMappedFile?.Dispose();
        _queueLock.Dispose();
        _initializationLock.Dispose();
        GC.SuppressFinalize(this);
    }
}