From 83deecd8f1983ea1c17578126830722a6db13f37 Mon Sep 17 00:00:00 2001
From: SirLynix <lynix680@gmail.com>
Date: Fri, 25 Mar 2022 12:55:32 +0100
Subject: [PATCH] Shader: Fix SPIRV shaders terminating before function last
 statement

---
 include/Nazara/Shader/SpirvAstVisitor.hpp |  2 ++
 src/Nazara/Shader/SpirvAstVisitor.cpp     | 25 +++++++++++++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/include/Nazara/Shader/SpirvAstVisitor.hpp b/include/Nazara/Shader/SpirvAstVisitor.hpp
index 967dccb89..580fbb1d2 100644
--- a/include/Nazara/Shader/SpirvAstVisitor.hpp
+++ b/include/Nazara/Shader/SpirvAstVisitor.hpp
@@ -138,6 +138,8 @@ namespace Nz
 			};
 
 		private:
+			void HandleStatementList(const std::vector<ShaderAst::StatementPtr>& statements);
+
 			void PushResultId(UInt32 value);
 			UInt32 PopResultId();
 
diff --git a/src/Nazara/Shader/SpirvAstVisitor.cpp b/src/Nazara/Shader/SpirvAstVisitor.cpp
index 69b549ff1..d7a7c4d62 100644
--- a/src/Nazara/Shader/SpirvAstVisitor.cpp
+++ b/src/Nazara/Shader/SpirvAstVisitor.cpp
@@ -655,8 +655,7 @@ namespace Nz
 			}
 		}
 
-		for (auto& statementPtr : node.statements)
-			statementPtr->Visit(*this);
+		HandleStatementList(node.statements);
 
 		// Add implicit return
 		if (!m_functionBlocks.back()->IsTerminated())
@@ -920,8 +919,7 @@ namespace Nz
 
 	void SpirvAstVisitor::Visit(ShaderAst::MultiStatement& node)
 	{
-		for (auto& statement : node.statements)
-			statement->Visit(*this);
+		HandleStatementList(node.statements);
 	}
 
 	void SpirvAstVisitor::Visit(ShaderAst::ReturnStatement& node)
@@ -1155,6 +1153,25 @@ namespace Nz
 		m_currentBlock = m_functionBlocks.back().get();
 	}
 
+	void SpirvAstVisitor::HandleStatementList(const std::vector<ShaderAst::StatementPtr>& statements)
+	{
+		for (auto& statement : statements)
+		{
+			// Handle termination statements
+			switch (statement->GetType())
+			{
+				case ShaderAst::NodeType::DiscardStatement:
+				case ShaderAst::NodeType::ReturnStatement:
+					statement->Visit(*this);
+					return; //< stop processing statements after this one
+
+				default:
+					statement->Visit(*this);
+					break;
+			}
+		}
+	}
+
 	void SpirvAstVisitor::PushResultId(UInt32 value)
 	{
 		m_resultIds.push_back(value);