From 61d082e61ba0e313c4811e69766cedc9ab7113a8 Mon Sep 17 00:00:00 2001
From: Lynix <lynix680@gmail.com>
Date: Sun, 16 May 2021 23:08:36 +0200
Subject: [PATCH] Shader/ShaderWrite: Add optimize bool

---
 examples/RenderTest/main.cpp               |  1 +
 include/Nazara/Shader/Ast/AstOptimizer.hpp |  3 +++
 include/Nazara/Shader/Ast/AstOptimizer.inl | 13 ++++++++++++-
 include/Nazara/Shader/ShaderWriter.hpp     |  1 +
 src/Nazara/Shader/GlslWriter.cpp           | 20 ++++++++++++++++----
 src/Nazara/Shader/SpirvWriter.cpp          | 15 ++++++++++++++-
 6 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/examples/RenderTest/main.cpp b/examples/RenderTest/main.cpp
index 4e8fff1cc..89986304b 100644
--- a/examples/RenderTest/main.cpp
+++ b/examples/RenderTest/main.cpp
@@ -101,6 +101,7 @@ int main()
 
 	Nz::ShaderWriter::States states;
 	states.enabledOptions = 0;
+	states.optimize = true;
 
 	auto fragVertShader = device->InstantiateShaderModule(Nz::ShaderStageType::Fragment | Nz::ShaderStageType::Vertex, Nz::ShaderLanguage::NazaraShader, shaderSource, sizeof(shaderSource), states);
 	if (!fragVertShader)
diff --git a/include/Nazara/Shader/Ast/AstOptimizer.hpp b/include/Nazara/Shader/Ast/AstOptimizer.hpp
index 8a60bce05..967d0cec9 100644
--- a/include/Nazara/Shader/Ast/AstOptimizer.hpp
+++ b/include/Nazara/Shader/Ast/AstOptimizer.hpp
@@ -42,6 +42,9 @@ namespace Nz::ShaderAst
 		private:
 			std::optional<UInt64> m_enabledOptions;
 	};
+
+	inline StatementPtr Optimize(const StatementPtr& ast);
+	inline StatementPtr Optimize(const StatementPtr& ast, UInt64 enabledConditions);
 }
 
 #include <Nazara/Shader/Ast/AstOptimizer.inl>
diff --git a/include/Nazara/Shader/Ast/AstOptimizer.inl b/include/Nazara/Shader/Ast/AstOptimizer.inl
index 9607f7af0..78b79324d 100644
--- a/include/Nazara/Shader/Ast/AstOptimizer.inl
+++ b/include/Nazara/Shader/Ast/AstOptimizer.inl
@@ -5,8 +5,19 @@
 #include <Nazara/Shader/Ast/AstOptimizer.hpp>
 #include <Nazara/Shader/Debug.hpp>
 
-namespace Nz
+namespace Nz::ShaderAst
 {
+	inline StatementPtr Optimize(const StatementPtr& ast)
+	{
+		AstOptimizer optimize;
+		return optimize.Optimise(ast);
+	}
+
+	inline StatementPtr Optimize(const StatementPtr& ast, UInt64 enabledConditions)
+	{
+		AstOptimizer optimize;
+		return optimize.Optimise(ast, enabledConditions);
+	}
 }
 
 #include <Nazara/Shader/DebugOff.hpp>
diff --git a/include/Nazara/Shader/ShaderWriter.hpp b/include/Nazara/Shader/ShaderWriter.hpp
index 5433e0842..e50bd8d50 100644
--- a/include/Nazara/Shader/ShaderWriter.hpp
+++ b/include/Nazara/Shader/ShaderWriter.hpp
@@ -27,6 +27,7 @@ namespace Nz
 			struct States
 			{
 				Nz::UInt64 enabledOptions = 0;
+				bool optimize = false;
 				bool sanitized = false;
 			};
 	};
diff --git a/src/Nazara/Shader/GlslWriter.cpp b/src/Nazara/Shader/GlslWriter.cpp
index c3a419fd1..13b952818 100644
--- a/src/Nazara/Shader/GlslWriter.cpp
+++ b/src/Nazara/Shader/GlslWriter.cpp
@@ -8,6 +8,7 @@
 #include <Nazara/Math/Algorithm.hpp>
 #include <Nazara/Shader/ShaderBuilder.hpp>
 #include <Nazara/Shader/Ast/AstCloner.hpp>
+#include <Nazara/Shader/Ast/AstOptimizer.hpp>
 #include <Nazara/Shader/Ast/AstRecursiveVisitor.hpp>
 #include <Nazara/Shader/Ast/AstUtils.hpp>
 #include <Nazara/Shader/Ast/SanitizeVisitor.hpp>
@@ -122,10 +123,21 @@ namespace Nz
 
 		ShaderAst::StatementPtr sanitizedAst = ShaderAst::Sanitize(shader, options);
 
+		ShaderAst::StatementPtr* targetAstPtr = &sanitizedAst;
+
+		ShaderAst::StatementPtr optimizedAst;
+		if (states.optimize)
+		{
+			optimizedAst = ShaderAst::Optimize(*targetAstPtr);
+			targetAstPtr = &optimizedAst;
+		}
+
+		ShaderAst::StatementPtr& targetAst = *targetAstPtr;
+
 		PreVisitor previsitor;
 		previsitor.enabledOptions = states.enabledOptions;
 		previsitor.selectedStage = shaderStage;
-		sanitizedAst->Visit(previsitor);
+		targetAst->Visit(previsitor);
 
 		if (!previsitor.entryPoint)
 			throw std::runtime_error("missing entry point");
@@ -167,7 +179,7 @@ namespace Nz
 		}
 	}
 
-	void GlslWriter::Append(const ShaderAst::IdentifierType& identifierType)
+	void GlslWriter::Append(const ShaderAst::IdentifierType& /*identifierType*/)
 	{
 		throw std::runtime_error("unexpected identifier type");
 	}
@@ -230,7 +242,7 @@ namespace Nz
 		Append(structDesc.name);
 	}
 
-	void GlslWriter::Append(const ShaderAst::UniformType& uniformType)
+	void GlslWriter::Append(const ShaderAst::UniformType& /*uniformType*/)
 	{
 		throw std::runtime_error("unexpected UniformType");
 	}
@@ -734,7 +746,7 @@ namespace Nz
 		LeaveScope();
 	}
 
-	void GlslWriter::Visit(ShaderAst::DeclareOptionStatement& node)
+	void GlslWriter::Visit(ShaderAst::DeclareOptionStatement& /*node*/)
 	{
 		/* nothing to do */
 	}
diff --git a/src/Nazara/Shader/SpirvWriter.cpp b/src/Nazara/Shader/SpirvWriter.cpp
index 5bed57c67..26d82c7d1 100644
--- a/src/Nazara/Shader/SpirvWriter.cpp
+++ b/src/Nazara/Shader/SpirvWriter.cpp
@@ -11,6 +11,7 @@
 #include <Nazara/Shader/SpirvData.hpp>
 #include <Nazara/Shader/SpirvSection.hpp>
 #include <Nazara/Shader/Ast/AstCloner.hpp>
+#include <Nazara/Shader/Ast/AstOptimizer.hpp>
 #include <Nazara/Shader/Ast/AstRecursiveVisitor.hpp>
 #include <Nazara/Shader/Ast/SanitizeVisitor.hpp>
 #include <tsl/ordered_map.h>
@@ -401,11 +402,23 @@ namespace Nz
 
 	std::vector<UInt32> SpirvWriter::Generate(ShaderAst::StatementPtr& shader, const States& states)
 	{
+		ShaderAst::StatementPtr* targetAstPtr = &shader;
+
 		ShaderAst::StatementPtr sanitizedAst;
 		if (!states.sanitized)
+		{
 			sanitizedAst = ShaderAst::Sanitize(shader);
+			targetAstPtr = &sanitizedAst;
+		}
 
-		ShaderAst::StatementPtr& targetAst = (states.sanitized) ? shader : sanitizedAst;
+		ShaderAst::StatementPtr optimizedAst;
+		if (states.optimize)
+		{
+			optimizedAst = ShaderAst::Optimize(*targetAstPtr);
+			targetAstPtr = &optimizedAst;
+		}
+
+		ShaderAst::StatementPtr& targetAst = *targetAstPtr;
 
 		m_context.states = &states;