From 7e7d2884c8098b9caffdec4896795175ae76319e Mon Sep 17 00:00:00 2001
From: niki <nikeinikei98@gmail.com>
Date: Mon, 9 Jan 2023 16:37:12 +0100
Subject: [PATCH] vulkan: fix shader crash

---
 src/modules/graphics/vulkan/Graphics.cpp |  6 ++++--
 src/modules/graphics/vulkan/Graphics.h   |  2 +-
 src/modules/graphics/vulkan/Shader.cpp   | 10 +++-------
 3 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/modules/graphics/vulkan/Graphics.cpp b/src/modules/graphics/vulkan/Graphics.cpp
index 768d4056e..8270faf28 100644
--- a/src/modules/graphics/vulkan/Graphics.cpp
+++ b/src/modules/graphics/vulkan/Graphics.cpp
@@ -1131,6 +1131,8 @@ void Graphics::beginFrame()
 
 	Vulkan::resetShaderSwitches();
 
+	for (const auto shader : usedShadersInFrame)
+		shader->newFrame();
 	usedShadersInFrame.clear();
 }
 
@@ -2478,9 +2480,9 @@ void Graphics::setComputeShader(Shader *shader)
 	computeShader = shader;
 }
 
-std::set<Shader*> &Graphics::getUsedShadersInFrame()
+void Graphics::markShaderUsed(Shader *shader)
 {
-	return usedShadersInFrame;
+	usedShadersInFrame.insert(shader);
 }
 
 VkSampler Graphics::getCachedSampler(const SamplerState &samplerState)
diff --git a/src/modules/graphics/vulkan/Graphics.h b/src/modules/graphics/vulkan/Graphics.h
index 7d88e58f6..8109ed3b1 100644
--- a/src/modules/graphics/vulkan/Graphics.h
+++ b/src/modules/graphics/vulkan/Graphics.h
@@ -293,7 +293,7 @@ public:
 	graphics::Texture *getDefaultTexture() const;
 	VkSampler getCachedSampler(const SamplerState &sampler);
 	void setComputeShader(Shader *computeShader);
-	std::set<Shader*> &getUsedShadersInFrame();
+	void markShaderUsed(Shader*);
 	graphics::Shader::BuiltinUniformData getCurrentBuiltinUniformData();
 	const OptionalDeviceFeatures &getEnabledOptionalDeviceExtensions() const;
 	VkSampleCountFlagBits getMsaaCount(int requestedMsaa) const;
diff --git a/src/modules/graphics/vulkan/Shader.cpp b/src/modules/graphics/vulkan/Shader.cpp
index d09cdb36a..7bf87ad0f 100644
--- a/src/modules/graphics/vulkan/Shader.cpp
+++ b/src/modules/graphics/vulkan/Shader.cpp
@@ -200,6 +200,7 @@ bool Shader::loadVolatile()
 	currentFrame = 0;
 	currentUsedUniformStreamBuffersCount = 0;
 	currentUsedDescriptorSetsCount = 0;
+	newFrame();
 
 	return true;
 }
@@ -384,7 +385,7 @@ void Shader::cmdPushDescriptorSets(VkCommandBuffer commandBuffer, VkPipelineBind
 	for (const auto &u : uniformInfos)
 	{
 		if (updatedUniforms.find(u.second.location) == updatedUniforms.end())
-			updateUniform(&u.second, u.second.count);
+			updateUniform(&u.second, u.second.count, true);
 	}
 
 	vkCmdBindDescriptorSets(commandBuffer, bindPoint, pipelineLayout, 0, 1, &currentDescriptorSet, 0, nullptr);
@@ -406,12 +407,7 @@ Shader::~Shader()
 
 void Shader::attach()
 {
-	auto &usedShadersInFrame = vgfx->getUsedShadersInFrame();
-	if (usedShadersInFrame.find(this) == usedShadersInFrame.end())
-	{
-		newFrame();
-		usedShadersInFrame.insert(this);
-	}
+	vgfx->markShaderUsed(this);
 
 	if (!isCompute)
 	{