LightingShaderResourceManager.h

#pragma once
 
// Standard.
#include <mutex>
#include <memory>
#include <functional>
#include <unordered_map>
#include <unordered_set>
 
// Custom.
#include "render/general/resources/UploadBuffer.h"
#include "render/general/resources/frame/FrameResourcesManager.h"
#if defined(WIN32)
#include "render/directx/resources/DirectXResource.h"
#include "shader/hlsl/RootSignatureGenerator.h"
#include "render/directx/pipeline/DirectXPso.h"
#endif
#include "shader/general/resources/ShaderLightArray.h"
 
// External.
#include "vulkan/vulkan_core.h"
#if defined(WIN32)
#include "directx/d3dx12.h"
#include <wrl.h>
#endif
 
namespace ne {
#if defined(WIN32)
    using namespace Microsoft::WRL;
#endif
 
    class Renderer;
    class Pipeline;
    class ComputeShaderInterface;
    class VulkanPipeline;
 
    struct ShadowPassLightShaderInfo {
        alignas(iVkVec4Alignment) glm::mat4 viewProjectionMatrix = glm::identity<glm::mat4>();
 
        alignas(iVkVec4Alignment) glm::vec4 position = glm::vec4(0.0F, 0.0F, 0.0F, 1.0F);
    };
 
    class LightingShaderResourceManager {
        // Only renderer is allowed to create and update resources of this manager.
        friend class Renderer;
 
    public:
        struct GeneralLightingShaderData {
            alignas(iVkVec4Alignment) glm::vec4 ambientLight = glm::vec4(0.0F, 0.0F, 0.0F, 1.0F);
 
            alignas(iVkScalarAlignment) unsigned int iPointLightCountInCameraFrustum = 0;
 
            alignas(iVkScalarAlignment) unsigned int iDirectionalLightCount = 0;
 
            alignas(iVkScalarAlignment) unsigned int iSpotLightCountInCameraFrustum = 0;
        };
 
        struct GpuData {
            GpuData() = default;
 
            std::array<std::unique_ptr<UploadBuffer>, FrameResourcesManager::getFrameResourcesCount()>
                vGeneralDataGpuResources;
 
            GeneralLightingShaderData generalData;
        };
 
        LightingShaderResourceManager() = delete;
 
        ~LightingShaderResourceManager();
 
        static std::string getGeneralLightingDataShaderResourceName();
 
        static std::string getPointLightsShaderResourceName();
 
        static std::string getDirectionalLightsShaderResourceName();
 
        static std::string getSpotlightsShaderResourceName();
 
        static std::string getPointLightsInCameraFrustumIndicesShaderResourceName();
 
        static std::string getSpotlightsInCameraFrustumIndicesShaderResourceName();
 
        static std::string getShadowPassLightInfoArrayShaderResourceName();
 
        static std::string getPointLightWorldLocationsShaderResourceName();
 
#if defined(WIN32)
        inline void setShadowPassLightInfoViewToCommandList(
            DirectXPso* pPso,
            const ComPtr<ID3D12GraphicsCommandList>& pCommandList,
            size_t iCurrentFrameResourceIndex) {
            // Lock internal resources.
            std::scoped_lock gpuDataGuard(mtxGpuData.first);
 
            auto& pipelineData = pPso->getInternalResources()->second;
 
#if defined(DEBUG)
            // Self check: make sure this resource is actually being used in the PSO.
            if (pipelineData.rootParameterIndices.find(sShadowPassLightInfoShaderResourceName) ==
                pipelineData.rootParameterIndices.end()) [[unlikely]] {
                Error error(std::format(
                    "shader resource \"{}\" is not used in the shaders of the specified PSO \"{}\" "
                    "but you are attempting to set this resource to a command list",
                    sShadowPassLightInfoShaderResourceName,
                    pPso->getPipelineIdentifier()));
                error.showError();
                throw std::runtime_error(error.getFullErrorMessage());
            }
#endif
 
            // Bind array.
            setLightingArrayViewToCommandList(
                pPso,
                pCommandList,
                iCurrentFrameResourceIndex,
                lightArrays.pShadowPassLightInfoArray,
                sShadowPassLightInfoShaderResourceName,
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::SHADOW_PASS_LIGHT_INFO)]);
        }
#endif
 
#if defined(WIN32)
        static inline void setLightingArrayViewToCommandList(
            DirectXPso* pPso,
            const ComPtr<ID3D12GraphicsCommandList>& pCommandList,
            size_t iCurrentFrameResourceIndex,
            const std::unique_ptr<ShaderLightArray>& pArray,
            const std::string& sArrayNameInShaders,
            UINT iArrayRootParameterIndex) {
            std::scoped_lock lightsGuard(pArray->mtxResources.first);
 
#if defined(DEBUG)
            // Self check: make sure this resource is actually being used in the PSO.
            if (pPso->getInternalResources()->second.rootParameterIndices.find(sArrayNameInShaders) ==
                pPso->getInternalResources()->second.rootParameterIndices.end()) [[unlikely]] {
                Error error(std::format(
                    "shader resource \"{}\" is not used in the shaders of the specified PSO \"{}\" "
                    "but you are attempting to set this resource to a command list",
                    sArrayNameInShaders,
                    pPso->getPipelineIdentifier()));
                error.showError();
                throw std::runtime_error(error.getFullErrorMessage());
            }
#endif
 
            // Bind lights array
            // (resource is guaranteed to be not `nullptr`, see field docs).
            pCommandList->SetGraphicsRootShaderResourceView(
                iArrayRootParameterIndex,
                reinterpret_cast<DirectXResource*>(
                    pArray->mtxResources.second.vGpuArrayLightDataResources[iCurrentFrameResourceIndex]
                        ->getInternalResource())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
        }
#endif
 
        ShaderLightArray* getPointLightDataArray() const;
 
        ShaderLightArray* getDirectionalLightDataArray() const;
 
        ShaderLightArray* getSpotlightDataArray() const;
 
        ShaderLightArray* getShadowPassLightInfoArray() const;
 
        [[nodiscard]] std::optional<Error> bindDescriptorsToRecreatedPipelineResources();
 
        [[nodiscard]] std::optional<Error> updateDescriptorsForPipelineResource(Pipeline* pPipeline);
 
        std::pair<std::recursive_mutex, GpuData>* getInternalResources();
 
#if defined(WIN32)
        inline void setResourceViewToCommandList(
            DirectXPso* pPso,
            const ComPtr<ID3D12GraphicsCommandList>& pCommandList,
            size_t iCurrentFrameResourceIndex) {
            // Lock internal resources.
            std::scoped_lock gpuDataGuard(mtxGpuData.first);
 
            auto& pipelineData = pPso->getInternalResources()->second;
 
#if defined(DEBUG)
            // Self check: make sure this resource is actually being used in the PSO.
            if (pipelineData.rootParameterIndices.find(sGeneralLightingDataShaderResourceName) ==
                pipelineData.rootParameterIndices.end()) [[unlikely]] {
                Error error(std::format(
                    "shader resource \"{}\" is not used in the shaders of the specified PSO \"{}\" "
                    "but you are attempting to set this resource to a command list",
                    sGeneralLightingDataShaderResourceName,
                    pPso->getPipelineIdentifier()));
                error.showError();
                throw std::runtime_error(error.getFullErrorMessage());
            }
#endif
 
            // Bind general lighting resources buffer.
            pCommandList->SetGraphicsRootConstantBufferView(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::GENERAL_LIGHTING)],
                reinterpret_cast<DirectXResource*>(
                    mtxGpuData.second.vGeneralDataGpuResources[iCurrentFrameResourceIndex]
                        ->getInternalResource())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
 
            // Bind point lights array.
            setLightingArrayViewToCommandList(
                pPso,
                pCommandList,
                iCurrentFrameResourceIndex,
                lightArrays.pPointLightDataArray,
                sPointLightsShaderResourceName,
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::POINT_LIGHTS)]);
 
            // Bind directional lights array.
            setLightingArrayViewToCommandList(
                pPso,
                pCommandList,
                iCurrentFrameResourceIndex,
                lightArrays.pDirectionalLightDataArray,
                sDirectionalLightsShaderResourceName,
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::DIRECTIONAL_LIGHTS)]);
 
            // Bind spotlights array.
            setLightingArrayViewToCommandList(
                pPso,
                pCommandList,
                iCurrentFrameResourceIndex,
                lightArrays.pSpotlightDataArray,
                sSpotlightsShaderResourceName,
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::SPOT_LIGHTS)]);
 
#if defined(DEBUG)
            static_assert(sizeof(LightArrays) == 32, "consider adding new arrays here"); // NOLINT
#endif
        }
 
        inline void setOpaqueLightGridResourcesViewToCommandList(
            DirectXPso* pPso, const ComPtr<ID3D12GraphicsCommandList>& pCommandList) const {
            auto& pipelineData = pPso->getInternalResources()->second;
 
            // Bind point light index list.
            pCommandList->SetGraphicsRootUnorderedAccessView(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_POINT_LIGHT_INDEX_LIST)],
                reinterpret_cast<DirectXResource*>(
                    lightCullingComputeShaderData.resources.pOpaquePointLightIndexList.get())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
 
            // Bind spotlight index list.
            pCommandList->SetGraphicsRootUnorderedAccessView(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_SPOT_LIGHT_INDEX_LIST)],
                reinterpret_cast<DirectXResource*>(
                    lightCullingComputeShaderData.resources.pOpaqueSpotLightIndexList.get())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
 
            // Bind point light grid.
            pCommandList->SetGraphicsRootDescriptorTable(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_POINT_LIGHT_GRID)],
                *reinterpret_cast<DirectXResource*>(
                     lightCullingComputeShaderData.resources.pOpaquePointLightGrid.get())
                     ->getBindedDescriptorGpuHandle(DirectXDescriptorType::UAV));
 
            // Bind spotlight light grid.
            pCommandList->SetGraphicsRootDescriptorTable(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_SPOT_LIGHT_GRID)],
                *reinterpret_cast<DirectXResource*>(
                     lightCullingComputeShaderData.resources.pOpaqueSpotLightGrid.get())
                     ->getBindedDescriptorGpuHandle(DirectXDescriptorType::UAV));
        }
 
        inline void setTransparentLightGridResourcesViewToCommandList(
            DirectXPso* pPso, const ComPtr<ID3D12GraphicsCommandList>& pCommandList) const {
            auto& pipelineData = pPso->getInternalResources()->second;
 
            // Bind point light index list.
            pCommandList->SetGraphicsRootUnorderedAccessView(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_POINT_LIGHT_INDEX_LIST)],
                reinterpret_cast<DirectXResource*>(
                    lightCullingComputeShaderData.resources.pTransparentPointLightIndexList.get())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
 
            // Bind spotlight index list.
            pCommandList->SetGraphicsRootUnorderedAccessView(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_SPOT_LIGHT_INDEX_LIST)],
                reinterpret_cast<DirectXResource*>(
                    lightCullingComputeShaderData.resources.pTransparentSpotLightIndexList.get())
                    ->getInternalResource()
                    ->GetGPUVirtualAddress());
 
            // Bind point light grid.
            pCommandList->SetGraphicsRootDescriptorTable(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_POINT_LIGHT_GRID)],
                *reinterpret_cast<DirectXResource*>(
                     lightCullingComputeShaderData.resources.pTransparentPointLightGrid.get())
                     ->getBindedDescriptorGpuHandle(DirectXDescriptorType::UAV));
 
            // Bind spotlight light grid.
            pCommandList->SetGraphicsRootDescriptorTable(
                pipelineData.vSpecialRootParameterIndices[static_cast<size_t>(
                    SpecialRootParameterSlot::LIGHT_CULLING_SPOT_LIGHT_GRID)],
                *reinterpret_cast<DirectXResource*>(
                     lightCullingComputeShaderData.resources.pTransparentSpotLightGrid.get())
                     ->getBindedDescriptorGpuHandle(DirectXDescriptorType::UAV));
        }
 
#endif
 
    private:
        struct LightArrays {
            std::unique_ptr<ShaderLightArray> pPointLightDataArray;
 
            std::unique_ptr<ShaderLightArray> pDirectionalLightDataArray;
 
            std::unique_ptr<ShaderLightArray> pSpotlightDataArray;
 
            std::unique_ptr<ShaderLightArray> pShadowPassLightInfoArray;
        };
 
        struct ComputeShaderData {
            struct Plane {
                glm::vec3 normal;
 
                float distanceFromOrigin;
            };
 
            struct Frustum {
                Plane planes[4];
            };
 
            struct FrustumGridComputeShader {
                struct ComputeInfo {
                    alignas(iVkScalarAlignment) unsigned int iThreadGroupCountX = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iThreadGroupCountY = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iTileCountX = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iTileCountY = 0;
                };
 
                struct ScreenToViewData {
                    alignas(iVkMat4Alignment) glm::mat4 inverseProjectionMatrix = glm::identity<glm::mat4>();
 
                    alignas(iVkScalarAlignment) unsigned int iRenderTargetWidth = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iRenderTargetHeight = 0;
                };
 
                struct ShaderResources {
                    std::unique_ptr<UploadBuffer> pComputeInfo;
 
                    std::unique_ptr<UploadBuffer> pScreenToViewData;
 
                    std::unique_ptr<GpuResource> pCalculatedFrustums;
                };
 
                struct ComputeShader {
                    [[nodiscard]] std::optional<Error> initialize(Renderer* pRenderer);
 
                    [[nodiscard]] std::optional<Error> updateDataAndSubmitShader(
                        Renderer* pRenderer,
                        const std::pair<unsigned int, unsigned int>& renderTargetSize,
                        const glm::mat4& inverseProjectionMatrix);
 
                    std::unique_ptr<ComputeShaderInterface> pComputeInterface;
 
                    ShaderResources resources;
 
                    unsigned int iLastUpdateTileCountX = 0;
 
                    unsigned int iLastUpdateTileCountY = 0;
 
                    bool bIsInitialized = false;
 
                    static inline const auto sComputeInfoShaderResourceName = "computeInfo";
 
                    static inline const auto sScreenToViewDataShaderResourceName = "screenToViewData";
 
                    static inline const auto sCalculatedFrustumsShaderResourceName = "calculatedFrustums";
                };
            };
 
            struct LightCullingComputeShader {
                struct GlobalCountersIntoLightIndexList {
                    alignas(iVkScalarAlignment) unsigned int iPointLightListOpaque = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iSpotlightListOpaque = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iPointLightListTransparent = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iSpotlightListTransparent = 0;
                };
 
                struct ThreadGroupCount {
                    alignas(iVkScalarAlignment) unsigned int iThreadGroupCountX = 0;
 
                    alignas(iVkScalarAlignment) unsigned int iThreadGroupCountY = 0;
                };
 
                struct ShaderResources {
                    std::unique_ptr<UploadBuffer> pThreadGroupCount;
 
                    std::unique_ptr<GpuResource> pGlobalCountersIntoLightIndexList;
 
                    std::unique_ptr<GpuResource> pOpaquePointLightIndexList;
 
                    std::unique_ptr<GpuResource> pOpaqueSpotLightIndexList;
 
                    std::unique_ptr<GpuResource> pTransparentPointLightIndexList;
 
                    std::unique_ptr<GpuResource> pTransparentSpotLightIndexList;
 
                    std::unique_ptr<GpuResource> pOpaquePointLightGrid;
 
                    std::unique_ptr<GpuResource> pOpaqueSpotLightGrid;
 
                    std::unique_ptr<GpuResource> pTransparentPointLightGrid;
 
                    std::unique_ptr<GpuResource> pTransparentSpotLightGrid;
 
                    size_t iLightGridTileCountX = 0;
 
                    size_t iLightGridTileCountY = 0;
                };
 
                struct ComputeShader {
                    [[nodiscard]] std::optional<Error> initialize(
                        Renderer* pRenderer,
                        const FrustumGridComputeShader::ComputeShader& frustumGridShader);
 
                    [[nodiscard]] std::variant<bool, Error> createLightIndexListsAndLightGrid(
                        Renderer* pRenderer, size_t iTileCountX, size_t iTileCountY);
 
                    [[nodiscard]] std::optional<Error> queueExecutionForNextFrame(
                        Renderer* pRenderer,
                        FrameResource* pCurrentFrameResource,
                        size_t iCurrentFrameResourceIndex,
                        GpuResource* pGeneralLightingData,
                        GpuResource* pPointLightArray,
                        GpuResource* pSpotlightArray,
                        GpuResource* pNonCulledPointLightsIndicesArray,
                        GpuResource* pNonCulledSpotlightsIndicesArray) const;
 
                    std::unique_ptr<ComputeShaderInterface> pComputeInterface;
 
                    ShaderResources resources;
 
                    ThreadGroupCount threadGroupCount;
 
                    bool bIsInitialized = false;
 
                    static inline const auto sDepthTextureShaderResourceName = "depthTexture";
 
                    static inline const auto sThreadGroupCountShaderResourceName = "threadGroupCount";
 
                    static inline const auto sGlobalCountersIntoLightIndexListShaderResourceName =
                        "globalCountersIntoLightIndexList";
                };
            };
        };
 
        static std::unique_ptr<LightingShaderResourceManager> create(Renderer* pRenderer);
 
        [[nodiscard]] static std::optional<Error> rebindBufferResourceToPipeline(
            VulkanPipeline* pVulkanPipeline,
            VkDevice pLogicalDevice,
            const std::string& sShaderResourceName,
            VkDeviceSize iResourceSize,
            VkDescriptorType descriptorType,
            std::array<VkBuffer, FrameResourcesManager::getFrameResourcesCount()> vBuffersToBind);
 
        [[nodiscard]] static std::optional<Error> rebindImageResourceToPipeline(
            VulkanPipeline* pVulkanPipeline,
            VkDevice pLogicalDevice,
            const std::string& sShaderResourceName,
            VkDescriptorType descriptorType,
            VkImageLayout imageLayout,
            VkSampler pSampler,
            std::array<VkImageView, FrameResourcesManager::getFrameResourcesCount()> vImagesToBind);
 
        LightingShaderResourceManager(Renderer* pRenderer);
 
        [[nodiscard]] std::optional<Error> recalculateLightTileFrustums(
            const std::pair<unsigned int, unsigned int>& renderTargetSize,
            const glm::mat4& inverseProjectionMatrix);
 
        void onEngineShadersCompiled();
 
        void setAmbientLight(const glm::vec3& ambientLight);
 
        void updateResources(FrameResource* pCurrentFrameResource, size_t iCurrentFrameResourceIndex);
 
        void onPointLightArraySizeChanged(size_t iNewSize);
 
        void onPointLightsInFrustumCulled(size_t iCurrentFrameResourceIndex);
 
        void onDirectionalLightArraySizeChanged(size_t iNewSize);
 
        void onSpotlightArraySizeChanged(size_t iNewSize);
 
        void onSpotlightsInFrustumCulled(size_t iCurrentFrameResourceIndex);
 
        void copyDataToGpu(size_t iCurrentFrameResourceIndex);
 
        [[nodiscard]] std::optional<Error> rebindGpuDataToAllPipelines();
 
        [[nodiscard]] std::optional<Error> rebindGpuDataToPipeline(Pipeline* pPipeline);
 
        LightArrays lightArrays;
 
        std::pair<std::recursive_mutex, GpuData> mtxGpuData;
 
        ComputeShaderData::FrustumGridComputeShader::ComputeShader frustumGridComputeShaderData;
 
        ComputeShaderData::LightCullingComputeShader::ComputeShader lightCullingComputeShaderData;
 
        std::unique_ptr<ComputeShaderInterface> pPrepareLightCullingComputeInterface;
 
        Renderer* pRenderer = nullptr;
 
        bool bEngineShadersCompiled = false;
 
        static inline const std::string sGeneralLightingDataShaderResourceName = "generalLightingData";
 
        static inline const std::string sPointLightsShaderResourceName = "pointLights";
 
        static inline const std::string sDirectionalLightsShaderResourceName = "directionalLights";
 
        static inline const std::string sSpotlightsShaderResourceName = "spotlights";
 
        static inline const std::string sPointLightsInCameraFrustumIndicesShaderResourceName =
            "pointLightsInCameraFrustumIndices";
 
        static inline const std::string sSpotlightsInCameraFrustumIndicesShaderResourceName =
            "spotlightsInCameraFrustumIndices";
 
        static inline const std::string sShadowPassLightInfoShaderResourceName = "shadowPassLightInfo";
 
        static constexpr auto generalLightingDataDescriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    };
}