PSO initialization: updated shader stage initialization

2 files changed, 21 insertions, 4 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
index 198acffd..f604f6ab 100644
--- a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
@@ -128,7 +128,7 @@ private:
     struct ShaderStageInfo
     {
         ShaderStageInfo() {}
-        ShaderStageInfo(SHADER_TYPE _Type, ShaderD3D12Impl* _pShader);
+        ShaderStageInfo(ShaderD3D12Impl* _pShader);
 
         void   Append(ShaderD3D12Impl* pShader);
         size_t Count() const;
diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index 187bab5a..afdeb741 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -336,14 +336,31 @@ TBindingMapPerStage ExtractResourceBindingMap(const RootSignatureBuilder&
 } // namespace
 
 
-PipelineStateD3D12Impl::ShaderStageInfo::ShaderStageInfo(SHADER_TYPE _Type, ShaderD3D12Impl* _pShader) :
-    Type{_Type},
+PipelineStateD3D12Impl::ShaderStageInfo::ShaderStageInfo(ShaderD3D12Impl* _pShader) :
+    Type{_pShader->GetDesc().ShaderType},
     Shaders{_pShader}
 {
 }
 
 void PipelineStateD3D12Impl::ShaderStageInfo::Append(ShaderD3D12Impl* pShader)
 {
+    VERIFY_EXPR(pShader != nullptr);
+    VERIFY(std::find(Shaders.begin(), Shaders.end(), pShader) == Shaders.end(),
+           "Shader '", pShader->GetDesc().Name, "' already exists in the stage. Shaders must be deduplicated.");
+
+    const auto NewShaderType = pShader->GetDesc().ShaderType;
+    if (Type == SHADER_TYPE_UNKNOWN)
+    {
+        VERIFY_EXPR(Shaders.empty());
+        Type = NewShaderType;
+    }
+    else
+    {
+        VERIFY(Type == NewShaderType, "The type (", GetShaderTypeLiteralName(NewShaderType),
+               ") of shader '", pShader->GetDesc().Name, "' being added to the stage is incosistent with the stage type (",
+               GetShaderTypeLiteralName(Type), ").");
+    }
+
     Shaders.push_back(pShader);
 }
 
@@ -706,7 +723,7 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
         if (FAILED(hr))
             LOG_ERROR_AND_THROW("Failed to create ray tracing state object");
 
-        GetShaderIdentifiers(m_pd3d12PSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex, m_pRayTracingPipelineData->Shaders);
+        GetShaderIdentifiers(m_pd3d12PSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex, m_pRayTracingPipelineData->ShaderHandles);
 
         if (*m_Desc.Name != 0)
         {