From 77973a48acbe29c2a64d43de7084574aaf7a26c4 Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Sat, 12 Mar 2022 00:11:49 +0200
Subject: [PATCH] Remove unused stages from shader programs

A stage is deemed unused if it has no main function after optimization,
or if a SPIR-V entry point's interface variables all become unused during
specialization.
---
 source/core/module.cpp   | 20 ++++++++++++++++++--
 source/glsl/compiler.cpp | 13 +++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/source/core/module.cpp b/source/core/module.cpp
index 11d11af5..f42a0da7 100644
--- a/source/core/module.cpp
+++ b/source/core/module.cpp
@@ -14,6 +14,7 @@ enum SpirVConstants
 	OP_NAME = 5,
 	OP_MEMBER_NAME = 6,
 	OP_ENTRY_POINT = 15,
+	OP_EXECUTION_MODE = 16,
 	OP_TYPE_VOID = 19,
 	OP_TYPE_BOOL = 20,
 	OP_TYPE_INT = 21,
@@ -269,6 +270,8 @@ SpirVModule *SpirVModule::specialize(const map<string, int> &spec_values) const
 		}
 	}
 
+	for(const EntryPoint &e: entry_points)
+		flags[e.id] = 0;
 	for(const Variable &v: variables)
 		flags[v.id] = 0;
 	for(const InstructionBlock &b: blocks)
@@ -307,7 +310,8 @@ SpirVModule *SpirVModule::specialize(const map<string, int> &spec_values) const
 				unsigned start = new_code.size();
 				new_code.push_back(opcode);
 				new_code.push_back(*(op+1));
-				new_code.push_back(*(op+2));
+				unsigned func_id = *(op+2);
+				new_code.push_back(func_id);
 
 				unsigned i=3;
 				while(i<word_count)
@@ -318,17 +322,29 @@ SpirVModule *SpirVModule::specialize(const map<string, int> &spec_values) const
 						break;
 				}
 
+				unsigned var_count = 0;
 				for(; i<word_count; ++i)
 				{
 					unsigned id = *(op+i);
 					if(flags[id])
+					{
+						++var_count;
 						new_code.push_back(id);
+					}
 				}
 
-				new_code[start] |= (new_code.size()-start)<<16;
+				if(var_count)
+				{
+					flags[func_id] = 1;
+					new_code[start] |= (new_code.size()-start)<<16;
+				}
+				else
+					new_code.resize(start);
 
 				copy = false;
 			}
+			else if(opcode==OP_EXECUTION_MODE)
+				copy = flags[*(op+1)];
 			else if(opcode==OP_SPEC_CONSTANT_TRUE || opcode==OP_SPEC_CONSTANT_FALSE)
 			{
 				unsigned id = *(op+2);
diff --git a/source/glsl/compiler.cpp b/source/glsl/compiler.cpp
index 4c8c1d03..05bf421c 100644
--- a/source/glsl/compiler.cpp
+++ b/source/glsl/compiler.cpp
@@ -107,6 +107,19 @@ void Compiler::compile(Mode mode)
 			++i;
 	}
 
+	Stage *prev_stage = 0;
+	for(auto i=module->stages.begin(); i!=module->stages.end(); )
+	{
+		if(i->functions.empty())
+			i = module->stages.erase(i);
+		else
+		{
+			i->previous = prev_stage;
+			prev_stage = &*i;
+			++i;
+		}
+	}
+
 	for(Stage &s: module->stages)
 	{
 		StructuralFeatureConverter().apply(s, features);
-- 
2.45.2