nzslc: Add support for struct reflection

bin/resources/Reflection.nzsl

@@ -0,0 +1,15 @@
+[nzsl_version("1.1")]
+[author("SirLynix"), desc("Test module")]
+[license("MIT")]
+module Shader;
+
+import * from DataStruct;
+
+[layout(std140)]
+struct Output
+{
+    color: vec4[f32],
+    normal: vec3[f32],
+    roughness: f32,
+    metalness: f32
+}

bin/resources/Reflection.nzsl.json

@@ -0,0 +1,30 @@
+{
+	"structs": [
+		{
+			"name": "Output",
+			"layout": "std140",
+			"members": [
+				{
+					"name": "color",
+					"type": "vec4[f32]",
+					"offset": 0
+				},
+				{
+					"name": "normal",
+					"type": "vec3[f32]",
+					"offset": 16
+				},
+				{
+					"name": "roughness",
+					"type": "f32",
+					"offset": 28
+				},
+				{
+					"name": "metalness",
+					"type": "f32",
+					"offset": 32
+				}
+			]
+		}
+	]
+}

src/ShaderCompiler/Compiler.cpp

@@ -13,6 +13,7 @@
 #include <NZSL/Lang/Errors.hpp>
 #include <NZSL/Lexer.hpp>
 #include <NZSL/Parser.hpp>
+#include <NZSL/Math/FieldOffsets.hpp>
 #include <NZSL/SpirV/SpirvPrinter.hpp>
 #include <NZSL/SpirvWriter.hpp>
 #include <NZSL/Serializer.hpp>
@@ -30,6 +31,7 @@
 #include <chrono>
 #include <fstream>
 #include <stdexcept>
+#include <unordered_set>
 
 namespace nzslc
 {
@@ -201,6 +203,9 @@
 
 				if (m_options.count("compile") > 0)
 					Step("Compiling"sv, __LINE__, &Compiler::Compile);
+
+				if (m_options.count("reflect") > 0)
+					Step("Reflecting"sv, __LINE__, &Compiler::Reflect);
 			});
 		}
 
@@ -240,6 +245,9 @@
 			("p,partial", "Allow partial compilation")
 			("skip-unchanged", "After compilation, compare the output with the current output file and skip writing if the content is the same", cxxopts::value<bool>()->default_value("false"));
 
+		options.add_options("reflection")
+			("r,reflect", "Outputs informations about a struct as a json", cxxopts::value<std::vector<std::string>>());
+
 		options.add_options("glsl output")
 			("gl-es", "Generate GLSL ES instead of GLSL", cxxopts::value<bool>()->default_value("false"))
 			("gl-version", "OpenGL version (310 being 3.1)", cxxopts::value<std::uint32_t>(), "version")
@@ -690,12 +698,120 @@
 			throw std::runtime_error(fmt::format("{} has unknown extension \"{}\"", Nz::PathToString(m_inputFilePath.filename()), Nz::PathToString(extension)));
 	}
 
+	void Compiler::Reflect()
+	{
+		m_outputHeader = false;
+
+		// if no output path has been provided, output in the same folder as the input file
+		std::filesystem::path outputFilePath = m_outputPath;
+		if (outputFilePath.empty())
+			outputFilePath = m_inputFilePath.parent_path();
+
+		outputFilePath /= m_inputFilePath.filename();
+		outputFilePath += Nz::Utf8Path(".json");
+
+		const std::vector<std::string>& reflectTypes = m_options["reflect"].as<std::vector<std::string>>();
+
+		std::unordered_set<std::string> remainingStructs(reflectTypes.begin(), reflectTypes.end());
+
+		nlohmann::ordered_json structArray = nlohmann::ordered_json::array();
+
+		nzsl::Ast::ReflectVisitor::Callbacks callbacks;
+		callbacks.onStructDeclaration = [&](const nzsl::Ast::DeclareStructStatement& structDecl)
+		{
+			auto it = remainingStructs.find(structDecl.description.name);
+			if (it == remainingStructs.end())
+				return;
+
+			remainingStructs.erase(it);
+
+			nlohmann::ordered_json structDoc;
+			structDoc["name"] = structDecl.description.name;
+			if (!structDecl.description.tag.empty())
+				structDoc["tag"] = structDecl.description.tag;
+
+			nlohmann::ordered_json structMemberArray = nlohmann::ordered_json::array();
+
+			std::optional<nzsl::FieldOffsets> fieldOffsets;
+			if (structDecl.description.layout.IsResultingValue())
+			{
+				switch (structDecl.description.layout.GetResultingValue())
+				{
+					case nzsl::Ast::MemoryLayout::Scalar:
+						structDoc["layout"] = "scalar";
+						fieldOffsets.emplace(nzsl::StructLayout::Scalar);
+						break;
+
+					case nzsl::Ast::MemoryLayout::Std140:
+						structDoc["layout"] = "std140";
+						fieldOffsets.emplace(nzsl::StructLayout::Std140);
+						break;
+
+					case nzsl::Ast::MemoryLayout::Std430:
+						structDoc["layout"] = "std430";
+						fieldOffsets.emplace(nzsl::StructLayout::Std430);
+						break;
+				}
+			}
+			else if (structDecl.description.layout.IsExpression())
+				structDoc["layout"] = "unresolved";
+
+			for (const auto& member : structDecl.description.members)
+			{
+				nlohmann::ordered_json memberDoc;
+				memberDoc["name"] = member.name;
+
+				if (member.type.IsResultingValue())
+				{
+					memberDoc["type"] = nzsl::Ast::ToString(member.type.GetResultingValue());
+					if (fieldOffsets)
+						memberDoc["offset"] = nzsl::Ast::RegisterStructField(*fieldOffsets, member.type.GetResultingValue());
+				}
+				else if (member.type.IsExpression())
+					memberDoc["type"] = "unresolved";
+
+				if (member.locationIndex.IsResultingValue())
+					memberDoc["location"] = member.locationIndex.GetResultingValue();
+				else if (member.locationIndex.IsExpression())
+					memberDoc["location"] = "unresolved";
+
+				structMemberArray.push_back(std::move(memberDoc));
+			}
+
+			structDoc["members"] = std::move(structMemberArray);
+
+			structArray.push_back(std::move(structDoc));
+
+			// TODO: Stop visit if remainingStructs.empty()
+		};
+
+		nzsl::Ast::ReflectVisitor reflectVisitor;
+		reflectVisitor.Reflect(*m_shaderModule, callbacks);
+
+		if (!remainingStructs.empty())
+			throw std::runtime_error(fmt::format("struct \"{}\" was not found", *remainingStructs.begin()));
+
+		nlohmann::ordered_json result;
+		result["structs"] = std::move(structArray);
+
+		if (m_skipOutput)
+			return;
+
+		if (m_outputToStdout)
+		{
+			OutputToStdout(result.dump(1, '\t'));
+			return;
+		}
+
+		std::string output = result.dump(1, '\t');
+		OutputFile(std::move(outputFilePath), output.data(), output.size());
+	}
+
 	void Compiler::Resolve()
 	{
 		using namespace std::literals;
 
-		nzsl::Ast::TransformerContext context;
-		context.partialCompilation = m_options.count("partial") > 0;
+		m_transformerContext.partialCompilation = m_options.count("partial") > 0;
 
 		nzsl::Ast::ResolveTransformer::Options resolverOpt;
 
@@ -726,8 +842,8 @@
 		nzsl::Ast::ResolveTransformer resolver;
 		nzsl::Ast::ValidationTransformer validation;
 
-		Step("AST processing"sv, __LINE__, [&] { resolver.Transform(*m_shaderModule, context, resolverOpt); });
-		Step("AST validation"sv, __LINE__, [&] { validation.Transform(*m_shaderModule, context); });
+		Step("AST processing"sv, __LINE__, [&] { resolver.Transform(*m_shaderModule, m_transformerContext, resolverOpt); });
+		Step("AST validation"sv, __LINE__, [&] { validation.Transform(*m_shaderModule, m_transformerContext); });
 	}
 
 	template<typename F, typename... Args>

src/ShaderCompiler/Compiler.hpp

@@ -10,6 +10,7 @@
 #include <NZSL/Config.hpp>
 #include <NZSL/Lang/Errors.hpp>
 #include <NZSL/Ast/Module.hpp>
+#include <NZSL/Ast/Transformations/TransformerContext.hpp>
 #include <cxxopts.hpp>
 #include <filesystem>
 #include <type_traits>
@@ -64,6 +65,7 @@ namespace nzslc
 			void OutputFile(std::filesystem::path filePath, const void* data, std::size_t size, bool disallowHeader = false);
 			void OutputToStdout(std::string_view str);
 			void ReadInput();
+			void Reflect();
 			void Resolve();
 			template<typename F, typename... Args> auto Step(std::enable_if_t<!std::is_member_function_pointer_v<F>, std::string_view> stepName, std::size_t uniqueIndex, F&& func, Args&&... args) -> decltype(std::invoke(func, std::forward<Args>(args)...));
 			template<typename F, typename... Args> auto Step(std::enable_if_t<std::is_member_function_pointer_v<F>, std::string_view> stepName, std::size_t uniqueIndex, F&& func, Args&&... args) -> decltype(std::invoke(func, this, std::forward<Args>(args)...));
@@ -88,6 +90,7 @@ namespace nzslc
 			std::vector<StepTime> m_steps;
 			LogFormat m_logFormat;
 			nzsl::Ast::ModulePtr m_shaderModule;
+			nzsl::Ast::TransformerContext m_transformerContext;
 			cxxopts::ParseResult& m_options;
 			bool m_isProfiling;
 			bool m_outputHeader;

tests/src/Tests/NzslcTests.cpp

@@ -59,4 +59,22 @@ TEST_CASE("Standalone compiler", "[NZSLC]")
 		// Generate the same shader a second time with --skip-unchanged and ensure file wasn't modified
 		ExecuteCommand("./nzslc --skip-unchanged --verbose --compile=spv --debug-level=regular -o test_files -m ../resources/modules/Color.nzslb  -m ../resources/modules/Data/OutputStruct.nzslb -m ../resources/modules/Data/DataStruct.nzslb ../resources/Shader.nzslb", "Skipped file .+Shader.spv");
 	}
+
+	WHEN("Performing reflection")
+	{
+		REQUIRE(std::filesystem::exists("../resources/Reflection.nzsl"));
+
+		auto Cleanup = []
+		{
+			if (std::filesystem::is_directory("test_files"))
+				std::filesystem::remove_all("test_files");
+		};
+
+		Cleanup();
+
+		Nz::CallOnExit cleanupOnExit(std::move(Cleanup));
+
+		ExecuteCommand("./nzslc --verbose --reflect=Output --partial -o test_files ../resources/Reflection.nzsl");
+		CheckFileMatch("../resources/Reflection.nzsl.json", "test_files/Reflection.nzsl.json");
+	}
 }