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>
/// </summary>
internal sealed unsafe class VulkanFrameBindings : IDisposable
{
    private readonly Silk.NET.Vulkan.Vk _vk;
    private readonly Device _device;
    private readonly DescriptorPool _pool;
    private readonly DescriptorSet _storageSet;
    private readonly DescriptorSet _uniformSet;

    private readonly uint[] _storageOffsets = new uint[GpuBindingModel.StorageBindingCount];
    private readonly uint[] _uniformOffsets = new uint[UniformBindingCount];
    private readonly Silk.NET.Vulkan.Buffer[] _storageBuffers =
        new Silk.NET.Vulkan.Buffer[GpuBindingModel.StorageBindingCount];
    private readonly Silk.NET.Vulkan.Buffer[] _uniformBuffers =
        new Silk.NET.Vulkan.Buffer[UniformBindingCount];

    private bool _disposed;

    /// <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..3 of set 1; only 1 (SceneLighting) and 3 (terrain tiling) are used.</summary>
    internal const int UniformBindingCount = 4;

    /// <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 const uint DynamicUniformBindingCount = 2;

    /// <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;
        ArgumentNullException.ThrowIfNull(layouts);
        ArgumentNullException.ThrowIfNull(ring);
        ArgumentNullException.ThrowIfNull(dummy);

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

        _storageSet = Allocate(layouts.Storage);
        _uniformSet = Allocate(layouts.Uniform);

        for (uint binding = 0; binding < GpuBindingModel.StorageBindingCount; binding++)
        {
            _storageBuffers[binding] = dummy.Handle;
            WriteStorage(
                binding,
                dummy.Handle,
                offsetBytes: 0,
                (uint)Math.Min(dummy.SizeBytes, MaxBindingRangeBytes),
                VulkanPipelineLayouts.IsDynamicStorageBinding(binding));
        }

        // Only the two bindings the layout declares exist; the rest of the
        // array is bookkeeping so the offsets stay index-aligned.
        WriteUniform(GpuBindingModel.UniformSceneLighting, dummy.Handle, (uint)Math.Min(dummy.SizeBytes, 65536));
        WriteUniform(GpuBindingModel.UniformTerrainTiling, dummy.Handle, (uint)Math.Min(dummy.SizeBytes, 65536));
        _uniformBuffers[GpuBindingModel.UniformSceneLighting] = dummy.Handle;
        _uniformBuffers[GpuBindingModel.UniformTerrainTiling] = dummy.Handle;

        Ring = ring;
        Dummy = dummy;
    }

    internal VulkanGpuBuffer Ring { get; }

    internal VulkanGpuBuffer Dummy { get; }

    /// <summary>
    /// Points a storage binding at a range.
    ///
    /// <para>A DYNAMIC binding re-writes its descriptor only when the BUFFER
    /// changes; the offset rides the bind call. A PLAIN binding has no such
    /// channel, so the descriptor itself carries the offset and is re-written
    /// when either moves.</para>
    /// </summary>
    internal void SetStorage(uint binding, VulkanGpuBuffer buffer, uint offsetBytes, uint sizeBytes)
    {
        ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(binding, GpuBindingModel.StorageBindingCount);
        bool dynamic = VulkanPipelineLayouts.IsDynamicStorageBinding(binding);
        bool bufferChanged = _storageBuffers[binding].Handle != buffer.Handle.Handle;
        bool offsetChanged = _storageOffsets[binding] != offsetBytes;
        if (bufferChanged || (!dynamic && offsetChanged))
        {
            _storageBuffers[binding] = buffer.Handle;
            WriteStorage(
                binding,
                buffer.Handle,
                dynamic ? 0 : offsetBytes,
                ClampRange(buffer, sizeBytes, dynamic ? 0 : offsetBytes),
                dynamic);
        }

        _storageOffsets[binding] = offsetBytes;
    }

    internal void SetUniform(uint binding, VulkanGpuBuffer buffer, uint offsetBytes, uint sizeBytes)
    {
        ArgumentOutOfRangeException.ThrowIfGreaterThanOrEqual(binding, (uint)UniformBindingCount);
        if (_uniformBuffers[binding].Handle != buffer.Handle.Handle)
        {
            _uniformBuffers[binding] = buffer.Handle;
            WriteUniform(binding, buffer.Handle, Math.Min(ClampRange(buffer, sizeBytes, offsetBytes: 0), 65536));
        }

        _uniformOffsets[binding] = offsetBytes;
    }

    /// <summary>Binds all three sets with the current dynamic offsets.</summary>
    internal void Bind(CommandBuffer commands, VulkanGpuDevice device)
    {
        DescriptorSet* sets = stackalloc DescriptorSet[3];
        sets[0] = _storageSet;
        sets[1] = _uniformSet;
        sets[2] = device.TextureTable.Set;

        int dynamicCount = DynamicStorageBindings.Length + 2;
        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] = _storageOffsets[DynamicStorageBindings[i]];
        offsets[DynamicStorageBindings.Length + 0] = _uniformOffsets[GpuBindingModel.UniformSceneLighting];
        offsets[DynamicStorageBindings.Length + 1] = _uniformOffsets[GpuBindingModel.UniformTerrainTiling];

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

    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(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(
        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 = _storageSet,
            DstBinding = binding,
            DescriptorCount = 1,
            DescriptorType = dynamic
                ? DescriptorType.StorageBufferDynamic
                : DescriptorType.StorageBuffer,
            PBufferInfo = &info,
        };
        _vk.UpdateDescriptorSets(_device, 1, &write, 0, null);
    }

    private void WriteUniform(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 = _uniformSet,
            DstBinding = binding,
            DescriptorCount = 1,
            DescriptorType = DescriptorType.UniformBufferDynamic,
            PBufferInfo = &info,
        };
        _vk.UpdateDescriptorSets(_device, 1, &write, 0, null);
    }

    public void Dispose()
    {
        if (_disposed)
            return;
        _disposed = true;
        if (_pool.Handle != 0)
            _vk.DestroyDescriptorPool(_device, _pool, null);
    }
}