From 557ce50737700883b622d333aecda78a55c2d51c Mon Sep 17 00:00:00 2001
From: Caleb Cornett <caleb.cornett@outlook.com>
Date: Fri, 9 Feb 2024 02:17:08 -0600
Subject: [PATCH] PushComputeShaderUniforms + CreateComputePipeline fix

---
 src/Refresh_Driver_D3D11.c | 38 +++++++++++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/Refresh_Driver_D3D11.c b/src/Refresh_Driver_D3D11.c
index 63f2a0d..74f1ae8 100644
--- a/src/Refresh_Driver_D3D11.c
+++ b/src/Refresh_Driver_D3D11.c
@@ -1132,7 +1132,6 @@ static Refresh_ComputePipeline* D3D11_CreateComputePipeline(
 
 	pipeline->numTextures = computeShaderInfo->imageBindingCount;
 	pipeline->numBuffers = computeShaderInfo->bufferBindingCount;
-	pipeline->computeUniformBlockSize = computeShaderInfo->uniformBufferSize;
 
 	if (shaderModule->shader == NULL)
 	{
@@ -1165,6 +1164,10 @@ static Refresh_ComputePipeline* D3D11_CreateComputePipeline(
 		ERROR_CHECK_RETURN("Could not create compute shader", NULL); /* FIXME: This leaks the pipeline! */
 	}
 	pipeline->computeShader = (ID3D11ComputeShader*) shaderModule->shader;
+	pipeline->computeUniformBlockSize = D3D11_INTERNAL_NextHighestAlignment(
+		(uint32_t) computeShaderInfo->uniformBufferSize,
+		256
+	);
 
 	return (Refresh_ComputePipeline*) pipeline;
 }
@@ -2130,8 +2133,37 @@ static uint32_t D3D11_PushComputeShaderUniforms(
 	void *data,
 	uint32_t dataLengthInBytes
 ) {
-	NOT_IMPLEMENTED
-	return 0;
+	D3D11Renderer *renderer = (D3D11Renderer*) driverData;
+	D3D11CommandBuffer *d3d11CommandBuffer = (D3D11CommandBuffer*) commandBuffer;
+	D3D11ComputePipeline *computePipeline = d3d11CommandBuffer->computePipeline;
+	uint32_t offset;
+
+	if (d3d11CommandBuffer->computeUniformBuffer->offset + computePipeline->computeUniformBlockSize >= UBO_BUFFER_SIZE)
+	{
+		/* Out of space! Get a new uniform buffer. */
+		D3D11_INTERNAL_AcquireUniformBuffer(
+			renderer,
+			d3d11CommandBuffer,
+			&d3d11CommandBuffer->computeUniformBuffer,
+			computePipeline->computeUniformBlockSize
+		);
+	}
+
+	offset = d3d11CommandBuffer->computeUniformBuffer->offset;
+
+	D3D11_INTERNAL_SetBufferData(
+		renderer,
+		d3d11CommandBuffer,
+		d3d11CommandBuffer->computeUniformBuffer->d3d11Buffer,
+		d3d11CommandBuffer->computeUniformBuffer->offset,
+		data,
+		dataLengthInBytes,
+		0 /* FIXME: Should be NoOverwrite! */
+	);
+
+	d3d11CommandBuffer->computeUniformBuffer->offset += (uint32_t) computePipeline->computeUniformBlockSize; /* FIXME: Is this cast safe? */
+
+	return offset;
 }
 
 /* Samplers */