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

using Silk.NET.Vulkan;

namespace AcDream.App.Rendering.Gpu.Vk;

/// <summary>
/// Campaign V slice V6c, plan §4.4: sets 0 and 1 for one flight slot, bound with
/// dynamic offsets so no descriptor is ever written mid-frame.
///
/// <para>The contract lets a renderer bind an arbitrary buffer range per draw,
/// and ring allocations mean that range moves every frame. The obvious
/// implementation — write a descriptor per bind — would put a
/// <c>vkUpdateDescriptorSets</c> in the hot path and reintroduce the exact cost
/// the texture table was designed to remove. So each ring-fed binding is a
/// <c>*_BUFFER_DYNAMIC</c> descriptor pointing at the whole ring, and the
/// per-draw offset travels in <c>vkCmdBindDescriptorSets</c>'s dynamic-offset
/// array, which is free.</para>
///
/// <para><b>Not every binding is dynamic.</b> Slice V6g split set 0 by
/// <see cref="VulkanPipelineLayouts.IsDynamicStorageBinding"/>, because ten
/// dynamic storage descriptors exceeded the device limit (plan §5.5.7 defect 1).
/// A plain binding carries its offset in the descriptor itself, so it is
/// rewritten when the range moves rather than when only the buffer changes — and
/// its slot in the dynamic-offset array does not exist. Getting that array's
/// length or ordering wrong is a validation error, so both are derived from the
/// same predicate the layout is built from rather than restated.</para>
///
/// <para><b>Every binding is always bound, whether a renderer uses it or
/// not.</b> Bindings a shader does not declare still need a live descriptor, so
/// unused ones point at a shared dummy range. That is what lets there be ONE
/// descriptor set layout and one pipeline layout rather than a permutation per
/// renderer — plan §4.4's requirement, and the thing that makes switching
/// pipelines mid-pass free.</para>
///
/// <para><b>Slice V6i: one set pair per renderer scope.</b> There is no longer a
/// single (set 0, set 1) pair per flight slot; there is an arena of them, and
/// <see cref="VulkanBindingScopeArena"/> decides which pair a bind belongs to and
/// whether its descriptors have to be written. That is what closes plan §5.5.8's
/// recorded one-binding-two-buffers hazard before the world arm fires it — see
/// the arena's own documentation for why the scope is derived from the
/// descriptor state rather than declared by the renderer.</para>
/// </summary>
internal sealed unsafe class VulkanFrameBindings : IDisposable
{
    private readonly Silk.NET.Vulkan.Vk _vk;
    private readonly Device _device;
    private readonly VulkanPipelineLayouts.Created _layouts;
    private readonly VulkanBindingScopeArena _arena;
    private readonly List<DescriptorPool> _pools = [];
    private readonly List<(DescriptorSet Storage, DescriptorSet Uniform)> _sets = [];

    private bool _disposed;

    /// <summary>
    /// Set pairs one descriptor pool serves. Distinct descriptor states in a
    /// frame are the world renderers plus the retained UI, so this is generous;
    /// exceeding it allocates another pool rather than failing.
    /// </summary>
    private const int PairsPerPool = 16;

    /// <summary>
    /// Set 0's dynamic-offset slots, in binding order — the order
    /// <c>vkCmdBindDescriptorSets</c> requires. A plain binding has no slot.
    /// </summary>
    private static readonly uint[] DynamicStorageBindings = BuildDynamicStorageBindings();

    private static uint[] BuildDynamicStorageBindings()
    {
        var bindings = new List<uint>((int)GpuBindingModel.StorageBindingCount);
        for (uint binding = 0; binding < GpuBindingModel.StorageBindingCount; binding++)
        {
            if (VulkanPipelineLayouts.IsDynamicStorageBinding(binding))
                bindings.Add(binding);
        }

        return [.. bindings];
    }

    /// <summary>
    /// Bindings 0..4 of set 1. Slice V6i-2 raised this from 4 when the layout
    /// gained binding 4 (sky params); binding 0 remains unused and is counted
    /// only so the bookkeeping arrays stay index-aligned with the binding number.
    /// Which of them the layout DECLARES is
    /// <see cref="VulkanPipelineLayouts.IsDeclaredUniformBinding"/>.
    /// </summary>
    internal const int UniformBindingCount = 5;

    /// <summary>
    /// How many of set 1's bindings the layout actually declares, all dynamic.
    /// Asserted against <c>maxDescriptorSetUniformBuffersDynamic</c> by the
    /// capability gate; Vulkan guarantees 8, so this is comfortable.
    /// </summary>
    internal static uint DynamicUniformBindingCount { get; } =
        (uint)VulkanPipelineLayouts.DeclaredUniformBindings.Length;

    /// <summary>
    /// Widest range any single binding may address. Dynamic descriptors take a
    /// static range at write time and slide it with an offset, so this bounds
    /// how much of the ring one binding can see at once.
    /// </summary>
    internal const uint MaxBindingRangeBytes = 4 * 1024 * 1024;

    internal VulkanFrameBindings(
        Silk.NET.Vulkan.Vk vk,
        Device device,
        VulkanPipelineLayouts.Created layouts,
        VulkanGpuBuffer ring,
        VulkanGpuBuffer dummy)
    {
        _vk = vk ?? throw new ArgumentNullException(nameof(vk));
        _device = device;
        _layouts = layouts ?? throw new ArgumentNullException(nameof(layouts));
        ArgumentNullException.ThrowIfNull(ring);
        ArgumentNullException.ThrowIfNull(dummy);

        Ring = ring;
        Dummy = dummy;
        _arena = new VulkanBindingScopeArena(
            (int)GpuBindingModel.StorageBindingCount,
            UniformBindingCount,
            VulkanPipelineLayouts.IsDynamicStorageBinding);

        uint dummyStorageRange = (uint)Math.Min(dummy.SizeBytes, MaxBindingRangeBytes);
        for (uint binding = 0; binding < GpuBindingModel.StorageBindingCount; binding++)
            _arena.SeedStorage(binding, dummy.Handle.Handle, offsetBytes: 0, dummyStorageRange);

        uint dummyUniformRange = (uint)Math.Min(dummy.SizeBytes, 65536);
        for (uint binding = 0; binding < UniformBindingCount; binding++)
            _arena.SeedUniform(binding, dummy.Handle.Handle, dummyUniformRange);
    }

    internal VulkanGpuBuffer Ring { get; }

    internal VulkanGpuBuffer Dummy { get; }

    /// <summary>
    /// Diagnostic: distinct descriptor states this flight slot materialised. One
    /// per renderer scope in a steady frame, so a number that keeps climbing is
    /// a renderer pointing a binding at a fresh buffer every draw.
    /// </summary>
    internal int ScopeCount => _arena.Count;

    /// <summary>
    /// Recycles the arena for a new frame on this slot. Safe because the slot's
    /// previous submission has retired before <c>BeginFrame</c> returns, which is
    /// the same guarantee that lets the ring rewind.
    /// </summary>
    internal void BeginFrame() => _arena.BeginFrame();

    internal void SetStorage(uint binding, VulkanGpuBuffer buffer, uint offsetBytes, uint sizeBytes)
    {
        ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(binding, GpuBindingModel.StorageBindingCount);
        uint descriptorOffset = VulkanPipelineLayouts.IsDynamicStorageBinding(binding) ? 0 : offsetBytes;
        _arena.SetStorage(
            binding,
            buffer.Handle.Handle,
            offsetBytes,
            ClampRange(buffer, sizeBytes, descriptorOffset));
    }

    internal void SetUniform(uint binding, VulkanGpuBuffer buffer, uint offsetBytes, uint sizeBytes)
    {
        ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(binding, (uint)UniformBindingCount);
        _arena.SetUniform(
            binding,
            buffer.Handle.Handle,
            offsetBytes,
            Math.Min(ClampRange(buffer, sizeBytes, offsetBytes: 0), 65536));
    }

    /// <summary>Binds all three sets with the current dynamic offsets.</summary>
    internal void Bind(CommandBuffer commands, VulkanGpuDevice device)
    {
        (int index, int slot, bool needsWrite) = _arena.Resolve();
        if (slot < 0)
        {
            slot = _sets.Count;
            _sets.Add(AllocatePair());
            _arena.AssignSlot(index, slot);
        }

        if (needsWrite)
            WritePair(_sets[slot]);

        DescriptorSet* sets = stackalloc DescriptorSet[3];
        sets[0] = _sets[slot].Storage;
        sets[1] = _sets[slot].Uniform;
        sets[2] = device.TextureTable.Set;

        uint[] declaredUniforms = VulkanPipelineLayouts.DeclaredUniformBindings;
        int dynamicCount = DynamicStorageBindings.Length + declaredUniforms.Length;
        uint* offsets = stackalloc uint[dynamicCount];
        // Dynamic offsets are ordered by set, then by binding number, and only
        // the DYNAMIC descriptors have a slot at all.
        for (int i = 0; i < DynamicStorageBindings.Length; i++)
            offsets[i] = _arena.StorageOffset(DynamicStorageBindings[i]);
        for (int i = 0; i < declaredUniforms.Length; i++)
            offsets[DynamicStorageBindings.Length + i] = _arena.UniformOffset(declaredUniforms[i]);

        _vk.CmdBindDescriptorSets(
            commands,
            PipelineBindPoint.Graphics,
            device.Layouts.PipelineLayout,
            0,
            3,
            sets,
            (uint)dynamicCount,
            offsets);
    }

    private void WritePair((DescriptorSet Storage, DescriptorSet Uniform) pair)
    {
        for (uint binding = 0; binding < GpuBindingModel.StorageBindingCount; binding++)
        {
            WriteStorage(
                pair.Storage,
                binding,
                new Silk.NET.Vulkan.Buffer(_arena.StorageBuffer(binding)),
                _arena.StorageDescriptorOffset(binding),
                _arena.StorageRange(binding),
                VulkanPipelineLayouts.IsDynamicStorageBinding(binding));
        }

        // Only the bindings the layout declares exist; the rest of the array is
        // bookkeeping so the offsets stay index-aligned with the binding number.
        foreach (uint binding in VulkanPipelineLayouts.DeclaredUniformBindings)
        {
            WriteUniform(
                pair.Uniform,
                binding,
                new Silk.NET.Vulkan.Buffer(_arena.UniformBuffer(binding)),
                _arena.UniformRange(binding));
        }
    }

    private (DescriptorSet Storage, DescriptorSet Uniform) AllocatePair()
    {
        if (_sets.Count % PairsPerPool == 0)
            _pools.Add(CreatePool());
        DescriptorPool pool = _pools[^1];
        return (Allocate(pool, _layouts.Storage), Allocate(pool, _layouts.Uniform));
    }

    private DescriptorPool CreatePool()
    {
        DescriptorPoolSize* sizes = stackalloc DescriptorPoolSize[3];
        sizes[0] = new DescriptorPoolSize
        {
            Type = DescriptorType.StorageBufferDynamic,
            DescriptorCount = VulkanPipelineLayouts.DynamicStorageBindingCount * PairsPerPool,
        };
        sizes[1] = new DescriptorPoolSize
        {
            Type = DescriptorType.StorageBuffer,
            DescriptorCount =
                (GpuBindingModel.StorageBindingCount - VulkanPipelineLayouts.DynamicStorageBindingCount)
                * PairsPerPool,
        };
        sizes[2] = new DescriptorPoolSize
        {
            Type = DescriptorType.UniformBufferDynamic,
            DescriptorCount = DynamicUniformBindingCount * PairsPerPool,
        };
        var poolCreate = new DescriptorPoolCreateInfo
        {
            SType = StructureType.DescriptorPoolCreateInfo,
            MaxSets = 2 * PairsPerPool,
            PoolSizeCount = 3,
            PPoolSizes = sizes,
        };
        VulkanInterop.Check(
            _vk.CreateDescriptorPool(_device, &poolCreate, null, out DescriptorPool pool),
            "vkCreateDescriptorPool (frame bindings)");
        return pool;
    }

    private static uint ClampRange(VulkanGpuBuffer buffer, uint requested, uint offsetBytes)
    {
        long remaining = buffer.SizeBytes - offsetBytes;
        if (remaining <= 0)
        {
            throw new ArgumentOutOfRangeException(
                nameof(offsetBytes),
                offsetBytes,
                $"A storage binding was pointed past the end of its {buffer.SizeBytes}-byte buffer. " +
                "A descriptor range of zero is not representable in Vulkan.");
        }

        uint available = (uint)Math.Min(remaining, MaxBindingRangeBytes);
        return requested == 0 ? available : Math.Min(Math.Max(requested, 16), available);
    }

    private DescriptorSet Allocate(DescriptorPool pool, DescriptorSetLayout layout)
    {
        DescriptorSetLayout handle = layout;
        var allocate = new DescriptorSetAllocateInfo
        {
            SType = StructureType.DescriptorSetAllocateInfo,
            DescriptorPool = pool,
            DescriptorSetCount = 1,
            PSetLayouts = &handle,
        };
        VulkanInterop.Check(
            _vk.AllocateDescriptorSets(_device, &allocate, out DescriptorSet set),
            "vkAllocateDescriptorSets (frame bindings)");
        return set;
    }

    private void WriteStorage(
        DescriptorSet set,
        uint binding,
        Silk.NET.Vulkan.Buffer buffer,
        uint offsetBytes,
        uint rangeBytes,
        bool dynamic)
    {
        var info = new DescriptorBufferInfo
        {
            Buffer = buffer,
            Offset = offsetBytes,
            Range = rangeBytes,
        };
        var write = new WriteDescriptorSet
        {
            SType = StructureType.WriteDescriptorSet,
            DstSet = set,
            DstBinding = binding,
            DescriptorCount = 1,
            DescriptorType = dynamic
                ? DescriptorType.StorageBufferDynamic
                : DescriptorType.StorageBuffer,
            PBufferInfo = &info,
        };
        _vk.UpdateDescriptorSets(_device, 1, &write, 0, null);
    }

    private void WriteUniform(
        DescriptorSet set,
        uint binding,
        Silk.NET.Vulkan.Buffer buffer,
        uint rangeBytes)
    {
        var info = new DescriptorBufferInfo
        {
            Buffer = buffer,
            Offset = 0,
            Range = rangeBytes,
        };
        var write = new WriteDescriptorSet
        {
            SType = StructureType.WriteDescriptorSet,
            DstSet = set,
            DstBinding = binding,
            DescriptorCount = 1,
            DescriptorType = DescriptorType.UniformBufferDynamic,
            PBufferInfo = &info,
        };
        _vk.UpdateDescriptorSets(_device, 1, &write, 0, null);
    }

    public void Dispose()
    {
        if (_disposed)
            return;
        _disposed = true;
        foreach (DescriptorPool pool in _pools)
        {
            if (pool.Handle != 0)
                _vk.DestroyDescriptorPool(_device, pool, null);
        }

        _pools.Clear();
        _sets.Clear();
    }
}