From ff24a98e35558397f755dbbc8558e043df39a792 Mon Sep 17 00:00:00 2001
From: Laura Hermanns <lukas.hermanns90@gmail.com>
Date: Sat, 29 Jun 2024 12:30:46 -0400
Subject: [PATCH] [D3D12] Fix crash during descriptor cache flush for PSO
 without descriptor heap.

D3D12CommandContext::PrepareStagingDescriptorHeaps() must be called for each PSO binding, not just for those with resource bindings.
Otherwise, the layout information is not updated and the next call to any of the Flush*StagingDescriptorTables() functions can crash the D3D runtime.
---
 .../Direct3D12/Command/D3D12CommandBuffer.cpp |  7 ++++-
 .../Command/D3D12CommandContext.cpp           | 30 +++++++++++--------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/sources/Renderer/Direct3D12/Command/D3D12CommandBuffer.cpp b/sources/Renderer/Direct3D12/Command/D3D12CommandBuffer.cpp
index c6509746ee..01e45f9d03 100644
--- a/sources/Renderer/Direct3D12/Command/D3D12CommandBuffer.cpp
+++ b/sources/Renderer/Direct3D12/Command/D3D12CommandBuffer.cpp
@@ -703,14 +703,19 @@ void D3D12CommandBuffer::SetPipelineState(PipelineState& pipelineState)
     /* Keep reference to pipeline layout */
     boundPipelineLayout_ = pipelineStateD3D.GetPipelineLayout();
 
-    /* Prepare staging descriptor heaps for bound pipeline layout */
     if (boundPipelineLayout_ != nullptr)
     {
+        /* Prepare staging descriptor heaps for bound pipeline layout */
         commandContext_.PrepareStagingDescriptorHeaps(
             boundPipelineLayout_->GetDescriptorHeapSetLayout(),
             boundPipelineLayout_->GetRootParameterIndices()
         );
     }
+    else
+    {
+        /* Reset staging descriptor layout to avoid undefined behavior in next Flush*StagingDescriptorTables() call */
+        commandContext_.PrepareStagingDescriptorHeaps({}, {});
+    }
 }
 
 void D3D12CommandBuffer::SetBlendFactor(const float color[4])
diff --git a/sources/Renderer/Direct3D12/Command/D3D12CommandContext.cpp b/sources/Renderer/Direct3D12/Command/D3D12CommandContext.cpp
index 59e19309d9..911d4411ce 100644
--- a/sources/Renderer/Direct3D12/Command/D3D12CommandContext.cpp
+++ b/sources/Renderer/Direct3D12/Command/D3D12CommandContext.cpp
@@ -368,19 +368,25 @@ void D3D12CommandContext::PrepareStagingDescriptorHeaps(
     stagingDescriptorSetLayout_ = layout;
     stagingDescriptorIndices_   = indices;
 
-    /* Bind shader-visible descriptor heaps */
-    ID3D12DescriptorHeap* const stagingDescriptorHeaps[2] =
+    if (stagingDescriptorSetLayout_.numHeapResourceViews > 0 ||
+        stagingDescriptorSetLayout_.numHeapSamplers      > 0 ||
+        stagingDescriptorSetLayout_.numResourceViews     > 0 ||
+        stagingDescriptorSetLayout_.numSamplers          > 0)
     {
-        stagingDescriptorPools_[currentAllocatorIndex_][0].GetDescriptorHeap(),
-        stagingDescriptorPools_[currentAllocatorIndex_][1].GetDescriptorHeap()
-    };
-    SetDescriptorHeaps(2, stagingDescriptorHeaps);
-
-    /* Reset descriptor cache for dynamic descriptors */
-    descriptorCaches_[currentAllocatorIndex_].Reset(
-        stagingDescriptorSetLayout_.numResourceViews,
-        stagingDescriptorSetLayout_.numSamplers
-    );
+        /* Bind shader-visible descriptor heaps */
+        ID3D12DescriptorHeap* const stagingDescriptorHeaps[2] =
+        {
+            stagingDescriptorPools_[currentAllocatorIndex_][0].GetDescriptorHeap(),
+            stagingDescriptorPools_[currentAllocatorIndex_][1].GetDescriptorHeap()
+        };
+        SetDescriptorHeaps(2, stagingDescriptorHeaps);
+
+        /* Reset descriptor cache for dynamic descriptors */
+        descriptorCaches_[currentAllocatorIndex_].Reset(
+            stagingDescriptorSetLayout_.numResourceViews,
+            stagingDescriptorSetLayout_.numSamplers
+        );
+    }
 }
 
 void D3D12CommandContext::SetGraphicsConstant(UINT parameterIndex, D3D12Constant value, UINT offset)