From fa79938e8aa115d48b028f82685bff1756a68822 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Tue, 1 Aug 2023 16:05:49 +0100
Subject: [PATCH] Ensure image layout is up to date mid-command buffer in pixel
 history

* Previously we were obtaining the image layout from before the current command
  buffer was recorded, but this is out of date if the image layout has changed
  during the command buffer. Worst case it could be e.g. UNDEFINED if the image
  was created and not used until the current command buffer.
---
 renderdoc/driver/vulkan/vk_core.h           | 24 +++++++++++
 renderdoc/driver/vulkan/vk_pixelhistory.cpp | 46 +++++++++++++++++----
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/renderdoc/driver/vulkan/vk_core.h b/renderdoc/driver/vulkan/vk_core.h
index 128311973..61590dae6 100644
--- a/renderdoc/driver/vulkan/vk_core.h
+++ b/renderdoc/driver/vulkan/vk_core.h
@@ -1269,6 +1269,30 @@ public:
     return ImageTransitionInfo(m_State, m_QueueFamilyIdx, SeparateDepthStencil());
   }
 
+  const ImageState *GetRecordingLayoutWithinActionCallback(ResourceId id) const
+  {
+    if(m_ActionCallback == NULL)
+    {
+      RDCERR("Attempting to get latest layout with no action callback active");
+      return NULL;
+    }
+
+    if(m_LastCmdBufferID == ResourceId())
+      return NULL;
+
+    auto cmdIt = m_BakedCmdBufferInfo.find(m_LastCmdBufferID);
+
+    if(cmdIt != m_BakedCmdBufferInfo.end())
+    {
+      auto it = cmdIt->second.imageStates.find(id);
+
+      if(it != cmdIt->second.imageStates.end())
+        return &it->second;
+    }
+
+    return NULL;
+  }
+
   // Device initialization
 
   IMPLEMENT_FUNCTION_SERIALISED(VkResult, vkCreateInstance, const VkInstanceCreateInfo *pCreateInfo,
diff --git a/renderdoc/driver/vulkan/vk_pixelhistory.cpp b/renderdoc/driver/vulkan/vk_pixelhistory.cpp
index 500ef7b28..7318aca19 100644
--- a/renderdoc/driver/vulkan/vk_pixelhistory.cpp
+++ b/renderdoc/driver/vulkan/vk_pixelhistory.cpp
@@ -674,6 +674,36 @@ struct VulkanPixelHistoryCallback : public VulkanActionCallback
   }
 
 protected:
+  VkImageLayout GetImageLayout(ResourceId id, VkImageAspectFlagBits aspect, const Subresource &sub)
+  {
+    const ImageState *latestState = m_pDriver->GetRecordingLayoutWithinActionCallback(id);
+
+    VkImageLayout ret = VK_IMAGE_LAYOUT_UNDEFINED;
+
+    if(latestState)
+    {
+      for(auto it = latestState->subresourceStates.begin();
+          it != latestState->subresourceStates.end(); ++it)
+      {
+        const ImageSubresourceRange &range = it->range();
+        if(VkImageAspectFlagBits(range.aspectMask & aspect) == aspect &&
+           // check slice (respecting layerCount == ~0U for 'all')
+           range.baseArrayLayer <= sub.slice &&
+           (sub.slice - range.baseArrayLayer) < range.layerCount &&
+           // check mip (respecting levelCount == ~0U for 'all')
+           range.baseMipLevel <= sub.mip && (sub.mip - range.baseMipLevel) < range.levelCount)
+        {
+          ret = it->state().newLayout;
+        }
+      }
+    }
+
+    if(ret == VK_IMAGE_LAYOUT_UNDEFINED)
+      ret = m_pDriver->GetDebugManager()->GetImageLayout(id, aspect, sub.mip, sub.slice);
+
+    return ret;
+  }
+
   // MakeIncrementStencilPipelineCI fills in the provided pipeCreateInfo
   // to create a graphics pipeline that is based on the original. The modifications
   // to the original pipeline: disables depth test and write, stencil is set
@@ -1935,9 +1965,8 @@ private:
       offset += offsetof(struct PixelHistoryValue, depth);
       aspect = VK_IMAGE_ASPECT_DEPTH_BIT;
     }
-    targetCopyParams.srcImageLayout = m_pDriver->GetDebugManager()->GetImageLayout(
-        GetResID(m_CallbackInfo.targetImage), aspect, m_CallbackInfo.targetSubresource.mip,
-        m_CallbackInfo.targetSubresource.slice);
+    targetCopyParams.srcImageLayout = GetImageLayout(GetResID(m_CallbackInfo.targetImage), aspect,
+                                                     m_CallbackInfo.targetSubresource);
     CopyImagePixel(cmd, targetCopyParams, offset);
 
     // If the target image is a depth/stencil attachment, we already
@@ -2736,9 +2765,9 @@ struct VulkanPixelHistoryPerFragmentCallback : VulkanPixelHistoryCallback
     else
     {
       // Use the layout of the image we are substituting for.
-      VkImageLayout srcImageLayout = m_pDriver->GetDebugManager()->GetImageLayout(
-          GetResID(m_CallbackInfo.targetImage), VK_IMAGE_ASPECT_COLOR_BIT,
-          m_CallbackInfo.targetSubresource.mip, m_CallbackInfo.targetSubresource.slice);
+      VkImageLayout srcImageLayout =
+          GetImageLayout(GetResID(m_CallbackInfo.targetImage), VK_IMAGE_ASPECT_COLOR_BIT,
+                         m_CallbackInfo.targetSubresource);
       colourCopyParams.srcImageLayout = srcImageLayout;
     }
 
@@ -2908,9 +2937,8 @@ struct VulkanPixelHistoryPerFragmentCallback : VulkanPixelHistoryCallback
     VkImageAspectFlagBits aspect = VK_IMAGE_ASPECT_COLOR_BIT;
     if(IsDepthOrStencilFormat(m_CallbackInfo.targetImageFormat))
       aspect = VK_IMAGE_ASPECT_DEPTH_BIT;
-    colourCopyParams.srcImageLayout = m_pDriver->GetDebugManager()->GetImageLayout(
-        GetResID(m_CallbackInfo.targetImage), aspect, m_CallbackInfo.targetSubresource.mip,
-        m_CallbackInfo.targetSubresource.slice);
+    colourCopyParams.srcImageLayout = GetImageLayout(GetResID(m_CallbackInfo.targetImage), aspect,
+                                                     m_CallbackInfo.targetSubresource);
 
     const ModificationValue &premod = m_EventPremods[eid];
     // For every fragment except the last one, retrieve post-modification