From 75100adf91972e325e77bfdf0e448053359925aa Mon Sep 17 00:00:00 2001
From: Tim <15017472+doodlum@users.noreply.github.com>
Date: Thu, 19 Dec 2024 13:22:44 +0000
Subject: [PATCH] chore: unregister buffers

---
 src/Raytracing.cpp | 52 ++++++++++++++++++++++++++++++++++++++--------
 src/Raytracing.h   |  3 +++
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/Raytracing.cpp b/src/Raytracing.cpp
index 450cf8d73..15a268135 100644
--- a/src/Raytracing.cpp
+++ b/src/Raytracing.cpp
@@ -340,24 +340,25 @@ struct int4
 	int w;
 };
 
-float4x4 GetBoneTransformMatrix(float4 bonePositions[240], int4 boneIndices, float4 pivot, float4 boneWeights)
+float4x4 GetBoneTransformMatrix(float4 bonePositions[240], float4 boneIndices, float4 pivot, float4 boneWeights)
 {
-	boneIndices.x *= 765.01;
-	boneIndices.y *= 765.01;
-	boneIndices.z *= 765.01;
-	boneIndices.w *= 765.01;
+	int4 boneIndicesInt;
+	
+	boneIndices *= 765.01f;
+
+	boneIndicesInt = { (int)boneIndices.x, (int)boneIndices.y, (int)boneIndices.z, (int)boneIndices.w };
 
 	float4 zeroes = { 0, 0, 0, 0 };
 	float4x4 pivotMatrix = float4x4(zeroes, zeroes, zeroes, pivot).Transpose();
 
 	float4x4 boneMatrix1 =
-		float4x4(bonePositions[boneIndices.x], bonePositions[boneIndices.x + 1], bonePositions[boneIndices.x + 2], zeroes);
+		float4x4(bonePositions[boneIndicesInt.x], bonePositions[boneIndicesInt.x + 1], bonePositions[boneIndicesInt.x + 2], zeroes);
 	float4x4 boneMatrix2 =
-		float4x4(bonePositions[boneIndices.y], bonePositions[boneIndices.y + 1], bonePositions[boneIndices.y + 2], zeroes);
+		float4x4(bonePositions[boneIndicesInt.y], bonePositions[boneIndicesInt.y + 1], bonePositions[boneIndicesInt.y + 2], zeroes);
 	float4x4 boneMatrix3 =
-		float4x4(bonePositions[boneIndices.z], bonePositions[boneIndices.z + 1], bonePositions[boneIndices.z + 2], zeroes);
+		float4x4(bonePositions[boneIndicesInt.z], bonePositions[boneIndicesInt.z + 1], bonePositions[boneIndicesInt.z + 2], zeroes);
 	float4x4 boneMatrix4 =
-		float4x4(bonePositions[boneIndices.w], bonePositions[boneIndices.w + 1], bonePositions[boneIndices.w + 2], zeroes);
+		float4x4(bonePositions[boneIndicesInt.w], bonePositions[boneIndicesInt.w + 1], bonePositions[boneIndicesInt.w + 2], zeroes);
 
 	float4 ones = { 1.0, 1.0, 1.0, 1.0 };
 	float4x4 unitMatrix = float4x4(ones, ones, ones, ones);
@@ -581,16 +582,39 @@ Raytracing::BufferData Raytracing::AllocateBuffer(const D3D11_BUFFER_DESC* pDesc
 	return data;
 }
 
+struct ID3D11Buffer_Release
+{
+	static void thunk(ID3D11Buffer* This)
+	{
+		Raytracing::GetSingleton()->UnregisterVertexBuffer(This);
+		Raytracing::GetSingleton()->UnregisterIndexBuffer(This);
+		func(This);
+	}
+	static inline REL::Relocation<decltype(thunk)> func;
+};
+
+
+bool hooked = false;
+
+
 void Raytracing::RegisterVertexBuffer(const D3D11_BUFFER_DESC* pDesc, const D3D11_SUBRESOURCE_DATA* pInitialData, ID3D11Buffer** ppBuffer)
 {
 	BufferData data = AllocateBuffer(pDesc, pInitialData);
 	vertexBuffers.insert({ *ppBuffer, data });
+	if (!hooked) {
+		stl::detour_vfunc<2, ID3D11Buffer_Release>(*ppBuffer);
+		hooked = true;
+	}
 }
 
 void Raytracing::RegisterIndexBuffer(const D3D11_BUFFER_DESC* pDesc, const D3D11_SUBRESOURCE_DATA* pInitialData, ID3D11Buffer** ppBuffer)
 {
 	BufferData data = AllocateBuffer(pDesc, pInitialData);
 	indexBuffers.insert({ *ppBuffer, data });
+	if (!hooked) {
+		stl::detour_vfunc<2, ID3D11Buffer_Release>(*ppBuffer);
+		hooked = true;
+	}
 }
 
 void Raytracing::RegisterInputLayout(ID3D11InputLayout* ppInputLayout, D3D11_INPUT_ELEMENT_DESC* pInputElementDescs, UINT NumElements)
@@ -624,6 +648,16 @@ void Raytracing::RegisterInputLayout(ID3D11InputLayout* ppInputLayout, D3D11_INP
 	}
 }
 
+void Raytracing::UnregisterVertexBuffer(ID3D11Buffer* ppBuffer)
+{
+	vertexBuffers.erase(ppBuffer);
+}
+
+void Raytracing::UnregisterIndexBuffer(ID3D11Buffer* ppBuffer)
+{
+	indexBuffers.erase(ppBuffer);
+}
+
 void Raytracing::TransitionResources(D3D12_RESOURCE_STATES stateBefore, D3D12_RESOURCE_STATES stateAfter)
 {
 	std::vector<D3D12_RESOURCE_BARRIER> barriers;
diff --git a/src/Raytracing.h b/src/Raytracing.h
index 87244cb5f..647c8a28c 100644
--- a/src/Raytracing.h
+++ b/src/Raytracing.h
@@ -166,6 +166,9 @@ class Raytracing
 
 	void RegisterInputLayout(ID3D11InputLayout* ppInputLayout, D3D11_INPUT_ELEMENT_DESC* pInputElementDescs, UINT NumElements);
 
+	void UnregisterVertexBuffer(ID3D11Buffer* ppBuffer);
+	void UnregisterIndexBuffer(ID3D11Buffer* ppBuffer);
+
 	void TransitionResources(D3D12_RESOURCE_STATES stateBefore, D3D12_RESOURCE_STATES stateAfter);
 
 	void InitFenceAndEvent();