From fbae4136204479f6f31bad6ce25745fb76b31d5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Leclercq?= <lynix680@gmail.com>
Date: Sun, 27 Feb 2022 18:42:41 +0100
Subject: [PATCH] Shader/GlslWriter: Improve removal of unused code

---
 .../Shader/Ast/EliminateUnusedPassVisitor.hpp | 11 ++++++++--
 .../Shader/Ast/EliminateUnusedPassVisitor.inl |  4 ++--
 .../Shader/Ast/EliminateUnusedPassVisitor.cpp | 21 ++++++++++++++++---
 src/Nazara/Shader/GlslWriter.cpp              |  6 +++++-
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.hpp b/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.hpp
index 1631b684b..854ee69df 100644
--- a/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.hpp
+++ b/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.hpp
@@ -17,16 +17,23 @@ namespace Nz::ShaderAst
 	class NAZARA_SHADER_API EliminateUnusedPassVisitor : AstCloner
 	{
 		public:
+			struct Config;
+
 			EliminateUnusedPassVisitor() = default;
 			EliminateUnusedPassVisitor(const EliminateUnusedPassVisitor&) = delete;
 			EliminateUnusedPassVisitor(EliminateUnusedPassVisitor&&) = delete;
 			~EliminateUnusedPassVisitor() = default;
 
-			StatementPtr Process(Statement& statement);
+			StatementPtr Process(Statement& statement, const Config& config = {});
 
 			EliminateUnusedPassVisitor& operator=(const EliminateUnusedPassVisitor&) = delete;
 			EliminateUnusedPassVisitor& operator=(EliminateUnusedPassVisitor&&) = delete;
 
+			struct Config
+			{
+				ShaderStageTypeFlags usedShaderStages = ShaderStageType_All;
+			};
+
 		private:
 			using AstCloner::Clone;
 			StatementPtr Clone(DeclareExternalStatement& node) override;
@@ -42,7 +49,7 @@ namespace Nz::ShaderAst
 			Context* m_context;
 	};
 
-	inline StatementPtr EliminateUnusedPass(Statement& ast);
+	inline StatementPtr EliminateUnusedPass(Statement& ast, const EliminateUnusedPassVisitor::Config& config = {});
 }
 
 #include <Nazara/Shader/Ast/EliminateUnusedPassVisitor.inl>
diff --git a/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.inl b/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.inl
index 08512f309..db8124116 100644
--- a/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.inl
+++ b/include/Nazara/Shader/Ast/EliminateUnusedPassVisitor.inl
@@ -7,10 +7,10 @@
 
 namespace Nz::ShaderAst
 {
-	inline StatementPtr EliminateUnusedPass(Statement& ast)
+	inline StatementPtr EliminateUnusedPass(Statement& ast, const EliminateUnusedPassVisitor::Config& config)
 	{
 		EliminateUnusedPassVisitor visitor;
-		return visitor.Process(ast);
+		return visitor.Process(ast, config);
 	}
 }
 
diff --git a/src/Nazara/Shader/Ast/EliminateUnusedPassVisitor.cpp b/src/Nazara/Shader/Ast/EliminateUnusedPassVisitor.cpp
index 9fccbceac..95afe355d 100644
--- a/src/Nazara/Shader/Ast/EliminateUnusedPassVisitor.cpp
+++ b/src/Nazara/Shader/Ast/EliminateUnusedPassVisitor.cpp
@@ -29,6 +29,11 @@ namespace Nz::ShaderAst
 
 		struct UsageChecker : AstRecursiveVisitor
 		{
+			UsageChecker(const EliminateUnusedPassVisitor::Config& Config) :
+			config(Config)
+			{
+			}
+
 			struct UsageSet;
 
 			void Resolve()
@@ -129,7 +134,11 @@ namespace Nz::ShaderAst
 				}
 
 				if (node.entryStage.HasValue())
-					globalUsage.usedFunctions.UnboundedSet(*node.funcIndex);
+				{
+					ShaderStageType shaderStage = node.entryStage.GetResultingValue();
+					if (shaderStage & config.usedShaderStages)
+						globalUsage.usedFunctions.UnboundedSet(*node.funcIndex);
+				}
 
 				currentFunctionIndex = node.funcIndex;
 				AstRecursiveVisitor::Visit(node);
@@ -195,6 +204,7 @@ namespace Nz::ShaderAst
 				Bitset<> usedVariables;
 			};
 
+			const EliminateUnusedPassVisitor::Config& config;
 			std::optional<std::size_t> currentFunctionIndex;
 			std::optional<std::size_t> currentVariableDeclIndex;
 			std::unordered_map<std::size_t, UsageSet> functionUsages;
@@ -207,12 +217,17 @@ namespace Nz::ShaderAst
 
 	struct EliminateUnusedPassVisitor::Context
 	{
+		Context(const Config& config) :
+		usageChecker(config)
+		{
+		}
+
 		UsageChecker usageChecker;
 	};
 
-	StatementPtr EliminateUnusedPassVisitor::Process(Statement& statement)
+	StatementPtr EliminateUnusedPassVisitor::Process(Statement& statement, const Config& config)
 	{
-		Context context;
+		Context context(config);
 		statement.Visit(context.usageChecker);
 		context.usageChecker.Resolve();
 
diff --git a/src/Nazara/Shader/GlslWriter.cpp b/src/Nazara/Shader/GlslWriter.cpp
index 5b05f1f09..6a60c6991 100644
--- a/src/Nazara/Shader/GlslWriter.cpp
+++ b/src/Nazara/Shader/GlslWriter.cpp
@@ -180,8 +180,12 @@ namespace Nz
 		{
 			ShaderAst::StatementPtr tempAst;
 
+			ShaderAst::EliminateUnusedPassVisitor::Config eliminateUnunsedConfig;
+			if (shaderStage)
+				eliminateUnunsedConfig.usedShaderStages = *shaderStage;
+
 			tempAst = ShaderAst::PropagateConstants(*targetAst);
-			optimizedAst = ShaderAst::EliminateUnusedPass(*tempAst);
+			optimizedAst = ShaderAst::EliminateUnusedPass(*tempAst, eliminateUnunsedConfig);
 
 			targetAst = optimizedAst.get();
 		}