acdream/src/AcDream.App/Rendering/Gpu/Vk/VulkanUploadQueue.cs

using Silk.NET.Vulkan;
using Buffer = Silk.NET.Vulkan.Buffer;

namespace AcDream.App.Rendering.Gpu.Vk;

/// <summary>
/// Campaign V slice V6a, plan §4.3 and §4.8: every transfer that has been asked
/// for but not yet recorded, plus the staging memory those transfers read from.
///
/// <para><b>Why a queue rather than an immediate copy.</b> Uploads arrive from
/// the streaming and texture-cache code at arbitrary moments, including with no
/// frame open at all. Vulkan copies must be recorded into a command buffer, and
/// they must be recorded OUTSIDE a dynamic-rendering block. So requests
/// accumulate here and are drained into the frame's command buffer at the one
/// moment both conditions hold: immediately before a pass begins. That is the
/// direct analogue of the GL backend's "flush immediately before every draw"
/// discipline, moved to the coarser granularity Vulkan actually needs.</para>
///
/// <para><b>One batched barrier, not one per copy.</b> The drain emits a single
/// buffer memory barrier covering every copy it recorded, moving the whole batch
/// from transfer writes to the vertex/index/indirect/shader reads that follow.
/// Plan §4.8 budgets four to six barriers per frame; this is one of them.</para>
///
/// <para><b>Staging exhaustion is not an error.</b> When the ring cannot serve a
/// request — the payload is larger than the whole ring, or unretired frames hold
/// the space — the upload takes a temporary dedicated staging buffer that is
/// retired through the ledger. The transfer is equally correct and equally
/// ordered; it just costs one allocation. See
/// <see cref="VulkanStagingRingState"/> for why that is a policy rather than a
/// workaround.</para>
/// </summary>
internal sealed unsafe class VulkanUploadQueue : IDisposable
{
    /// <summary>Plan §4.3: a 48 MiB persistently mapped staging ring.</summary>
    internal const ulong DefaultStagingCapacityBytes = 48UL * 1024 * 1024;

    private readonly Silk.NET.Vulkan.Vk _vk;
    private readonly Device _device;
    private readonly VulkanDeviceMemoryAllocator _allocator;
    private readonly VulkanFrameFlightController _flights;
    private readonly VulkanStagingRingState _ringState;
    private readonly VulkanDebugNames _debugNames;

    private readonly Buffer _stagingBuffer;
    private readonly VulkanAllocation _stagingAllocation;

    private readonly List<BufferCopy2> _bufferCopies = [];
    private readonly List<ImageCopy2> _imageCopies = [];
    private readonly List<TemporaryStaging> _temporaries = [];

    private bool _disposed;

    private readonly record struct BufferCopy2(Buffer Source, Buffer Destination, ulong SourceOffset, ulong DestinationOffset, ulong SizeBytes);

    private readonly record struct ImageCopy2(
        Buffer Source,
        ulong SourceOffset,
        Image Destination,
        uint MipLevel,
        uint Layer,
        uint Width,
        uint Height);

    private readonly record struct TemporaryStaging(Buffer Buffer, VulkanAllocation Allocation);

    internal VulkanUploadQueue(
        Silk.NET.Vulkan.Vk vk,
        Device device,
        VulkanDeviceMemoryAllocator allocator,
        VulkanFrameFlightController flights,
        VulkanDebugNames debugNames,
        ulong stagingCapacityBytes = DefaultStagingCapacityBytes)
    {
        _vk = vk ?? throw new ArgumentNullException(nameof(vk));
        _device = device;
        _allocator = allocator ?? throw new ArgumentNullException(nameof(allocator));
        _flights = flights ?? throw new ArgumentNullException(nameof(flights));
        _debugNames = debugNames ?? throw new ArgumentNullException(nameof(debugNames));
        _ringState = new VulkanStagingRingState(stagingCapacityBytes);

        (_stagingBuffer, _stagingAllocation) = CreateHostBuffer(
            stagingCapacityBytes,
            BufferUsageFlags.TransferSrcBit,
            "vk-staging-ring");
    }

    /// <summary>
    /// Images touched by this batch, and the layout each is in on entry.
    ///
    /// <para>The entry layout is not always <c>UNDEFINED</c>, and that
    /// distinction is load-bearing. <c>UNDEFINED</c> lets the driver discard the
    /// existing contents, which is exactly right for the first upload into a
    /// fresh image and exactly wrong for the incremental array-layer fills that
    /// mirror <c>ManagedGLTextureArray</c> — discarding there would erase every
    /// layer uploaded earlier. The first writer of a batch records the layout it
    /// found the image in, and that is what the barrier names.</para>
    /// </summary>
    private readonly Dictionary<Image, ImageLayout> _imageEntryLayouts = [];

    /// <summary>Mip chains to generate with <c>vkCmdBlitImage</c> after this batch's copies land.</summary>
    private readonly List<MipBlitRequest> _mipBlits = [];

    private readonly record struct MipBlitRequest(
        Image Image,
        int Width,
        int Height,
        int MipLevelCount,
        int LayerCount);

    internal int PendingBufferCopyCount => _bufferCopies.Count;

    internal int PendingImageCopyCount => _imageCopies.Count;

    internal ulong StagingLiveBytes => _ringState.LiveBytes;

    /// <summary>Stages <paramref name="data"/> and queues a copy into <paramref name="destination"/>.</summary>
    internal void StageBufferWrite(
        Buffer destination,
        ulong destinationOffsetBytes,
        ReadOnlySpan<byte> data,
        string ownerName)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (data.IsEmpty)
            return;

        (Buffer source, ulong sourceOffset) = Stage(data, ownerName);
        _bufferCopies.Add(new BufferCopy2(
            source,
            destination,
            sourceOffset,
            destinationOffsetBytes,
            (ulong)data.Length));
    }

    /// <summary>Stages <paramref name="data"/> and queues a copy into one mip level of one array layer.</summary>
    internal void StageImageWrite(
        Image destination,
        int mipLevel,
        int layer,
        int width,
        int height,
        ImageLayout entryLayout,
        ReadOnlySpan<byte> data,
        string ownerName)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (data.IsEmpty)
            return;

        // Vulkan requires a buffer-to-image copy's source offset to be a
        // multiple of 4 and of the texel/block size; 16 covers every format
        // acdream uses (BC3's 16-byte block is the largest).
        (Buffer source, ulong sourceOffset) = Stage(data, ownerName, alignmentBytes: 16);
        _imageCopies.Add(new ImageCopy2(
            source,
            sourceOffset,
            destination,
            (uint)mipLevel,
            (uint)layer,
            (uint)width,
            (uint)height));
        RecordEntryLayout(destination, entryLayout);
    }

    /// <summary>
    /// Queues a <c>vkCmdBlitImage</c> mip chain for an uncompressed image. BC
    /// images cannot use this — a compressed image is not a legal blit
    /// destination — and take a CPU-built chain instead (plan §4.3).
    /// </summary>
    internal void EnqueueMipBlit(
        Image image,
        int width,
        int height,
        int mipLevelCount,
        int layerCount,
        ImageLayout entryLayout)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (mipLevelCount <= 1)
            return;
        _mipBlits.Add(new MipBlitRequest(image, width, height, mipLevelCount, layerCount));
        RecordEntryLayout(image, entryLayout);
    }

    private void RecordEntryLayout(Image image, ImageLayout entryLayout)
    {
        // First writer of the batch wins: a later writer that found the image
        // already in TRANSFER_DST is describing this batch's own effect, not the
        // layout the batch started from.
        _imageEntryLayouts.TryAdd(image, entryLayout);
    }

    /// <summary>Queues a device-side buffer copy — the mesh arena's grow-and-copy migration.</summary>
    internal void EnqueueBufferCopy(
        Buffer source,
        Buffer destination,
        ulong sourceOffsetBytes,
        ulong destinationOffsetBytes,
        ulong byteCount)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (byteCount == 0)
            return;
        _bufferCopies.Add(new BufferCopy2(
            source,
            destination,
            sourceOffsetBytes,
            destinationOffsetBytes,
            byteCount));
    }

    /// <summary>
    /// Records every pending transfer into <paramref name="commands"/> and
    /// clears the queue. Must be called outside a dynamic-rendering block.
    /// Returns false when there was nothing to do.
    /// </summary>
    internal bool Record(CommandBuffer commands)
    {
        if (_bufferCopies.Count == 0 && _imageCopies.Count == 0 && _mipBlits.Count == 0)
            return false;

        if (_imageEntryLayouts.Count > 0)
            TransitionImagesToTransfer(commands);

        foreach (BufferCopy2 copy in _bufferCopies)
        {
            var region = new BufferCopy
            {
                SrcOffset = copy.SourceOffset,
                DstOffset = copy.DestinationOffset,
                Size = copy.SizeBytes,
            };
            _vk.CmdCopyBuffer(commands, copy.Source, copy.Destination, 1, &region);
        }

        foreach (ImageCopy2 copy in _imageCopies)
        {
            var region = new BufferImageCopy
            {
                BufferOffset = copy.SourceOffset,
                BufferRowLength = 0,
                BufferImageHeight = 0,
                ImageSubresource = new ImageSubresourceLayers
                {
                    AspectMask = ImageAspectFlags.ColorBit,
                    MipLevel = copy.MipLevel,
                    BaseArrayLayer = copy.Layer,
                    LayerCount = 1,
                },
                ImageOffset = new Offset3D(0, 0, 0),
                ImageExtent = new Extent3D(copy.Width, copy.Height, 1),
            };
            _vk.CmdCopyBufferToImage(
                commands,
                copy.Source,
                copy.Destination,
                ImageLayout.TransferDstOptimal,
                1,
                &region);
        }

        foreach (MipBlitRequest blit in _mipBlits)
            RecordMipBlit(commands, blit);

        if (_imageEntryLayouts.Count > 0)
            TransitionImagesToShaderRead(commands);

        // One buffer barrier for the whole batch: transfer writes become
        // readable by every consumer stage a copied buffer can feed.
        if (_bufferCopies.Count > 0)
        {
            var barrier = new MemoryBarrier2
            {
                SType = StructureType.MemoryBarrier2,
                SrcStageMask = PipelineStageFlags2.AllTransferBit,
                SrcAccessMask = AccessFlags2.TransferWriteBit,
                DstStageMask = PipelineStageFlags2.VertexInputBit
                    | PipelineStageFlags2.VertexShaderBit
                    | PipelineStageFlags2.FragmentShaderBit
                    | PipelineStageFlags2.DrawIndirectBit,
                DstAccessMask = AccessFlags2.VertexAttributeReadBit
                    | AccessFlags2.IndexReadBit
                    | AccessFlags2.ShaderReadBit
                    | AccessFlags2.UniformReadBit
                    | AccessFlags2.IndirectCommandReadBit,
            };
            var dependency = new DependencyInfo
            {
                SType = StructureType.DependencyInfo,
                MemoryBarrierCount = 1,
                PMemoryBarriers = &barrier,
            };
            _vk.CmdPipelineBarrier2(commands, &dependency);
        }

        _bufferCopies.Clear();
        _imageCopies.Clear();
        _mipBlits.Clear();
        _imageEntryLayouts.Clear();
        return true;
    }

    private static ImageSubresourceRange WholeColorImage => new()
    {
        AspectMask = ImageAspectFlags.ColorBit,
        BaseMipLevel = 0,
        LevelCount = Silk.NET.Vulkan.Vk.RemainingMipLevels,
        BaseArrayLayer = 0,
        LayerCount = Silk.NET.Vulkan.Vk.RemainingArrayLayers,
    };

    private void TransitionImagesToTransfer(CommandBuffer commands)
    {
        var barriers = new ImageMemoryBarrier2[_imageEntryLayouts.Count];
        int index = 0;
        foreach ((Image image, ImageLayout entryLayout) in _imageEntryLayouts)
        {
            barriers[index++] = new ImageMemoryBarrier2
            {
                SType = StructureType.ImageMemoryBarrier2,
                SrcStageMask = PipelineStageFlags2.AllCommandsBit,
                SrcAccessMask = AccessFlags2.None,
                DstStageMask = PipelineStageFlags2.AllTransferBit,
                DstAccessMask = AccessFlags2.TransferWriteBit | AccessFlags2.TransferReadBit,
                OldLayout = entryLayout,
                NewLayout = ImageLayout.TransferDstOptimal,
                SrcQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
                DstQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
                Image = image,
                SubresourceRange = WholeColorImage,
            };
        }

        SubmitBarriers(commands, barriers);
    }

    private void TransitionImagesToShaderRead(CommandBuffer commands)
    {
        var barriers = new ImageMemoryBarrier2[_imageEntryLayouts.Count];
        int index = 0;
        foreach (Image image in _imageEntryLayouts.Keys)
        {
            barriers[index++] = new ImageMemoryBarrier2
            {
                SType = StructureType.ImageMemoryBarrier2,
                SrcStageMask = PipelineStageFlags2.AllTransferBit,
                SrcAccessMask = AccessFlags2.TransferWriteBit,
                DstStageMask = PipelineStageFlags2.FragmentShaderBit | PipelineStageFlags2.VertexShaderBit,
                DstAccessMask = AccessFlags2.ShaderReadBit,
                OldLayout = ImageLayout.TransferDstOptimal,
                NewLayout = ImageLayout.ShaderReadOnlyOptimal,
                SrcQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
                DstQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
                Image = image,
                SubresourceRange = WholeColorImage,
            };
        }

        SubmitBarriers(commands, barriers);
    }

    private void SubmitBarriers(CommandBuffer commands, ImageMemoryBarrier2[] barriers)
    {
        if (barriers.Length == 0)
            return;
        fixed (ImageMemoryBarrier2* first = barriers)
        {
            var dependency = new DependencyInfo
            {
                SType = StructureType.DependencyInfo,
                ImageMemoryBarrierCount = (uint)barriers.Length,
                PImageMemoryBarriers = first,
            };
            _vk.CmdPipelineBarrier2(commands, &dependency);
        }
    }

    /// <summary>
    /// Halves level N into level N+1 with a linear blit, all layers at once.
    ///
    /// <para>Each source level is moved to TRANSFER_SRC for its blit and then
    /// moved BACK to TRANSFER_DST. Leaving the chain in mixed layouts would be
    /// one barrier cheaper and would then need the batch's final
    /// shader-read transition to name a different old layout per level; ending
    /// every level in the same layout is what lets that final transition stay
    /// one barrier per image.</para>
    /// </summary>
    private void RecordMipBlit(CommandBuffer commands, in MipBlitRequest request)
    {
        int width = request.Width;
        int height = request.Height;

        for (uint level = 1; level < request.MipLevelCount; level++)
        {
            int nextWidth = Math.Max(1, width / 2);
            int nextHeight = Math.Max(1, height / 2);

            TransitionMipLevel(
                commands,
                request.Image,
                level - 1,
                ImageLayout.TransferDstOptimal,
                ImageLayout.TransferSrcOptimal,
                AccessFlags2.TransferWriteBit,
                AccessFlags2.TransferReadBit);

            var blit = new ImageBlit2
            {
                SType = StructureType.ImageBlit2,
                SrcSubresource = new ImageSubresourceLayers
                {
                    AspectMask = ImageAspectFlags.ColorBit,
                    MipLevel = level - 1,
                    BaseArrayLayer = 0,
                    LayerCount = (uint)request.LayerCount,
                },
                DstSubresource = new ImageSubresourceLayers
                {
                    AspectMask = ImageAspectFlags.ColorBit,
                    MipLevel = level,
                    BaseArrayLayer = 0,
                    LayerCount = (uint)request.LayerCount,
                },
            };
            blit.SrcOffsets.Element0 = new Offset3D(0, 0, 0);
            blit.SrcOffsets.Element1 = new Offset3D(width, height, 1);
            blit.DstOffsets.Element0 = new Offset3D(0, 0, 0);
            blit.DstOffsets.Element1 = new Offset3D(nextWidth, nextHeight, 1);

            var info = new BlitImageInfo2
            {
                SType = StructureType.BlitImageInfo2,
                SrcImage = request.Image,
                SrcImageLayout = ImageLayout.TransferSrcOptimal,
                DstImage = request.Image,
                DstImageLayout = ImageLayout.TransferDstOptimal,
                RegionCount = 1,
                PRegions = &blit,
                Filter = Filter.Linear,
            };
            _vk.CmdBlitImage2(commands, &info);

            TransitionMipLevel(
                commands,
                request.Image,
                level - 1,
                ImageLayout.TransferSrcOptimal,
                ImageLayout.TransferDstOptimal,
                AccessFlags2.TransferReadBit,
                AccessFlags2.TransferWriteBit);

            width = nextWidth;
            height = nextHeight;
        }
    }

    private void TransitionMipLevel(
        CommandBuffer commands,
        Image image,
        uint level,
        ImageLayout oldLayout,
        ImageLayout newLayout,
        AccessFlags2 sourceAccess,
        AccessFlags2 destinationAccess)
    {
        var barrier = new ImageMemoryBarrier2
        {
            SType = StructureType.ImageMemoryBarrier2,
            SrcStageMask = PipelineStageFlags2.AllTransferBit,
            SrcAccessMask = sourceAccess,
            DstStageMask = PipelineStageFlags2.AllTransferBit,
            DstAccessMask = destinationAccess,
            OldLayout = oldLayout,
            NewLayout = newLayout,
            SrcQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
            DstQueueFamilyIndex = Silk.NET.Vulkan.Vk.QueueFamilyIgnored,
            Image = image,
            SubresourceRange = new ImageSubresourceRange
            {
                AspectMask = ImageAspectFlags.ColorBit,
                BaseMipLevel = level,
                LevelCount = 1,
                BaseArrayLayer = 0,
                LayerCount = Silk.NET.Vulkan.Vk.RemainingArrayLayers,
            },
        };
        var dependency = new DependencyInfo
        {
            SType = StructureType.DependencyInfo,
            ImageMemoryBarrierCount = 1,
            PImageMemoryBarriers = &barrier,
        };
        _vk.CmdPipelineBarrier2(commands, &dependency);
    }

    /// <summary>Reclaims staging bytes and temporary buffers belonging to completed frames.</summary>
    internal void ReleaseCompleted(long completedSerial) => _ringState.Release(completedSerial);

    private (Buffer Buffer, ulong Offset) Stage(
        ReadOnlySpan<byte> data,
        string ownerName,
        ulong alignmentBytes = 4)
    {
        long serial = _flights.OpenSerial != 0 ? _flights.OpenSerial : _flights.SubmittedSerial + 1;
        if (_ringState.TryAllocate(data.Length, alignmentBytes, serial, out ulong offset))
        {
            data.CopyTo(_stagingAllocation.AsSpan().Slice((int)offset, data.Length));
            return (_stagingBuffer, offset);
        }

        (Buffer temporary, VulkanAllocation allocation) = CreateHostBuffer(
            (ulong)data.Length,
            BufferUsageFlags.TransferSrcBit,
            $"vk-staging-temp-{ownerName}");
        data.CopyTo(allocation.AsSpan());
        _temporaries.Add(new TemporaryStaging(temporary, allocation));

        Buffer captured = temporary;
        VulkanAllocation capturedAllocation = allocation;
        _flights.Retire(() =>
        {
            _vk.DestroyBuffer(_device, captured, null);
            _allocator.Free(capturedAllocation);
            _temporaries.RemoveAll(entry => entry.Buffer.Handle == captured.Handle);
        });
        return (temporary, 0);
    }

    private (Buffer Buffer, VulkanAllocation Allocation) CreateHostBuffer(
        ulong sizeBytes,
        BufferUsageFlags usage,
        string name)
    {
        var create = new BufferCreateInfo
        {
            SType = StructureType.BufferCreateInfo,
            Size = sizeBytes,
            Usage = usage,
            SharingMode = SharingMode.Exclusive,
        };
        VulkanInterop.Check(
            _vk.CreateBuffer(_device, &create, null, out Buffer buffer),
            $"vkCreateBuffer ({name})");

        _vk.GetBufferMemoryRequirements(_device, buffer, out MemoryRequirements requirements);
        VulkanAllocation allocation = _allocator.Allocate(
            requirements,
            GpuMemoryResidency.HostWritable,
            name);
        VulkanInterop.Check(
            _vk.BindBufferMemory(_device, buffer, allocation.Memory, allocation.OffsetBytes),
            $"vkBindBufferMemory ({name})");
        _debugNames.NameBuffer(buffer, name);
        return (buffer, allocation);
    }

    public void Dispose()
    {
        if (_disposed)
            return;
        _disposed = true;

        foreach (TemporaryStaging temporary in _temporaries)
        {
            _vk.DestroyBuffer(_device, temporary.Buffer, null);
            _allocator.Free(temporary.Allocation);
        }
        _temporaries.Clear();

        _vk.DestroyBuffer(_device, _stagingBuffer, null);
        _allocator.Free(_stagingAllocation);
        _ringState.Reset();
        _bufferCopies.Clear();
        _imageCopies.Clear();
        
    }
}