From 06b7db14d58aed8e61c735d5768787107158a29c Mon Sep 17 00:00:00 2001
From: Kah Wei <msntkw@msn.com>
Date: Tue, 3 Jan 2023 18:57:14 +0800
Subject: [PATCH] Added code for bone matrices

---
 .../Graphics/MiddleEnd/Batching/SHBatch.cpp   | 101 +++++++++++++-----
 .../src/Graphics/MiddleEnd/Batching/SHBatch.h |   1 +
 2 files changed, 75 insertions(+), 27 deletions(-)

diff --git a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
index bb717ea5..7a1d46ed 100644
--- a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
+++ b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.cpp
@@ -449,13 +449,13 @@ namespace SHADE
         if (CURR_INSTANCES > 0)
         {
           drawData.emplace_back(vk::DrawIndexedIndirectCommand
-           {
-               .indexCount = subBatch.Mesh->IndexCount,
-               .instanceCount = CURR_INSTANCES,
-               .firstIndex = subBatch.Mesh->FirstIndex,
-               .vertexOffset = subBatch.Mesh->FirstVertex,
-               .firstInstance = nextInstanceIndex
-           });
+          {
+              .indexCount = subBatch.Mesh->IndexCount,
+              .instanceCount = CURR_INSTANCES,
+              .firstIndex = subBatch.Mesh->FirstIndex,
+              .vertexOffset = subBatch.Mesh->FirstVertex,
+              .firstInstance = nextInstanceIndex
+          });
         }
         nextInstanceIndex += CURR_INSTANCES;
 
@@ -520,7 +520,6 @@ namespace SHADE
             const auto& BONE_MATRICES = animator->GetBoneMatrices();
             boneMatrixData.insert(boneMatrixData.end(), BONE_MATRICES.cbegin(), BONE_MATRICES.cend());
           }
-
         }
       }
 
@@ -558,23 +557,7 @@ namespace SHADE
       // - Material Properties Buffer
       rebuildMaterialBuffers(frameIndex, descPool);
       // - Bone Buffers
-	  if (!boneMatrixIndices.empty())
-      {
-        const uint32_t BONE_IDX_DATA_BYTES = static_cast<uint32_t>(boneMatrixIndices.size() * sizeof(uint32_t));
-        SHVkUtil::EnsureBufferAndCopyHostVisibleData
-        (
-          device, boneFirstIndexBuffer[frameIndex], boneMatrixIndices.data(), BONE_IDX_DATA_BYTES,
-          BuffUsage::eVertexBuffer,
-          "Batch Bone Index Buffer"
-        );
-        const uint32_t BONE_MTX_DATA_BYTES = static_cast<uint32_t>(boneMatrixData.size() * sizeof(uint32_t));
-        SHVkUtil::EnsureBufferAndCopyHostVisibleData
-        (
-          device, boneMatrixBuffer[frameIndex], boneMatrixData.data(), BONE_MTX_DATA_BYTES,
-          BuffUsage::eStorageBuffer,
-          "Batch Bone Matrix Buffer"
-        );
-      }
+      rebuildBoneBuffers(frameIndex, descPool);
     }
 
     // Mark this frame as no longer dirty
@@ -605,7 +588,7 @@ namespace SHADE
     if (matPropsDescSet[frameIndex])
     {
       cmdBuffer->BindDescriptorSet
-      (
+      ( 
         matPropsDescSet[frameIndex],
         SH_PIPELINE_TYPE::GRAPHICS,
         SHGraphicsConstants::DescriptorSetIndex::PER_INSTANCE,
@@ -615,7 +598,13 @@ namespace SHADE
     if (boneMatrixBuffer[frameIndex] && boneFirstIndexBuffer[frameIndex])
     {
       cmdBuffer->BindVertexBuffer(SHGraphicsConstants::VertexBufferBindings::BONE_INDICES, boneFirstIndexBuffer[frameIndex], 0);
-      // TODO: Bind storage buffer
+      cmdBuffer->BindDescriptorSet
+      (
+        boneMatricesDescSet[frameIndex],
+        SH_PIPELINE_TYPE::GRAPHICS,
+        SHGraphicsConstants::DescriptorSetIndex::PER_INSTANCE,
+        dynamicOffset
+      );
     }
     cmdBuffer->DrawMultiIndirect(drawDataBuffer[frameIndex], static_cast<uint32_t>(drawData.size()));
     cmdBuffer->EndLabeledSegment();
@@ -673,4 +662,62 @@ namespace SHADE
     }
   }
 
+  void SHBatch::rebuildBoneBuffers(uint32_t frameIndex, Handle<SHVkDescriptorPool> descPool)
+  {
+    // Using Declarations
+    using BuffUsage = vk::BufferUsageFlagBits;
+
+    // Nothing to rebuild
+    if (boneMatrixData.empty())
+      return;
+
+    // Update GPU Buffers
+    const uint32_t BONE_IDX_DATA_BYTES = static_cast<uint32_t>(boneMatrixIndices.size() * sizeof(uint32_t));
+    SHVkUtil::EnsureBufferAndCopyHostVisibleData
+    (
+      device, boneFirstIndexBuffer[frameIndex], boneMatrixIndices.data(), BONE_IDX_DATA_BYTES,
+      BuffUsage::eVertexBuffer,
+      "Batch Bone Index Buffer"
+    );
+    const uint32_t BONE_MTX_DATA_BYTES = static_cast<uint32_t>(boneMatrixData.size() * sizeof(uint32_t));
+    SHVkUtil::EnsureBufferAndCopyHostVisibleData
+    (
+      device, boneMatrixBuffer[frameIndex], boneMatrixData.data(), BONE_MTX_DATA_BYTES,
+      BuffUsage::eStorageBuffer,
+      "Batch Bone Matrix Buffer"
+    );
+
+    // Update descriptor set buffer
+    if (!boneMatricesDescSet[frameIndex])
+    {
+      boneMatricesDescSet[frameIndex] = descPool->Allocate
+      (
+        { SHGraphicsGlobalData::GetDescSetLayouts()[SHGraphicsConstants::DescriptorSetIndex::PER_INSTANCE] },
+        { 0 }
+      );
+
+#ifdef _DEBUG
+      const auto& DESC_SETS = boneMatricesDescSet[frameIndex]->GetVkHandle();
+      for (auto descSet : DESC_SETS)
+      {
+        SET_VK_OBJ_NAME(device, vk::ObjectType::eDescriptorSet, descSet, "[Descriptor Set] Batch Bone Data");
+      }
+#endif
+    }
+    std::array<Handle<SHVkBuffer>, 1>  bufferList = { boneMatrixBuffer[frameIndex] };
+    boneMatricesDescSet[frameIndex]->ModifyWriteDescBuffer
+    (
+      SHGraphicsConstants::DescriptorSetIndex::PER_INSTANCE,
+      SHGraphicsConstants::DescriptorSetBindings::BATCHED_PER_INST_DATA,
+      bufferList,
+      0, static_cast<uint32_t>(boneMatrixData.size() * sizeof(SHMatrix))
+    );
+    boneMatricesDescSet[frameIndex]->UpdateDescriptorSetBuffer
+    (
+      SHGraphicsConstants::DescriptorSetIndex::PER_INSTANCE,
+      SHGraphicsConstants::DescriptorSetBindings::BATCHED_PER_INST_DATA
+    );
+
+    // TODO: Need to merge this with the material one
+  }
 }
diff --git a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
index 477a5720..66260c7d 100644
--- a/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
+++ b/SHADE_Engine/src/Graphics/MiddleEnd/Batching/SHBatch.h
@@ -141,5 +141,6 @@ namespace SHADE
     /*-----------------------------------------------------------------------------*/
     void setAllDirtyFlags();
     void rebuildMaterialBuffers(uint32_t frameIndex, Handle<SHVkDescriptorPool> descPool);
+    void rebuildBoneBuffers(uint32_t frameIndex, Handle<SHVkDescriptorPool> descPool);
   };
 }