From 16e410d310a9c370d58eb08e247e2d35140da451 Mon Sep 17 00:00:00 2001
From: Benson Joeris <bjoeris@google.com>
Date: Thu, 30 May 2019 15:10:11 -0400
Subject: [PATCH] Vulkan: Add image subresource tracking to
 VulkanResourceManager

Change-Id: I9e14cb8bbb98142c8b6ff194c0d989e22f040a06
---
 renderdoc/driver/vulkan/vk_manager.cpp | 40 ++++++++++++++++++++++++++
 renderdoc/driver/vulkan/vk_manager.h   |  8 ++++++
 2 files changed, 48 insertions(+)

diff --git a/renderdoc/driver/vulkan/vk_manager.cpp b/renderdoc/driver/vulkan/vk_manager.cpp
index f3c653dab..6b302f15c 100644
--- a/renderdoc/driver/vulkan/vk_manager.cpp
+++ b/renderdoc/driver/vulkan/vk_manager.cpp
@@ -731,6 +731,20 @@ ResourceId VulkanResourceManager::GetFirstIDForHandle(uint64_t handle)
   return ResourceId();
 }
 
+void VulkanResourceManager::MarkImageFrameReferenced(const VkResourceRecord *img,
+                                                     const ImageRange &range, FrameRefType refType)
+{
+  MarkImageFrameReferenced(img->GetResourceID(), img->resInfo->imageInfo, range, refType);
+}
+
+void VulkanResourceManager::MarkImageFrameReferenced(ResourceId img, const ImageInfo &imageInfo,
+                                                     const ImageRange &range, FrameRefType refType)
+{
+  FrameRefType maxRef = MarkImageReferenced(m_ImgFrameRefs, img, imageInfo, range, refType);
+  MarkResourceFrameReferenced(
+      img, maxRef, [](FrameRefType x, FrameRefType y) -> FrameRefType { return std::max(x, y); });
+}
+
 void VulkanResourceManager::MarkMemoryFrameReferenced(ResourceId mem, VkDeviceSize offset,
                                                       VkDeviceSize size, FrameRefType refType)
 {
@@ -741,6 +755,18 @@ void VulkanResourceManager::MarkMemoryFrameReferenced(ResourceId mem, VkDeviceSi
       mem, maxRef, [](FrameRefType x, FrameRefType y) -> FrameRefType { return std::max(x, y); });
 }
 
+void VulkanResourceManager::MergeReferencedImages(std::map<ResourceId, ImgRefs> &imgRefs)
+{
+  for(auto j = imgRefs.begin(); j != imgRefs.end(); j++)
+  {
+    auto i = m_ImgFrameRefs.find(j->first);
+    if(i == m_ImgFrameRefs.end())
+      m_ImgFrameRefs.insert(*j);
+    else
+      i->second.Merge(j->second);
+  }
+}
+
 void VulkanResourceManager::MergeReferencedMemory(std::map<ResourceId, MemRefs> &memRefs)
 {
   SCOPED_LOCK(m_Lock);
@@ -755,6 +781,11 @@ void VulkanResourceManager::MergeReferencedMemory(std::map<ResourceId, MemRefs>
   }
 }
 
+void VulkanResourceManager::ClearReferencedImages()
+{
+  m_ImgFrameRefs.clear();
+}
+
 void VulkanResourceManager::ClearReferencedMemory()
 {
   SCOPED_LOCK(m_Lock);
@@ -771,6 +802,15 @@ MemRefs *VulkanResourceManager::FindMemRefs(ResourceId mem)
     return NULL;
 }
 
+ImgRefs *VulkanResourceManager::FindImgRefs(ResourceId img)
+{
+  auto it = m_ImgFrameRefs.find(img);
+  if(it != m_ImgFrameRefs.end())
+    return &it->second;
+  else
+    return NULL;
+}
+
 bool VulkanResourceManager::Prepare_InitialState(WrappedVkRes *res)
 {
   return m_Core->Prepare_InitialState(res);
diff --git a/renderdoc/driver/vulkan/vk_manager.h b/renderdoc/driver/vulkan/vk_manager.h
index 398568248..c818d5153 100644
--- a/renderdoc/driver/vulkan/vk_manager.h
+++ b/renderdoc/driver/vulkan/vk_manager.h
@@ -422,12 +422,19 @@ public:
 
   void SetInternalResource(ResourceId id);
 
+  void MarkImageFrameReferenced(const VkResourceRecord *img, const ImageRange &range,
+                                FrameRefType refType);
+  void MarkImageFrameReferenced(ResourceId img, const ImageInfo &imageInfo, const ImageRange &range,
+                                FrameRefType refType);
   void MarkMemoryFrameReferenced(ResourceId mem, VkDeviceSize start, VkDeviceSize end,
                                  FrameRefType refType);
 
   void MergeReferencedMemory(std::map<ResourceId, MemRefs> &memRefs);
+  void MergeReferencedImages(std::map<ResourceId, ImgRefs> &imgRefs);
+  void ClearReferencedImages();
   void ClearReferencedMemory();
   MemRefs *FindMemRefs(ResourceId mem);
+  ImgRefs *FindImgRefs(ResourceId img);
 
   inline bool OptimizeInitialState() { return m_OptimizeInitialState; }
 private:
@@ -444,5 +451,6 @@ private:
   CaptureState m_State;
   WrappedVulkan *m_Core;
   std::map<ResourceId, MemRefs> m_MemFrameRefs;
+  std::map<ResourceId, ImgRefs> m_ImgFrameRefs;
   bool m_OptimizeInitialState = false;
 };