From eab2f2d54ac686f31e9b9aab4264dfe52463e8ec Mon Sep 17 00:00:00 2001
From: Kah Wei <msntkw@msn.com>
Date: Sun, 29 Jan 2023 16:03:08 +0800
Subject: [PATCH] Fixed bug where replacing an animator's rig causes the GPU to
 be lost

---
 .../Graphics/MiddleEnd/Batching/SHBatch.cpp   | 66 ++++++++++---------
 .../src/Graphics/MiddleEnd/Batching/SHBatch.h |  7 ++
 2 files changed, 41 insertions(+), 32 deletions(-)

diff --git a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
index 330c395c..685511fe 100644
--- a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
+++ b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
@@ -428,11 +428,7 @@ namespace SHADE
     }
 
     // Update GPU Buffers
-    const uint32_t BONE_MTX_DATA_BYTES = static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix));
-    if (boneMatrixBuffer[frameIndex])
-    {
-      boneMatrixBuffer[frameIndex]->WriteToMemory(boneMatrixData.data(), BONE_MTX_DATA_BYTES, 0, 0);
-    }
+    rebuildBoneMatrixDescSetBuffer(frameIndex);
   }
 
   void SHBatch::Build(Handle<SHVkLogicalDevice> _device, Handle<SHVkDescriptorPool> descPool, uint32_t frameIndex)
@@ -734,7 +730,7 @@ namespace SHADE
     // Using Declarations and constants
     using BuffUsage = vk::BufferUsageFlagBits;
     using PreDefDescLayoutType = SHGraphicsPredefinedData::PredefinedDescSetLayoutTypes;
-    static constexpr uint32_t MATERIAL_DESC_SET_INDEX = 0;
+    
 
     /* Create Descriptor Sets if Needed */
     PreDefDescLayoutType layoutTypes = {};
@@ -801,35 +797,41 @@ namespace SHADE
     /* Animation Bone Data */
     if (MUST_BUILD_BONE_DESC)
     {
-      // Update GPU Buffers
-      const uint32_t BONE_MTX_DATA_BYTES = static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix));
-      SHVkUtil::EnsureBufferAndCopyHostVisibleData
-      (
-        device, boneMatrixBuffer[frameIndex], boneMatrixData.data(), BONE_MTX_DATA_BYTES,
-        BuffUsage::eStorageBuffer,
-        "Batch Bone Matrix Buffer"
-      );
-
-      // Update descriptor set buffer
-      std::array<Handle<SHVkBuffer>, 1>  bufferList = { boneMatrixBuffer[frameIndex] };
-      instanceDataDescSet[frameIndex]->ModifyWriteDescBuffer
-      (
-        MATERIAL_DESC_SET_INDEX,
-        SHGraphicsConstants::DescriptorSetBindings::PER_INST_BONE_DATA,
-        bufferList,
-        0, 
-        static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix))
-      );
-
-      // Update the descriptor set buffer
-      instanceDataDescSet[frameIndex]->UpdateDescriptorSetBuffer
-      (
-        MATERIAL_DESC_SET_INDEX,
-        SHGraphicsConstants::DescriptorSetBindings::PER_INST_BONE_DATA
-      );
+      rebuildBoneMatrixDescSetBuffer(frameIndex);
     }
   }
 
+  void SHBatch::rebuildBoneMatrixDescSetBuffer(uint32_t frameIndex)
+  {
+    using BuffUsage = vk::BufferUsageFlagBits;
+    // Update GPU Buffers
+    const uint32_t BONE_MTX_DATA_BYTES = static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix));
+    SHVkUtil::EnsureBufferAndCopyHostVisibleData
+    (
+      device, boneMatrixBuffer[frameIndex], boneMatrixData.data(), BONE_MTX_DATA_BYTES,
+      BuffUsage::eStorageBuffer,
+      "Batch Bone Matrix Buffer"
+    );
+
+    // Update descriptor set buffer
+    std::array<Handle<SHVkBuffer>, 1>  bufferList = { boneMatrixBuffer[frameIndex] };
+    instanceDataDescSet[frameIndex]->ModifyWriteDescBuffer
+    (
+      MATERIAL_DESC_SET_INDEX,
+      SHGraphicsConstants::DescriptorSetBindings::PER_INST_BONE_DATA,
+      bufferList,
+      0, 
+      static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix))
+    );
+
+    // Update the descriptor set buffer
+    instanceDataDescSet[frameIndex]->UpdateDescriptorSetBuffer
+    (
+      MATERIAL_DESC_SET_INDEX,
+      SHGraphicsConstants::DescriptorSetBindings::PER_INST_BONE_DATA
+    );
+  }
+
   bool SHBatch::checkIfIsAnimatedPipeline(Handle<SHVkPipeline> pipeline)
   {
     if (!pipeline || !pipeline->GetPipelineLayout())
diff --git a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
index 86e17035..b9192663 100644
--- a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
+++ b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
@@ -106,6 +106,11 @@ namespace SHADE
     using TripleBool    = std::array<bool, SHGraphicsConstants::NUM_FRAME_BUFFERS>;
     using TripleBuffer  = std::array<Handle<SHVkBuffer>, SHGraphicsConstants::NUM_FRAME_BUFFERS>;
     using TripleDescSet = std::array<Handle<SHVkDescriptorSetGroup>, SHGraphicsConstants::NUM_FRAME_BUFFERS>;
+    
+    /*---------------------------------------------------------------------------------*/
+    /* Constants                                                                       */
+    /*---------------------------------------------------------------------------------*/
+    static constexpr uint32_t MATERIAL_DESC_SET_INDEX = 0;
 
     /*---------------------------------------------------------------------------------*/
     /* Data Members                                                                    */
@@ -146,6 +151,8 @@ namespace SHADE
     /*-----------------------------------------------------------------------------*/
     void setAllDirtyFlags();
     void rebuildDescriptorSetBuffers(uint32_t frameIndex, Handle<SHVkDescriptorPool> descPool);
+    void rebuildBoneMatrixDescSetBuffer(uint32_t frameIndex);
     static bool checkIfIsAnimatedPipeline(Handle<SHVkPipeline> pipeline);
+
   };
 }