From 3706949733ed34c48fef14c83bca958f9238007f Mon Sep 17 00:00:00 2001
From: VReaperV <reaper9977@gmail.com>
Date: Wed, 29 Jan 2025 00:36:09 +0300
Subject: [PATCH] Fix glsl_restart with material system

---
 src/engine/renderer/Material.cpp | 17 +++++++++++++++++
 src/engine/renderer/Material.h   |  6 ++++++
 src/engine/renderer/tr_init.cpp  |  6 ++++++
 3 files changed, 29 insertions(+)

diff --git a/src/engine/renderer/Material.cpp b/src/engine/renderer/Material.cpp
index 07fcaa61d2..18c4520811 100644
--- a/src/engine/renderer/Material.cpp
+++ b/src/engine/renderer/Material.cpp
@@ -1360,6 +1360,13 @@ void MaterialSystem::ProcessStage( drawSurf_t* drawSurf, shaderStage_t* pStage,
 	pStage->materialProcessor( &material, pStage, drawSurf );
 	pStage->paddedSize = material.shader->GetPaddedSize();
 
+	// HACK: Copy the shaderStage_t and drawSurf_t that we need into the material, so we can use it with glsl_restart
+	material.refStage = pStage;
+	material.refDrawSurf = *drawSurf;
+	material.refDrawSurf.entity = nullptr;
+	material.refDrawSurf.depthSurface = nullptr;
+	material.refDrawSurf.fogSurface = nullptr;
+
 	std::vector<Material>& materials = materialPacks[materialPack].materials;
 	std::vector<Material>::iterator currentSearchIt = materials.begin();
 	std::vector<Material>::iterator materialIt;
@@ -1507,6 +1514,16 @@ void MaterialSystem::AddAllWorldSurfaces() {
 	generatingWorldCommandBuffer = false;
 }
 
+void MaterialSystem::GLSLRestart() {
+	for ( MaterialPack& materialPack : materialPacks ) {
+		for ( Material& material : materialPack.materials ) {
+			/* We only really need to reset material.shader here,
+			but we keep a copy of the original shaderStage_t and drawSurf_t used for it so we don't change the actual material data */
+			material.refStage->materialProcessor( &material, material.refStage, &material.refDrawSurf );
+		}
+	}
+}
+
 void MaterialSystem::AddStageTextures( drawSurf_t* drawSurf, const uint32_t stage, Material* material ) {
 	TextureData textureData;
 	const shaderStage_t* pStage = &drawSurf->shader->stages[stage];
diff --git a/src/engine/renderer/Material.h b/src/engine/renderer/Material.h
index ea49e676b1..dc4b7b9121 100644
--- a/src/engine/renderer/Material.h
+++ b/src/engine/renderer/Material.h
@@ -93,6 +93,10 @@ struct Material {
 	GLuint program = 0;
 	GLShader* shader;
 
+	// Used only for glsl_restart
+	shaderStage_t* refStage;
+	drawSurf_t refDrawSurf;
+
 	int deformIndex;
 	bool tcGenEnvironment;
 	bool tcGen_Lightmap;
@@ -369,6 +373,8 @@ class MaterialSystem {
 
 	void AddAllWorldSurfaces();
 
+	void GLSLRestart();
+
 	void Free();
 
 	private:
diff --git a/src/engine/renderer/tr_init.cpp b/src/engine/renderer/tr_init.cpp
index c69eaad77b..ae01cd72f6 100644
--- a/src/engine/renderer/tr_init.cpp
+++ b/src/engine/renderer/tr_init.cpp
@@ -1128,6 +1128,12 @@ ScreenshotCmd screenshotPNGRegistration("screenshotPNG", ssFormat_t::SSF_PNG, "p
 					stage->deformIndex = deformIndex;
 				}
 			}
+
+			if ( glConfig2.usingMaterialSystem ) {
+				/* GLSL shaders linked to materials will be invalidated by glsl_restart,
+				so we need to reset them here */
+				materialSystem.GLSLRestart();
+			}
 		}
 	};
 	static GlslRestartCmd glslRestartCmdRegistration;