vk_api.h

#pragma once

#include <slang-com-helper.h>
#include <slang-com-ptr.h>
#include <slang.h>

#include <vk_mem_alloc.h>

#include "backend/vulkan/vk_conversion.h"

namespace vk {

static constexpr u32 MAX_RESOURCES = 400'000;
static constexpr u32 MAX_SAMPLERS = 64;

struct SlotAllocator {
    u32 maxSlots;
    u32 nextSlot = 0;
    std::vector<u32> freeSlots;

    explicit SlotAllocator(u32 maxSlots) : maxSlots(maxSlots) {
    }

    u32 allocate() {
        if (!freeSlots.empty()) {
            u32 slot = freeSlots.back();
            freeSlots.pop_back();
            return slot;
        }

        if (nextSlot >= maxSlots) {
            return UINT32_MAX;
        }

        return nextSlot++;
    }

    void free(u32 slot) {
        freeSlots.push_back(slot);
    }
};

static SlotAllocator SAMPLER_SLOTS = SlotAllocator{u64(MAX_SAMPLERS)};
static SlotAllocator RESOURCE_SLOTS = SlotAllocator{u64(MAX_RESOURCES)};

// no native view, sentinel marks valid
inline const VkBufferView NON_TEXEL_BUFFER_VIEW_SENTINEL = (VkBufferView)(uintptr_t)1;

void create_global_descriptors(Device &device);
void destroy_global_descriptors(rhi::Device device);

u64 write_sampler_descriptor(Device &device, Sampler sampler);
u64 write_image_view(Device &device, VkImageView view, bool storage);
u64 write_buffer_view(Device &device, rhi::Buffer *buffer, u64 offset, u64 size, bool writeable);
void free_resource_slot(u64 slot);
void free_sampler_slot(u64 slot);

} // namespace vk

struct Platform;

namespace rhi {

struct Memory {
    VmaAllocator allocator;
};

struct Queue {
    VkQueue handle = VK_NULL_HANDLE;
    u32 family_index = UINT32_MAX;
    u32 queue_index = 0;
};

struct Buffer {
    rhi::BufferDesc desc{};
    VkBuffer _handle{};
    VmaAllocation _allocation{};
    void *mapped = nullptr;
    u64 address = 0;
    rhi::ResourceState state = rhi::ResourceState::Idle;
    std::vector<std::pair<BufferViewDesc, VkBufferView>> view_cache;
};

struct Image {
    rhi::ImageDesc desc;
    VkImage _handle{};
    VmaAllocation _alloc{};
    rhi::ResourceState state = rhi::ResourceState::Idle;
    std::vector<std::pair<ImageViewDesc, VkImageView>> view_cache;
};

struct CmdPool {
    VkCommandPool handle = VK_NULL_HANDLE;
};

struct CmdBuffer {
    VkCommandBuffer handle = VK_NULL_HANDLE;
    bool is_rendering = false;
};

struct Sync {
    VkSemaphore timeline = VK_NULL_HANDLE;
    VkDevice device = VK_NULL_HANDLE;
};

struct RenderPipeline {
    Device *device = nullptr;
    VkPipeline _handle = VK_NULL_HANDLE;
};

struct ComputePipeline {
    Device *device = nullptr;
    VkPipeline _handle = VK_NULL_HANDLE;
};

struct Sampler {
    VkSampler _handle = VK_NULL_HANDLE;
    u64 slot = UINT64_MAX;
};

struct ImageView {
    ImageViewDesc desc{};
    VkImageView _handle = nullptr;
    u64 slot = UINT64_MAX;
};

struct BufferView {
    Buffer *buffer = nullptr;
    BufferViewDesc desc{};
    VkBufferView _handle = nullptr;
    u64 slot = UINT64_MAX;
};

struct QueryPool {
    VkQueryPool _handle = VK_NULL_HANDLE;
    u32 _type;
    rhi::QueryPoolDesc desc;
};

struct AccelStruct {
    VkAccelerationStructureKHR handle = VK_NULL_HANDLE;
    Buffer *buffer = nullptr;
};

struct Swapchain {
    VkSurfaceKHR surface = VK_NULL_HANDLE;
    VkSwapchainKHR handle = VK_NULL_HANDLE;
    rhi::ImageFormat format = rhi::ImageFormat::UNDEFINED;
    rhi::ImageFormat depth_format = rhi::ImageFormat::UNDEFINED;
    rhi::Extent2D extent = {};

    bool recreate = false;

    std::vector<Image> backbuffers;
    std::vector<VkSemaphore> present;
    VkFence acquire = VK_NULL_HANDLE;
};

struct Device {
    VkInstance instance = VK_NULL_HANDLE;
    VkPhysicalDevice physical = VK_NULL_HANDLE;
    VkDevice logical = VK_NULL_HANDLE;

    VkDescriptorPool descriptor_pool = VK_NULL_HANDLE;
    VkDescriptorSetLayout descriptor_set_layout = VK_NULL_HANDLE;
    VkDescriptorSet descriptor_set = VK_NULL_HANDLE;
    VkPipelineLayout pipeline_layout = VK_NULL_HANDLE;

    Queue graphics;

    Memory memory;

    u32 max_sampled_textures = 1024;
    u32 max_storage_buffers = 1024;
    u32 max_storage_images = 1024;
    u32 max_uniform_buffers = 128;
    u32 max_samplers = 32;
    u32 max_acceleration_structures = 1024;

    u32 min_ubo_alignment = 16;
    u64 timestamp_period = 1;

    VkSurfaceKHR surface = VK_NULL_HANDLE;

    VkPipelineCache pipeline_cache = VK_NULL_HANDLE;
    VkDebugUtilsMessengerEXT debug_messenger = VK_NULL_HANDLE;
};

struct ShaderCompiler {
    Slang::ComPtr<slang::IGlobalSession> global;
    Slang::ComPtr<slang::ISession> session;

    bool recreate_session();
};

struct ShaderModule {
    VkShaderModule handle = VK_NULL_HANDLE;
    rhi::ShaderStage stage;
};

ShaderModule
load_shader(ShaderCompiler &sc, Device &device, const char *path, const char *entry_name, ShaderStage stage);

struct SyncPoint {
    Sync *sync = nullptr;
    u64 value = 0;
};

struct QueueWait {
    SyncPoint point;
    rhi::PipelineStages stage_mask;
};

struct QueueSignal {
    SyncPoint point;
};

struct QueueSubmitDesc {
    u32 wait_count = 0;
    const QueueWait *waits = nullptr;
    u32 cmd_count = 0;
    const CmdBuffer *cmds = nullptr;
    u32 signal_count = 0;
    const QueueSignal *signals = nullptr;
};

struct RenderPipelineDesc {
    Device *device = nullptr;

    std::string name;

    std::vector<rhi::ImageFormat> color_attachment_formats;
    rhi::ImageFormat depth_attachment_format = rhi::ImageFormat::UNDEFINED;

    std::string shader_path;
    std::string vs_entry, fs_entry;
    std::vector<ShaderModule> shader_modules;

    std::vector<VertexAttribute> vertex_attributes;
    std::vector<VertexBinding> vertex_bindings;

    rhi::PipelineTopology topology = rhi::PipelineTopology::TRIANGLES;
    rhi::PipelineFillMode fill_mode = rhi::PipelineFillMode::SOLID;
    rhi::PipelineCullMode cull_mode = rhi::PipelineCullMode::BACK;
    rhi::PipelineTriangleWindingOrder winding = rhi::PipelineTriangleWindingOrder::CCW;

    rhi::SampleCount samples = rhi::SampleCount::Sample1;

    bool depth_test_enable = true;
    bool depth_write_enable = true;
    bool depth_bias_enable = false;
    f32 depth_bias_constant = 0.0;
    f32 depth_bias_clamp = 0.0;
    f32 depth_bias_slope = 0.0;
    bool depth_clamp_enable = false;
    rhi::PipelineCompareOp depth_compare_op = rhi::PipelineCompareOp::LESS;
    bool enable_dynamic_polygon_mode = false;
    bool enable_dynamic_depth_write = false;

    bool blend_enable = false;
    rhi::PipelineBlendFactor src_color_blend_factor = rhi::PipelineBlendFactor::SRC_ALPHA;
    rhi::PipelineBlendFactor dst_color_blend_factor = rhi::PipelineBlendFactor::ONE_MINUS_SRC_ALPHA;
    rhi::PipelineBlendOp color_blend_op = rhi::PipelineBlendOp::ADD;
    rhi::PipelineBlendFactor src_alpha_blend_factor = rhi::PipelineBlendFactor::SRC_ALPHA;
    rhi::PipelineBlendFactor dst_alpha_blend_factor = rhi::PipelineBlendFactor::ONE_MINUS_SRC_ALPHA;
    rhi::PipelineBlendOp alpha_blend_op = rhi::PipelineBlendOp::ADD;

    RenderPipelineDesc() = default;
    RenderPipelineDesc(const RenderPipelineDesc &) = default;
    RenderPipelineDesc &operator=(const RenderPipelineDesc &) = default;

    RenderPipelineDesc &set_shaders(ShaderCompiler &sc, const char *path, const char *vs, const char *fs);
    RenderPipelineDesc &set_input_topology(rhi::PipelineTopology topology);
    RenderPipelineDesc &add_vertex_binding(u32 binding, u32 stride = 0, bool is_instanced = false);
    RenderPipelineDesc &add_vertex_attribute(u32 location, u32 binding, ImageFormat format, u32 offset = 0);
    RenderPipelineDesc &set_polygon_mode(rhi::PipelineFillMode mode);
    RenderPipelineDesc &set_color_format(rhi::ImageFormat format);
    RenderPipelineDesc &set_depth_format(rhi::ImageFormat format);
    RenderPipelineDesc &set_cull_mode(rhi::PipelineCullMode cull_mode, rhi::PipelineTriangleWindingOrder winding);
    RenderPipelineDesc &set_multisampling(rhi::SampleCount count = rhi::SampleCount::Sample1);
    RenderPipelineDesc &set_depth_testing(bool do_test, bool do_write, rhi::PipelineCompareOp op);
    RenderPipelineDesc &set_depth_bias(bool enabled, f32 constant, f32 clamp, f32 slope);
    RenderPipelineDesc &set_depth_clamp(bool enabled);
    RenderPipelineDesc &set_dynamic_polygon_mode(bool enable = true) {
        enable_dynamic_polygon_mode = enable;
        return *this;
    }
    RenderPipelineDesc &set_dynamic_depth_write(bool enable = true) {
        enable_dynamic_depth_write = enable;
        return *this;
    }
    RenderPipelineDesc &set_blending(
        rhi::PipelineBlendFactor src_color = rhi::PipelineBlendFactor::SRC_ALPHA,
        rhi::PipelineBlendFactor dst_color = rhi::PipelineBlendFactor::ONE_MINUS_SRC_ALPHA,
        rhi::PipelineBlendOp color_op = rhi::PipelineBlendOp::ADD,
        rhi::PipelineBlendFactor src_alpha = rhi::PipelineBlendFactor::SRC_ALPHA,
        rhi::PipelineBlendFactor dst_alpha = rhi::PipelineBlendFactor::ONE_MINUS_SRC_ALPHA,
        rhi::PipelineBlendOp alpha_op = rhi::PipelineBlendOp::ADD
    );

    RenderPipeline build();
};

struct ComputePipelineDesc {
    Device *device = nullptr;

    std::string name;

    std::string shader_path;
    std::string cs_entry;
    std::vector<ShaderModule> shader_modules;

    ComputePipelineDesc() = default;
    ComputePipelineDesc(const ComputePipelineDesc &) = default;
    ComputePipelineDesc &operator=(const ComputePipelineDesc &) = default;

    ComputePipelineDesc &set_shader(ShaderCompiler &sc, const char *path, const char *cs);

    ComputePipeline build();
};

} // namespace rhi