RootSignatureGenerator.h

#pragma once
 
// Standard.
#include <variant>
#include <optional>
#include <unordered_map>
#include <stdexcept>
#include <array>
#include <set>
 
// Custom.
#include "misc/Error.h"
#include "SpecialRootParameterSlot.hpp"
 
// External.
#include "directx/d3dx12.h"
#include <DirectXShaderCompiler/inc/d3d12shader.h>
 
// OS.
#include <wrl/client.h>
 
namespace ne {
    class HlslShader;
    class DirectXRenderer;
 
    using namespace Microsoft::WRL;
 
    enum class SamplerType {
        BASIC,     //< Usual `SamplerState` type in HLSL.
        COMPARISON //< `SamplerComparisonState` type in HLSL.
    };
 
    class RootSignatureGenerator {
    public:
        RootSignatureGenerator() = delete;
        RootSignatureGenerator(const RootSignatureGenerator&) = delete;
        RootSignatureGenerator& operator=(const RootSignatureGenerator&) = delete;
 
        class RootParameter {
        public:
            enum class Type { CONSTANTS, CBV, SRV, UAV };
 
            RootParameter() = default;
 
            RootParameter(UINT iBindPoint, UINT iSpace, Type type, bool bIsTable = false, UINT iCount = 1);
 
            CD3DX12_ROOT_PARAMETER generateSingleDescriptorDescription() const;
 
            CD3DX12_DESCRIPTOR_RANGE generateTableRange() const;
 
            D3D12_SHADER_VISIBILITY getVisibility() const;
 
            bool isTable() const;
 
        private:
            UINT iBindPoint = 0;
 
            UINT iSpace = 0;
 
            Type type;
 
            UINT iCount = 0;
 
            bool bIsTable = false;
 
            D3D12_SHADER_VISIBILITY visibility = D3D12_SHADER_VISIBILITY_ALL;
        };
 
        struct CollectedInfo {
            std::set<SamplerType> staticSamplers;
 
            std::vector<RootParameter> vRootParameters;
 
            std::unordered_map<std::string, std::pair<UINT, RootParameter>> rootParameterIndices;
 
            std::unordered_map<std::string, size_t> rootConstantOffsets;
        };
 
        struct Generated {
            ComPtr<ID3D12RootSignature> pRootSignature;
 
            std::unordered_map<std::string, UINT> rootParameterIndices;
 
            std::array<UINT, static_cast<unsigned int>(SpecialRootParameterSlot::SIZE)>
                vSpecialRootParameterIndices;
 
            std::unordered_map<std::string, size_t> rootConstantOffsets;
        };
 
        static std::variant<CollectedInfo, Error> collectInfoFromReflection(
            ID3D12Device* pDevice, const ComPtr<ID3D12ShaderReflection>& pShaderReflection);
 
        static std::variant<Generated, Error>
        generateGraphics(DirectXRenderer* pRenderer, HlslShader* pVertexShader, HlslShader* pPixelShader);
 
        static std::variant<Generated, Error>
        generateCompute(DirectXRenderer* pRenderer, HlslShader* pComputeShader);
 
    private:
        static std::variant<SamplerType, Error>
        findStaticSamplerForSamplerResource(const D3D12_SHADER_INPUT_BIND_DESC& samplerResourceDescription);
 
        static void addSpecialResourceRootParametersIfUsed(
            std::unordered_map<std::string, std::pair<UINT, RootSignatureGenerator::RootParameter>>&
                shaderRootParameterIndices,
            std::vector<CD3DX12_ROOT_PARAMETER>& vRootParameters,
            std::vector<CD3DX12_DESCRIPTOR_RANGE>& vTableRanges,
            std::set<std::string>& addedRootParameterNames,
            std::unordered_map<std::string, UINT>& rootParameterIndices,
            std::array<UINT, static_cast<unsigned int>(SpecialRootParameterSlot::SIZE)>&
                vSpecialRootParameterIndices);
 
        static std::optional<Error> addUniquePairResourceNameRootParameterIndex(
            std::unordered_map<std::string, std::pair<UINT, RootParameter>>& mapToAddTo,
            const std::string& sResourceName,
            UINT iRootParameterIndex,
            const RootParameter& parameter);
 
        static std::optional<Error> addCbufferRootParameter(
            std::vector<RootParameter>& vRootParameters,
            std::unordered_map<std::string, std::pair<UINT, RootParameter>>& rootParameterIndices,
            const D3D12_SHADER_INPUT_BIND_DESC& resourceDescription);
 
        static std::optional<Error> addTexture2DRootParameter(
            std::vector<RootParameter>& vRootParameters,
            std::unordered_map<std::string, std::pair<UINT, RootParameter>>& rootParameterIndices,
            const D3D12_SHADER_INPUT_BIND_DESC& resourceDescription,
            bool bIsReadWrite);
 
        static std::optional<Error> addStructuredBufferRootParameter(
            std::vector<RootParameter>& vRootParameters,
            std::unordered_map<std::string, std::pair<UINT, RootParameter>>& rootParameterIndices,
            const D3D12_SHADER_INPUT_BIND_DESC& resourceDescription,
            bool bIsReadWrite);
 
        static std::variant<bool, Error> processRootConstantsIfFound(
            const ComPtr<ID3D12ShaderReflection>& pShaderReflection,
            const D3D12_SHADER_INPUT_BIND_DESC& resourceDescription,
            std::unordered_map<std::string, size_t>& rootConstantOffsets,
            std::vector<RootParameter>& vRootParameters,
            std::unordered_map<std::string, std::pair<UINT, RootParameter>>& rootParameterIndices);
 
        static inline const std::string sRootConstantsVariableName = "constants";
 
        static inline const std::string sRootConstantsTypeName = "RootConstants";
    };
} // namespace ne