From a5cc91594826927c3e67cfc5ad5138e61e54252c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Leclercq?= <lynix680@gmail.com>
Date: Thu, 23 Dec 2021 17:36:40 +0100
Subject: [PATCH] Shader/SPIRV: Fix structs being always decorated as block

They are now only decorated as blocks when used as uniform buffers, which fixes structure nesting
---
 include/Nazara/Shader/SpirvConstantCache.hpp |  4 +++-
 src/Nazara/Shader/SpirvConstantCache.cpp     | 17 +++++++++++++++--
 src/Nazara/Shader/SpirvWriter.cpp            | 20 ++++++++++++++++++--
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/include/Nazara/Shader/SpirvConstantCache.hpp b/include/Nazara/Shader/SpirvConstantCache.hpp
index 86abefcf6..d3055cd55 100644
--- a/include/Nazara/Shader/SpirvConstantCache.hpp
+++ b/include/Nazara/Shader/SpirvConstantCache.hpp
@@ -105,6 +105,7 @@ namespace Nz
 
 				std::string name;
 				std::vector<Member> members;
+				std::vector<SpirvDecoration> decorations;
 			};
 
 			using AnyType = std::variant<Bool, Float, Function, Image, Integer, Matrix, Pointer, SampledImage, Structure, Vector, Void>;
@@ -166,6 +167,7 @@ namespace Nz
 			TypePtr BuildFunctionType(const ShaderAst::ExpressionType& retType, const std::vector<ShaderAst::ExpressionType>& parameters) const;
 			TypePtr BuildPointerType(const ShaderAst::PrimitiveType& type, SpirvStorageClass storageClass) const;
 			TypePtr BuildPointerType(const ShaderAst::ExpressionType& type, SpirvStorageClass storageClass) const;
+			TypePtr BuildPointerType(const TypePtr& type, SpirvStorageClass storageClass) const;
 			TypePtr BuildType(const ShaderAst::ExpressionType& type) const;
 			TypePtr BuildType(const ShaderAst::IdentifierType& type) const;
 			TypePtr BuildType(const ShaderAst::MatrixType& type) const;
@@ -173,7 +175,7 @@ namespace Nz
 			TypePtr BuildType(const ShaderAst::PrimitiveType& type) const;
 			TypePtr BuildType(const ShaderAst::SamplerType& type) const;
 			TypePtr BuildType(const ShaderAst::StructType& type) const;
-			TypePtr BuildType(const ShaderAst::StructDescription& structDesc) const;
+			TypePtr BuildType(const ShaderAst::StructDescription& structDesc, std::vector<SpirvDecoration> decorations = {}) const;
 			TypePtr BuildType(const ShaderAst::VectorType& type) const;
 			TypePtr BuildType(const ShaderAst::UniformType& type) const;
 
diff --git a/src/Nazara/Shader/SpirvConstantCache.cpp b/src/Nazara/Shader/SpirvConstantCache.cpp
index 0dab16799..53441fb51 100644
--- a/src/Nazara/Shader/SpirvConstantCache.cpp
+++ b/src/Nazara/Shader/SpirvConstantCache.cpp
@@ -87,6 +87,9 @@ namespace Nz
 			if (lhs.name != rhs.name)
 				return false;
 
+			if (lhs.decorations != rhs.decorations)
+				return false;
+
 			if (!Compare(lhs.members, rhs.members))
 				return false;
 
@@ -484,6 +487,14 @@ namespace Nz
 		});
 	}
 
+	auto SpirvConstantCache::BuildPointerType(const TypePtr& type, SpirvStorageClass storageClass) const -> TypePtr
+	{
+		return std::make_shared<Type>(Pointer{
+			type,
+			storageClass
+		});
+	}
+
 	auto SpirvConstantCache::BuildPointerType(const ShaderAst::PrimitiveType& type, SpirvStorageClass storageClass) const -> TypePtr
 	{
 		return std::make_shared<Type>(Pointer{
@@ -583,10 +594,11 @@ namespace Nz
 		return BuildType(m_internal->structCallback(type.structIndex));
 	}
 
-	auto SpirvConstantCache::BuildType(const ShaderAst::StructDescription& structDesc) const -> TypePtr
+	auto SpirvConstantCache::BuildType(const ShaderAst::StructDescription& structDesc, std::vector<SpirvDecoration> decorations) const -> TypePtr
 	{
 		Structure sType;
 		sType.name = structDesc.name;
+		sType.decorations = std::move(decorations);
 
 		for (const auto& member : structDesc.members)
 		{
@@ -864,7 +876,8 @@ namespace Nz
 
 		debugInfos.Append(SpirvOp::OpName, resultId, structData.name);
 
-		annotations.Append(SpirvOp::OpDecorate, resultId, SpirvDecoration::Block);
+		for (SpirvDecoration decoration : structData.decorations)
+			annotations.Append(SpirvOp::OpDecorate, resultId, decoration);
 
 		FieldOffsets structOffsets(StructLayout::Std140);
 
diff --git a/src/Nazara/Shader/SpirvWriter.cpp b/src/Nazara/Shader/SpirvWriter.cpp
index 81920c6ef..fff1cbfb8 100644
--- a/src/Nazara/Shader/SpirvWriter.cpp
+++ b/src/Nazara/Shader/SpirvWriter.cpp
@@ -143,8 +143,24 @@ namespace Nz
 					{
 						SpirvConstantCache::Variable variable;
 						variable.debugName = extVar.name;
-						variable.storageClass = (ShaderAst::IsSamplerType(extVar.type)) ? SpirvStorageClass::UniformConstant : SpirvStorageClass::Uniform;
-						variable.type = m_constantCache.BuildPointerType(extVar.type, variable.storageClass);
+
+						if (ShaderAst::IsSamplerType(extVar.type))
+						{
+							variable.storageClass = SpirvStorageClass::UniformConstant;
+							variable.type = m_constantCache.BuildPointerType(extVar.type, variable.storageClass);
+						}
+						else
+						{
+							assert(ShaderAst::IsUniformType(extVar.type));
+							const auto& uniformType = std::get<ShaderAst::UniformType>(extVar.type);
+							assert(std::holds_alternative<ShaderAst::StructType>(uniformType.containedType));
+							const auto& structType = std::get<ShaderAst::StructType>(uniformType.containedType);
+							assert(structType.structIndex < declaredStructs.size());
+							const auto& type = m_constantCache.BuildType(*declaredStructs[structType.structIndex], { SpirvDecoration::Block });
+
+							variable.storageClass = SpirvStorageClass::Uniform;
+							variable.type = m_constantCache.BuildPointerType(type, variable.storageClass);
+						}
 
 						assert(extVar.bindingIndex.IsResultingValue());