DescriptorSet.h

#pragma once
 
#include "spark/render/Buffer.h"
#include "spark/render/DescriptorBinding.h"
#include "spark/render/Export.h"
#include "spark/render/Sampler.h"
#include "spark/render/ShaderStages.h"
 
#include <algorithm>
#include <memory>
 
namespace spark::render
{
    enum class DescriptorType
    {
        // ReSharper disable CppInconsistentNaming
 
        ConstantBuffer = 0x00000001,
 
        StructuredBuffer = 0x00000002,
 
        RWStructuredBuffer = 0x00000012,
 
        Texture = 0x00000003,
 
        RWTexture = 0x00000013,
 
        Sampler = 0x00000004,
 
        InputAttachment = 0x00000005,
 
        Buffer = 0x00000006,
 
        RWBuffer = 0x00000016,
 
        ByteAddressBuffer = 0x00000007,
 
        RWByteAddressBuffer = 0x00000017,
 
        // ReSharper restore CppInconsistentNaming
    };
 
    class SPARK_RENDER_EXPORT IDescriptorLayout : public IBufferLayout
    {
    public:
        ~IDescriptorLayout() noexcept override = default;
 
        [[nodiscard]] virtual DescriptorType descriptorType() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int descriptors() const noexcept = 0;
 
        [[nodiscard]] virtual const ISampler* staticSampler() const noexcept = 0;
    };
 
    class SPARK_RENDER_EXPORT IDescriptorSet
    {
    public:
        virtual ~IDescriptorSet() noexcept = default;
 
        void update(unsigned int binding, const IBuffer& buffer, unsigned int buffer_element = 0, unsigned int elements = 0, unsigned int first_descriptor = 0) const
        {
            genericUpdate(binding, buffer, buffer_element, elements, first_descriptor);
        }
 
        void update(unsigned int binding,
                    const IImage& texture,
                    unsigned int descriptor = 0,
                    unsigned int first_level = 0,
                    unsigned int levels = 0,
                    unsigned int first_layer = 0,
                    unsigned int layers = 0) const { genericUpdate(binding, texture, descriptor, first_level, levels, first_layer, layers); }
 
        void update(unsigned int binding, const ISampler& sampler, unsigned int descriptor = 0) const { genericUpdate(binding, sampler, descriptor); }
 
        void attach(unsigned int binding, const IImage& image) const { genericAttach(binding, image); }
 
    private:
        virtual void genericUpdate(unsigned int binding, const IBuffer& buffer, unsigned int buffer_element, unsigned int elements, unsigned int first_descriptor) const = 0;
        virtual void genericUpdate(unsigned int binding,
                                   const IImage& texture,
                                   unsigned int descriptor,
                                   unsigned int first_level,
                                   unsigned int levels,
                                   unsigned int first_layer,
                                   unsigned int layers) const = 0;
        virtual void genericUpdate(unsigned int binding, const ISampler& sampler, unsigned int descriptor) const = 0;
        virtual void genericAttach(unsigned int binding, const IImage& image) const = 0;
    };
 
    template <typename BufferType, typename ImageType, typename SamplerType>
    class DescriptorSet : public IDescriptorSet
    {
    public:
        using buffer_type = BufferType;
        using image_type = ImageType;
        using sampler_type = SamplerType;
 
    public:
        virtual void update(unsigned int binding,
                            const buffer_type& buffer,
                            unsigned int buffer_element = 0,
                            unsigned int elements = 0,
                            unsigned int first_descriptor = 0) const = 0;
 
        virtual void update(unsigned int binding,
                            const image_type& texture,
                            unsigned int descriptor = 0,
                            unsigned int first_level = 0,
                            unsigned int levels = 0,
                            unsigned int first_layer = 0,
                            unsigned int layers = 0) const = 0;
 
        virtual void update(unsigned int binding, const sampler_type& sampler, unsigned int descriptor = 0) const = 0;
 
        virtual void attach(unsigned int binding, const image_type& image) const = 0;
 
    private:
        void genericUpdate(unsigned binding, const IBuffer& buffer, unsigned buffer_element, unsigned elements, unsigned first_descriptor) const final
        {
            update(binding, dynamic_cast<const buffer_type&>(buffer), buffer_element, elements, first_descriptor);
        }
 
        void genericUpdate(unsigned binding,
                           const IImage& texture,
                           unsigned descriptor,
                           unsigned first_level,
                           unsigned levels,
                           unsigned first_layer,
                           unsigned layers) const override { update(binding, dynamic_cast<const image_type&>(texture), descriptor, first_level, levels, first_layer, layers); }
 
        void genericUpdate(unsigned binding, const ISampler& sampler, unsigned descriptor) const final { update(binding, dynamic_cast<const sampler_type&>(sampler), descriptor); }
        void genericAttach(unsigned binding, const IImage& image) const final { attach(binding, dynamic_cast<const image_type&>(image)); }
    };
 
    class SPARK_RENDER_EXPORT IDescriptorSetLayout
    {
    public:
        virtual ~IDescriptorSetLayout() noexcept = default;
 
        [[nodiscard]] std::vector<const IDescriptorLayout*> descriptors() const noexcept { return genericDescriptors(); }
 
        [[nodiscard]] virtual const IDescriptorLayout* descriptor(unsigned int binding) const = 0;
 
        [[nodiscard]] virtual unsigned int space() const noexcept = 0;
 
        [[nodiscard]] virtual ShaderStage shaderStage() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int uniforms() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int storages() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int images() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int buffers() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int samplers() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int staticSamplers() const noexcept = 0;
 
        [[nodiscard]] virtual unsigned int inputAttachments() const noexcept = 0;
 
        [[nodiscard]] std::unique_ptr<IDescriptorSet> allocate(const std::vector<DescriptorBinding>& bindings = {}) const { return genericAllocate(0, bindings); }
 
        [[nodiscard]] std::unique_ptr<IDescriptorSet> allocate(unsigned int descriptors, const std::vector<DescriptorBinding>& bindings = {}) const
        {
            return genericAllocate(descriptors, bindings);
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> allocateMultiple(unsigned int descriptors_sets,
                                                                                    const std::vector<std::vector<DescriptorBinding>>& bindings = {}) const noexcept
        {
            return allocateMultiple(descriptors_sets, 0, bindings);
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> allocateMultiple(unsigned int descriptors_sets,
                                                                                    const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory) const noexcept
        {
            return allocateMultiple(descriptors_sets, 0, binding_factory);
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> allocateMultiple(unsigned int descriptors_sets,
                                                                                    unsigned int descriptors,
                                                                                    const std::vector<std::vector<DescriptorBinding>>& bindings = {}) const noexcept
        {
            return genericAllocate(descriptors_sets, descriptors, bindings);
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> allocateMultiple(unsigned int descriptors_sets,
                                                                                    unsigned int descriptors,
                                                                                    const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory = {}) const
            noexcept { return genericAllocate(descriptors_sets, descriptors, binding_factory); }
 
        void free(const IDescriptorSet& descriptor_set) const noexcept { genericFree(descriptor_set); }
 
    private:
        [[nodiscard]] virtual std::vector<const IDescriptorLayout*> genericDescriptors() const noexcept = 0;
        [[nodiscard]] virtual std::unique_ptr<IDescriptorSet> genericAllocate(unsigned int descriptors, const std::vector<DescriptorBinding>& bindings) const noexcept = 0;
        [[nodiscard]] virtual std::vector<std::unique_ptr<IDescriptorSet>> genericAllocate(unsigned descriptor_sets,
                                                                                           unsigned descriptors,
                                                                                           const std::vector<std::vector<DescriptorBinding>>& bindings) const = 0;
        [[nodiscard]] virtual std::vector<std::unique_ptr<IDescriptorSet>> genericAllocate(unsigned descriptor_sets,
                                                                                           unsigned descriptors,
                                                                                           const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory) const =
        0;
        virtual void genericFree(const IDescriptorSet& descriptor_set) const noexcept = 0;
    };
 
    template <typename DescriptorLayoutType, typename DescriptorSetType>
    class DescriptorSetLayout : public IDescriptorSetLayout
    {
    public:
        using descriptor_layout_type = DescriptorLayoutType;
        using descriptor_set_type = DescriptorSetType;
 
    public:
        [[nodiscard]] virtual std::vector<const descriptor_layout_type*> descriptors() const noexcept = 0;
 
        [[nodiscard]] const descriptor_layout_type* descriptor(unsigned binding) const override = 0;
 
        [[nodiscard]] virtual std::unique_ptr<descriptor_set_type> allocate(const std::vector<DescriptorBinding>& bindings = {}) const = 0;
 
        [[nodiscard]] virtual std::unique_ptr<descriptor_set_type> allocate(unsigned descriptors, const std::vector<DescriptorBinding>& bindings = { }) const = 0;
 
        [[nodiscard]] virtual std::vector<std::unique_ptr<descriptor_set_type>> allocateMultiple(unsigned descriptor_sets, const std::vector<std::vector<DescriptorBinding>>& bindings = { }) const = 0;
 
        [[nodiscard]] virtual std::vector<std::unique_ptr<descriptor_set_type>> allocateMultiple(unsigned descriptor_sets, const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory) const = 0;
 
        [[nodiscard]] virtual std::vector<std::unique_ptr<descriptor_set_type>> allocateMultiple(unsigned descriptor_sets, unsigned descriptors, const std::vector<std::vector<DescriptorBinding>>& bindings = { }) const = 0;
 
        [[nodiscard]] virtual std::vector<std::unique_ptr<descriptor_set_type>> allocateMultiple(unsigned descriptor_sets, unsigned descriptors, const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory) const = 0;
 
        virtual void free(const descriptor_set_type& descriptor_set) const noexcept = 0;
 
    private:
        [[nodiscard]] std::vector<const IDescriptorLayout*> genericDescriptors() const noexcept final
        {
            auto tmp = descriptors();
            std::vector<const IDescriptorLayout*> result;
            result.reserve(tmp.size());
            std::ranges::transform(tmp, std::back_inserter(result), [](auto& ptr) { return static_cast<const IDescriptorLayout*>(ptr); });
            return result;
        }
 
        [[nodiscard]] std::unique_ptr<IDescriptorSet> genericAllocate(unsigned descriptors, const std::vector<DescriptorBinding>& bindings) const noexcept final
        {
            return allocate(descriptors, bindings);
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> genericAllocate(unsigned descriptor_sets,
                                                                                   unsigned descriptors,
                                                                                   const std::vector<std::vector<DescriptorBinding>>& bindings) const final
        {
            auto tmp = allocateMultiple(descriptor_sets, descriptors, bindings);
            std::vector<std::unique_ptr<IDescriptorSet>> result;
            result.reserve(tmp.size());
            std::ranges::transform(tmp, std::back_inserter(result), [](auto& ptr) { return std::unique_ptr<IDescriptorSet>(ptr.release()); });
            return result;
        }
 
        [[nodiscard]] std::vector<std::unique_ptr<IDescriptorSet>> genericAllocate(unsigned descriptor_sets,
                                                                                   unsigned descriptors,
                                                                                   const std::function<std::vector<DescriptorBinding>(unsigned)>& binding_factory) const final
        {
            auto tmp = allocateMultiple(descriptor_sets, descriptors, binding_factory);
            std::vector<std::unique_ptr<IDescriptorSet>> result;
            result.reserve(tmp.size());
            std::ranges::transform(tmp, std::back_inserter(result), [](auto& ptr) { return std::unique_ptr<IDescriptorSet>(ptr.release()); });
            return result;
        }
 
        void genericFree(const IDescriptorSet& descriptor_set) const noexcept final { free(dynamic_cast<const descriptor_set_type&>(descriptor_set)); }
    };
}
 
template <>
struct std::formatter<spark::render::DescriptorType> : std::formatter<std::string_view>
{
    static constexpr auto parse(format_parse_context& ctx) { return ctx.begin(); }
 
    static constexpr auto format(const spark::render::DescriptorType type, auto& ctx)
    {
        switch (type)
        {
        case spark::render::DescriptorType::ConstantBuffer:
            return std::format_to(ctx.out(), "ConstantBuffer");
        case spark::render::DescriptorType::StructuredBuffer:
            return std::format_to(ctx.out(), "StructuredBuffer");
        case spark::render::DescriptorType::RWStructuredBuffer:
            return std::format_to(ctx.out(), "RWStructuredBuffer");
        case spark::render::DescriptorType::Texture:
            return std::format_to(ctx.out(), "Texture");
        case spark::render::DescriptorType::RWTexture:
            return std::format_to(ctx.out(), "RWTexture");
        case spark::render::DescriptorType::Sampler:
            return std::format_to(ctx.out(), "Sampler");
        case spark::render::DescriptorType::InputAttachment:
            return std::format_to(ctx.out(), "InputAttachment");
        case spark::render::DescriptorType::Buffer:
            return std::format_to(ctx.out(), "Buffer");
        case spark::render::DescriptorType::RWBuffer:
            return std::format_to(ctx.out(), "RWBuffer");
        case spark::render::DescriptorType::ByteAddressBuffer:
            return std::format_to(ctx.out(), "ByteAddressBuffer");
        case spark::render::DescriptorType::RWByteAddressBuffer:
            return std::format_to(ctx.out(), "RWByteAddressBuffer");
        }
        return std::format_to(ctx.out(), "Unknown");
    }
};