DescriptorsHandler.hpp

#pragma once
 
#include "PixelBullet/Graphics/Buffers/PushHandler.hpp"
#include "PixelBullet/Graphics/Buffers/StorageHandler.hpp"
#include "PixelBullet/Graphics/Buffers/UniformHandler.hpp"
#include "PixelBullet/Graphics/Descriptors/DescriptorSet.hpp"
#include "PixelBullet/Graphics/Pipelines/Shader.hpp"
 
#include <iomanip>
#include <memory>
 
namespace PixelBullet
{
    class DescriptorsHandler
    {
    public:
        DescriptorsHandler() = default;
        explicit DescriptorsHandler(const Pipeline& pipeline);
 
        template <typename T>
        void Push(const std::string& descriptorName, const T& descriptor,
                  const std::optional<OffsetSize>& offsetSize = std::nullopt)
        {
            if (!m_Shader)
            {
                return;
            }
 
            auto it = m_Descriptors.find(descriptorName);
 
            if (it != m_Descriptors.end())
            {
                if (it->second.m_Descriptor == std::to_address(descriptor) && it->second.m_OffsetSize == offsetSize)
                {
                    return;
                }
 
                m_Descriptors.erase(it);
            }
 
            if (!std::to_address(descriptor))
            {
                return;
            }
 
            auto location = m_Shader->GetDescriptorLocation(descriptorName);
 
            if (!location)
            {
#ifdef PB_DEBUG
                if (m_Shader->ReportedNotFound(descriptorName, true))
                {
                    Log::Error("Could not find descriptor '{}' in shader '{}'", descriptorName,
                               m_Shader->GetName().c_str());
                }
#endif
                return;
            }
 
            auto descriptorType = m_Shader->GetDescriptorType(*location);
 
            if (!descriptorType)
            {
#ifdef PB_DEBUG
                if (m_Shader->ReportedNotFound(descriptorName, true))
                {
                    Log::Error("Could not find descriptor '{}' at location {} in shader '{}'", descriptorName,
                               *location, m_Shader->GetName().c_str());
                }
#endif
                return;
            }
 
            auto writeDescriptor =
                std::to_address(descriptor)->GetWriteDescriptor(*location, *descriptorType, offsetSize);
            m_Descriptors.emplace(descriptorName, DescriptorValue{ std::to_address(descriptor),
                                                                   std::move(writeDescriptor), offsetSize, *location });
            m_Changed = true;
        }
 
        template <typename T>
        void Push(const std::string& descriptorName, const T& descriptor, WriteDescriptorSet writeDescriptorSet)
        {
            if (!m_Shader)
            {
                return;
            }
 
            if (auto it = m_Descriptors.find(descriptorName); it != m_Descriptors.end())
            {
                m_Descriptors.erase(it);
            }
 
            auto location = m_Shader->GetDescriptorLocation(descriptorName);
 
            m_Descriptors.emplace(
                descriptorName,
                DescriptorValue{ std::to_address(descriptor), std::move(writeDescriptorSet), std::nullopt, *location });
            m_Changed = true;
        }
 
        void Push(const std::string& descriptorName, UniformHandler& uniformHandler,
                  const std::optional<OffsetSize>& offsetSize = std::nullopt);
        void Push(const std::string& descriptorName, StorageHandler& storageHandler,
                  const std::optional<OffsetSize>& offsetSize = std::nullopt);
        void Push(const std::string& descriptorName, PushHandler& pushHandler,
                  const std::optional<OffsetSize>& offsetSize = std::nullopt);
 
        bool Update(const Pipeline& pipeline);
 
        void BindDescriptor(const CommandBuffer& commandBuffer, const Pipeline& pipeline);
 
        const DescriptorSet* GetDescriptorSet() const
        {
            return m_DescriptorSet.get();
        }
 
    private:
        class DescriptorValue
        {
        public:
            const Descriptor* m_Descriptor;
            WriteDescriptorSet m_WriteDescriptor;
            std::optional<OffsetSize> m_OffsetSize;
            uint32_t m_Location;
        };
 
        const Shader* m_Shader = nullptr;
        bool m_PushDescriptors = false;
        std::unique_ptr<DescriptorSet> m_DescriptorSet;
 
        std::map<std::string, DescriptorValue> m_Descriptors;
        std::vector<VkWriteDescriptorSet> m_WriteDescriptorSets;
        bool m_Changed = false;
    };
} // namespace PixelBullet