//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// Where DescriptorHeapVulkan::Initialize puts the buffer descriptors and the
// reserved ranges (#163). A descriptor_heap array indexes the heap from byte 0
// in strides of its descriptor size, so buffer slot 0 has to sit on a whole
// buffer descriptor at or past the end of the image descriptors. Its index
// was rounded down instead: where the image descriptors were not a whole
// number of buffer descriptors, the first buffer descriptors overwrote the
// last image descriptors, and where they were less than one, the heap was
// sized as if the buffers began right after the images, so the last buffer
// descriptors landed in the implementation's reserved range. The reserved
// range sat at a multiple of imageDescriptorAlignment, where the spec asks
// for bufferDescriptorAlignment too, an image region past 65535 buffer
// descriptors wrapped bufferStartElement around, and on a driver that
// reserves nothing (llvmpipe) a heap without sampler slots was a zero-size
// buffer. Most of it needs descriptor sizes neither NVIDIA nor llvmpipe has,
// so the test feeds DescriptorHeapVulkan::ComputeLayout every combination of
// sizes and alignments the spec allows (powers of two, each alignment at
// most its size) and checks where everything lands, that no heap is empty,
// that those two drivers keep the layouts they had, and that slot counts
// whose buffer indices do not fit 16 bits throw while image-only heaps,
// which have no buffer index to fit, do not.
//
// Pure CPU logic, no Vulkan device.

#include <cstdlib>
#include "vulkan/vulkan.h"

import Crafter.Graphics;
import std;

using namespace Crafter;

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

using Layout = DescriptorHeapVulkan::HeapLayout;

VkPhysicalDeviceDescriptorHeapPropertiesEXT Properties(VkDeviceSize imageSize, VkDeviceSize imageAlignment,
                                                       VkDeviceSize bufferSize, VkDeviceSize bufferAlignment,
                                                       VkDeviceSize samplerSize, VkDeviceSize samplerAlignment,
                                                       VkDeviceSize resourceReserved, VkDeviceSize samplerReserved) {
    return {
        .sType                        = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_HEAP_PROPERTIES_EXT,
        .minSamplerHeapReservedRange  = samplerReserved,
        .minResourceHeapReservedRange = resourceReserved,
        .samplerDescriptorSize        = samplerSize,
        .imageDescriptorSize          = imageSize,
        .bufferDescriptorSize         = bufferSize,
        .samplerDescriptorAlignment   = samplerAlignment,
        .imageDescriptorAlignment     = imageAlignment,
        .bufferDescriptorAlignment    = bufferAlignment,
    };
}

// The development machine's two drivers, from vulkaninfo: NVIDIA 615.78.08
// on an RTX 4090, and llvmpipe from Mesa 26.2.5.
const VkPhysicalDeviceDescriptorHeapPropertiesEXT kNvidia   = Properties(32, 32, 16, 8, 32, 32, 0x17a00, 0x200);
const VkPhysicalDeviceDescriptorHeapPropertiesEXT kLlvmpipe = Properties(64, 4, 16, 4, 32, 4, 0, 0);

// 1, 2, 4, ... up to `max`.
std::vector<VkDeviceSize> PowersOfTwo(VkDeviceSize max) {
    std::vector<VkDeviceSize> powers;
    for (VkDeviceSize v = 1; v <= max; v *= 2) {
        powers.push_back(v);
    }
    return powers;
}

std::string Describe(const VkPhysicalDeviceDescriptorHeapPropertiesEXT& p) {
    return std::format("{}-byte image descriptors aligned to {}, {}-byte buffer descriptors aligned to {}, a {}-byte reserved range",
                       p.imageDescriptorSize, p.imageDescriptorAlignment, p.bufferDescriptorSize, p.bufferDescriptorAlignment,
                       p.minResourceHeapReservedRange);
}

std::string Format(const Layout& l) {
    return std::format("buffers from descriptor {} (byte {}), the resource heap's reserved range at {} of {} bytes, the sampler heap's at {} of {}",
                       l.bufferStartElement, l.bufferStartOffset, l.resourceReservedRangeOffset, l.resourceHeapSize,
                       l.samplerReservedRangeOffset, l.samplerHeapSize);
}

// What is wrong with where the resource heap's descriptors and reserved range
// land, or "" when nothing is.
std::string ResourceProblem(const VkPhysicalDeviceDescriptorHeapPropertiesEXT& p, std::uint16_t images, std::uint16_t buffers) {
    const Layout l = DescriptorHeapVulkan::ComputeLayout(p, images, buffers, 0);
    const std::uint64_t imageEnd = images * p.imageDescriptorSize;
    const std::uint64_t bufferEnd = l.bufferStartOffset + buffers * p.bufferDescriptorSize;
    if (buffers == 0) {
        if (l.bufferStartElement != 0 || l.bufferStartOffset != 0) {
            return std::format("with no buffer slots the buffer region starts at descriptor {}, byte {}, not at 0", l.bufferStartElement, l.bufferStartOffset);
        }
    } else {
        if (l.bufferStartOffset != l.bufferStartElement * p.bufferDescriptorSize) {
            return std::format("buffer slot 0 is at byte {}, not at buffer descriptor {}, where a shader reads it", l.bufferStartOffset, l.bufferStartElement);
        }
        if (l.bufferStartOffset < imageEnd) {
            return std::format("the buffer descriptors start at byte {}, over the image descriptors, which end at {}", l.bufferStartOffset, imageEnd);
        }
        if (l.bufferStartOffset >= imageEnd + p.bufferDescriptorSize) {
            return std::format("the buffer descriptors start at byte {}, a whole descriptor past the images' end at {}", l.bufferStartOffset, imageEnd);
        }
        if (l.bufferStartOffset % p.bufferDescriptorAlignment != 0) {
            return std::format("the buffer descriptors start at byte {}, not a multiple of bufferDescriptorAlignment", l.bufferStartOffset);
        }
    }
    const std::uint64_t descriptorEnd = std::max(imageEnd, bufferEnd);
    if (l.resourceReservedRangeOffset < descriptorEnd) {
        return std::format("the reserved range starts at byte {}, before the descriptors end at {}", l.resourceReservedRangeOffset, descriptorEnd);
    }
    if (l.resourceReservedRangeOffset % p.imageDescriptorAlignment != 0 || l.resourceReservedRangeOffset % p.bufferDescriptorAlignment != 0) {
        return std::format("the reserved range starts at byte {}, not a multiple of both descriptor alignments", l.resourceReservedRangeOffset);
    }
    if (l.resourceHeapSize == 0) {
        return "the heap is zero bytes, which no buffer can be";
    }
    if (l.resourceHeapSize != l.resourceReservedRangeOffset + p.minResourceHeapReservedRange) {
        return std::format("the heap is {} bytes, where the reserved range runs from {} to {}", l.resourceHeapSize,
                           l.resourceReservedRangeOffset, l.resourceReservedRangeOffset + p.minResourceHeapReservedRange);
    }
    return "";
}

std::string SamplerProblem(const VkPhysicalDeviceDescriptorHeapPropertiesEXT& p, std::uint16_t samplers) {
    const Layout l = DescriptorHeapVulkan::ComputeLayout(p, 0, 0, samplers);
    const std::uint64_t samplerEnd = samplers * p.samplerDescriptorSize;
    if (l.samplerReservedRangeOffset < samplerEnd) {
        return std::format("the reserved range starts at byte {}, before the sampler descriptors end at {}", l.samplerReservedRangeOffset, samplerEnd);
    }
    if (l.samplerReservedRangeOffset % p.samplerDescriptorAlignment != 0) {
        return std::format("the reserved range starts at byte {}, not a multiple of samplerDescriptorAlignment", l.samplerReservedRangeOffset);
    }
    if (l.samplerHeapSize == 0) {
        return "the heap is zero bytes, which no buffer can be";
    }
    if (l.samplerHeapSize != l.samplerReservedRangeOffset + p.minSamplerHeapReservedRange) {
        return std::format("the heap is {} bytes, where the reserved range runs from {} to {}", l.samplerHeapSize,
                           l.samplerReservedRangeOffset, l.samplerReservedRangeOffset + p.minSamplerHeapReservedRange);
    }
    return "";
}

std::string Thrown(std::uint16_t images, std::uint16_t buffers) {
    try {
        (void)DescriptorHeapVulkan::ComputeLayout(kLlvmpipe, images, buffers, 0);
    } catch (const std::runtime_error& e) {
        return e.what();
    }
    return "";
}

} // namespace

int main() {
    // ── The two layouts the issue describes ─────────────────────────────────
    // Image descriptors smaller than buffer descriptors. The spec makes every
    // size a power of two, so this is the only way the image region can end
    // inside a buffer descriptor.
    {
        const VkPhysicalDeviceDescriptorHeapPropertiesEXT p = Properties(32, 32, 64, 64, 32, 32, 0x17a00, 0x200);
        // Three images end at byte 96, halfway through buffer descriptor 1.
        const Layout overlap = DescriptorHeapVulkan::ComputeLayout(p, 3, 1, 0);
        Check(overlap.bufferStartElement == 2 && overlap.bufferStartOffset == 128,
              std::format("3 32-byte images, then 64-byte buffers: buffer slot 0 is descriptor 2, byte 128, past the images (got descriptor {}, byte {})",
                          overlap.bufferStartElement, overlap.bufferStartOffset));
        // One image is less than one buffer descriptor: the four buffers are
        // descriptors 1 to 4, bytes 64 to 320, and the reserved range follows.
        const Layout small = DescriptorHeapVulkan::ComputeLayout(p, 1, 4, 0);
        Check(small.bufferStartOffset == 64 && small.resourceReservedRangeOffset == 320 && small.resourceHeapSize == 320 + 0x17a00,
              std::format("1 32-byte image, then 4 64-byte buffers: the buffers end at byte 320 and the reserved range starts there (got buffers from byte {}, the reserved range at {}, a {}-byte heap)",
                          small.bufferStartOffset, small.resourceReservedRangeOffset, small.resourceHeapSize));
    }

    // ── The reserved range on a multiple of both alignments ─────────────────
    // VUID-vkCmdBindResourceHeapEXT-pBindInfo-11435 and -11436. Five 16-byte
    // images end at byte 80, a multiple of their own alignment but not of the
    // buffer descriptors' 64.
    {
        const VkPhysicalDeviceDescriptorHeapPropertiesEXT p = Properties(16, 16, 64, 64, 32, 32, 0x17a00, 0x200);
        const Layout l = DescriptorHeapVulkan::ComputeLayout(p, 5, 0, 0);
        Check(l.resourceReservedRangeOffset == 128,
              std::format("5 16-byte images and no buffers, with buffer descriptors aligned to 64: the reserved range starts at byte 128 (got {})",
                          l.resourceReservedRangeOffset));
    }

    // ── Every descriptor size and alignment the spec allows ─────────────────
    // Powers of two from 1 to 256 bytes, each alignment a power of two no
    // larger than its size, with and without a reserved range.
    {
        std::uint64_t layouts = 0;
        std::string first;
        for (const VkDeviceSize reserved : {VkDeviceSize{0}, VkDeviceSize{0x17a00}}) {
            for (const VkDeviceSize imageSize : PowersOfTwo(256)) {
                for (const VkDeviceSize imageAlignment : PowersOfTwo(imageSize)) {
                    for (const VkDeviceSize bufferSize : PowersOfTwo(256)) {
                        for (const VkDeviceSize bufferAlignment : PowersOfTwo(bufferSize)) {
                            const VkPhysicalDeviceDescriptorHeapPropertiesEXT p =
                                Properties(imageSize, imageAlignment, bufferSize, bufferAlignment, 32, 32, reserved, 0x200);
                            for (const std::uint16_t images : {0, 1, 2, 3, 5, 7, 8, 9, 31, 64}) {
                                for (const std::uint16_t buffers : {0, 1, 2, 3, 16}) {
                                    ++layouts;
                                    const std::string problem = ResourceProblem(p, images, buffers);
                                    if (!problem.empty() && first.empty()) {
                                        first = std::format("{} images and {} buffers with {}: {}", images, buffers, Describe(p), problem);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        Check(first.empty(), std::format("{} resource heap layouts over every legal descriptor size and alignment ({})",
                                         layouts, first.empty() ? "every one sound" : first));
    }
    {
        std::uint64_t layouts = 0;
        std::string first;
        for (const VkDeviceSize reserved : {VkDeviceSize{0}, VkDeviceSize{0x200}}) {
            for (const VkDeviceSize size : PowersOfTwo(256)) {
                for (const VkDeviceSize alignment : PowersOfTwo(size)) {
                    const VkPhysicalDeviceDescriptorHeapPropertiesEXT p = Properties(32, 32, 16, 16, size, alignment, 0x17a00, reserved);
                    for (const std::uint16_t samplers : {0, 1, 2, 3, 7, 16}) {
                        ++layouts;
                        const std::string problem = SamplerProblem(p, samplers);
                        if (!problem.empty() && first.empty()) {
                            first = std::format("{} {}-byte samplers aligned to {}, a {}-byte reserved range: {}", samplers, size, alignment, reserved, problem);
                        }
                    }
                }
            }
        }
        Check(first.empty(), std::format("and {} sampler heap layouts ({})", layouts, first.empty() ? "every one sound" : first));
    }

    // ── No heap is empty ────────────────────────────────────────────────────
    // llvmpipe reserves nothing, so a heap with no slots of its kind was a
    // zero-size buffer, which Vulkan does not allow, and every Initialize(1, 1,
    // 0) made one.
    {
        const Layout l = DescriptorHeapVulkan::ComputeLayout(kLlvmpipe, 1, 1, 0);
        Check(l.samplerHeapSize == 32, std::format("llvmpipe, no sampler slots: the sampler heap holds one 32-byte descriptor (got {} bytes)", l.samplerHeapSize));
        const Layout none = DescriptorHeapVulkan::ComputeLayout(kLlvmpipe, 0, 0, 0);
        Check(none.resourceHeapSize > 0, std::format("and with no slots at all the resource heap is not empty either ({} bytes)", none.resourceHeapSize));
    }

    // ── The drivers checked keep their layouts ──────────────────────────────
    // On NVIDIA and llvmpipe an image descriptor is a whole number of buffer
    // descriptors, and buffer descriptors are aligned no more strictly than
    // image ones, so nothing rounds: these are the layouts Initialize gave
    // them before, for slot counts the examples and tests use, except that a
    // heap with no sampler slots now holds one sampler descriptor.
    {
        struct Case {
            std::string_view driver;
            VkPhysicalDeviceDescriptorHeapPropertiesEXT properties;
            std::uint16_t images;
            std::uint16_t buffers;
            std::uint16_t samplers;
            Layout expected;
        };
        const Case cases[] = {
            { "NVIDIA", kNvidia, 1, 1, 0,
              { .bufferStartElement = 2, .bufferStartOffset = 32, .resourceReservedRangeOffset = 64, .resourceHeapSize = 64 + 0x17a00,
                .samplerReservedRangeOffset = 32, .samplerHeapSize = 32 + 0x200 } },
            { "NVIDIA", kNvidia, 3, 3, 1,
              { .bufferStartElement = 6, .bufferStartOffset = 96, .resourceReservedRangeOffset = 160, .resourceHeapSize = 160 + 0x17a00,
                .samplerReservedRangeOffset = 32, .samplerHeapSize = 32 + 0x200 } },
            { "NVIDIA", kNvidia, 16, 16, 4,
              { .bufferStartElement = 32, .bufferStartOffset = 512, .resourceReservedRangeOffset = 768, .resourceHeapSize = 768 + 0x17a00,
                .samplerReservedRangeOffset = 128, .samplerHeapSize = 128 + 0x200 } },
            { "llvmpipe", kLlvmpipe, 1, 1, 0,
              { .bufferStartElement = 4, .bufferStartOffset = 64, .resourceReservedRangeOffset = 80, .resourceHeapSize = 80,
                .samplerReservedRangeOffset = 32, .samplerHeapSize = 32 } },
            { "llvmpipe", kLlvmpipe, 8, 16, 4,
              { .bufferStartElement = 32, .bufferStartOffset = 512, .resourceReservedRangeOffset = 768, .resourceHeapSize = 768,
                .samplerReservedRangeOffset = 128, .samplerHeapSize = 128 } },
            { "llvmpipe", kLlvmpipe, 16, 16, 4,
              { .bufferStartElement = 64, .bufferStartOffset = 1024, .resourceReservedRangeOffset = 1280, .resourceHeapSize = 1280,
                .samplerReservedRangeOffset = 128, .samplerHeapSize = 128 } },
        };
        for (const Case& c : cases) {
            const Layout l = DescriptorHeapVulkan::ComputeLayout(c.properties, c.images, c.buffers, c.samplers);
            Check(l == c.expected, std::format("{}, Initialize({}, {}, {}): {}{}", c.driver, c.images, c.buffers, c.samplers,
                                               Format(c.expected), l == c.expected ? "" : std::format(" (got {})", Format(l))));
        }
    }

    // ── Buffer indices past 16 bits throw ───────────────────────────────────
    // Shaders and BufferSlot carry absolute buffer indices in 16 bits, 0xFFFF
    // being the empty slot. llvmpipe's 64-byte image descriptors are four
    // 16-byte buffer descriptors each, so 16383 images put buffer slot 0 at
    // index 65532.
    {
        const std::string fits = Thrown(16383, 3);
        Check(fits.empty(), std::format("16383 images and 3 buffers on llvmpipe take buffer indices up to 65534 ({})", fits.empty() ? "they fit" : fits));
        const std::string past = Thrown(16383, 4);
        Check(!past.empty(), std::format("a fourth buffer would take index 65535, the empty slot, and throws ({})", past.empty() ? "it did not" : past));
        const std::string wrapped = Thrown(20000, 1);
        Check(!wrapped.empty(), std::format("20000 images put buffer slot 0 at index 80000 and throw rather than wrap around ({})",
                                            wrapped.empty() ? "it did not" : wrapped));
        // Without buffer slots no buffer index has to fit, however far the
        // images reach.
        const std::string imageOnly = Thrown(0xFFFF, 0);
        Check(imageOnly.empty(), std::format("65535 images and no buffers are accepted ({})", imageOnly.empty() ? "they were" : imageOnly));
        if (imageOnly.empty()) {
            const Layout l = DescriptorHeapVulkan::ComputeLayout(kLlvmpipe, 0xFFFF, 0, 0);
            Check(l.bufferStartElement == 0 && l.resourceReservedRangeOffset == 0xFFFF * 64,
                  std::format("with the empty buffer region at 0 and the reserved range right after the images, at byte {} (got buffers from descriptor {}, the reserved range at {})",
                              0xFFFF * 64, l.bufferStartElement, l.resourceReservedRangeOffset));
        }
    }

    std::println("{}", failures == 0 ? "all checks passed" : "check(s) failed");
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}