From b28e61d00d1f1eda36a5ad5759767735534cde76 Mon Sep 17 00:00:00 2001
From: bjorn <tie.372@gmail.com>
Date: Sat, 5 Jun 2021 18:59:02 -0600
Subject: [PATCH] Shader:getLocalThreadgroupSize;

---
 src/modules/graphics/Shader.cpp      | 21 ++++++++++++++++++++
 src/modules/graphics/Shader.h        |  3 +++
 src/modules/graphics/wrap_Shader.cpp | 29 +++++++++++++++++++++++-----
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/modules/graphics/Shader.cpp b/src/modules/graphics/Shader.cpp
index 2a89fd8c8..e287ecf60 100644
--- a/src/modules/graphics/Shader.cpp
+++ b/src/modules/graphics/Shader.cpp
@@ -675,6 +675,13 @@ void Shader::validateDrawState(PrimitiveType primtype, Texture *maintex) const
 	}
 }
 
+void Shader::getLocalThreadgroupSize(int *x, int *y, int *z)
+{
+	*x = validationReflection.localThreadgroupSize[0];
+	*y = validationReflection.localThreadgroupSize[1];
+	*z = validationReflection.localThreadgroupSize[2];
+}
+
 bool Shader::validate(StrongRef<ShaderStage> stages[], std::string& err)
 {
 	ValidationReflection reflection;
@@ -710,6 +717,20 @@ bool Shader::validateInternal(StrongRef<ShaderStage> stages[], std::string &err,
 		reflection.usesPointSize = vertintermediate->inIoAccessed("gl_PointSize");
 	}
 
+	if (stages[SHADERSTAGE_COMPUTE] != nullptr)
+	{
+		for (int i = 0; i < 3; i++)
+		{
+			reflection.localThreadgroupSize[i] = program.getLocalSize(i);
+
+			if (reflection.localThreadgroupSize[i] <= 0)
+			{
+				err = "Shader validation error:\nNegative local threadgroup size.";
+				return false;
+			}
+		}
+	}
+
 	for (int i = 0; i < program.getNumBufferBlocks(); i++)
 	{
 		const glslang::TObjectReflection &info = program.getBufferBlock(i);
diff --git a/src/modules/graphics/Shader.h b/src/modules/graphics/Shader.h
index 85137a57d..93074c7ab 100644
--- a/src/modules/graphics/Shader.h
+++ b/src/modules/graphics/Shader.h
@@ -216,6 +216,8 @@ public:
 	TextureType getMainTextureType() const;
 	void validateDrawState(PrimitiveType primtype, Texture *maintexture) const;
 
+	void getLocalThreadgroupSize(int *x, int *y, int *z);
+
 	static SourceInfo getSourceInfo(const std::string &src);
 	static std::string createShaderStageCode(Graphics *gfx, ShaderStageType stage, const std::string &code, const SourceInfo &info);
 
@@ -243,6 +245,7 @@ protected:
 	struct ValidationReflection
 	{
 		std::map<std::string, BufferReflection> storageBuffers;
+		int localThreadgroupSize[3];
 		bool usesPointSize;
 	};
 
diff --git a/src/modules/graphics/wrap_Shader.cpp b/src/modules/graphics/wrap_Shader.cpp
index bd0fbd5c1..3f8434423 100644
--- a/src/modules/graphics/wrap_Shader.cpp
+++ b/src/modules/graphics/wrap_Shader.cpp
@@ -503,13 +503,32 @@ int w_Shader_hasStage(lua_State* L)
 	return 1;
 }
 
+int w_Shader_getLocalThreadgroupSize(lua_State* L)
+{
+	Shader *shader = luax_checkshader(L, 1);
+
+	if (!shader->hasStage(SHADERSTAGE_COMPUTE))
+	{
+		lua_pushnil(L);
+		return 1;
+	}
+
+	int x, y, z;
+	shader->getLocalThreadgroupSize(&x, &y, &z);
+	lua_pushinteger(L, x);
+	lua_pushinteger(L, y);
+	lua_pushinteger(L, z);
+	return 3;
+}
+
 static const luaL_Reg w_Shader_functions[] =
 {
-	{ "getWarnings", w_Shader_getWarnings },
-	{ "send",        w_Shader_send },
-	{ "sendColor",   w_Shader_sendColors },
-	{ "hasUniform",  w_Shader_hasUniform },
-	{ "hasStage",    w_Shader_hasStage },
+	{ "getWarnings",             w_Shader_getWarnings },
+	{ "send",                    w_Shader_send },
+	{ "sendColor",               w_Shader_sendColors },
+	{ "hasUniform",              w_Shader_hasUniform },
+	{ "hasStage",                w_Shader_hasStage },
+	{ "getLocalThreadgroupSize", w_Shader_getLocalThreadgroupSize },
 	{ 0, 0 }
 };