More shader decompiler work

include/PICA/shader.hpp

@@ -1,6 +1,8 @@
 #pragma once
 #include <algorithm>
 #include <array>
+#include <cassert>
+#include <cstddef>
 #include <cstring>
 
 #include "PICA/float_types.hpp"
@@ -90,9 +92,12 @@ class PICAShader {
   public:
 	// These are placed close to the temp registers and co because it helps the JIT generate better code
 	u32 entrypoint = 0;  // Initial shader PC
-	u32 boolUniform;
+
-	std::array<std::array<u8, 4>, 4> intUniforms;
+	// We want these registers in this order & with this alignment for uploading them directly to a UBO
+	// When emulating shaders on the GPU
 	alignas(16) std::array<vec4f, 96> floatUniforms;
+	alignas(16) std::array<std::array<u8, 4>, 4> intUniforms;
+	u32 boolUniform;
 
 	alignas(16) std::array<vec4f, 16> fixedAttributes;  // Fixed vertex attributes
 	alignas(16) std::array<vec4f, 16> inputs;           // Attributes passed to the shader
@@ -291,4 +296,9 @@ class PICAShader {
 
 	Hash getCodeHash();
 	Hash getOpdescHash();
-};
+};
+
+static_assert(
+	offsetof(PICAShader, intUniforms) == offsetof(PICAShader, floatUniforms) + 96 * sizeof(float) * 4 &&
+	offsetof(PICAShader, boolUniform) == offsetof(PICAShader, intUniforms) + 4 * sizeof(u8) * 4
+);

include/PICA/shader_decompiler.hpp

@@ -2,6 +2,7 @@
 #include <set>
 #include <string>
 #include <tuple>
+#include <map>
 #include <vector>
 
 #include "PICA/shader.hpp"
@@ -13,6 +14,15 @@ namespace PICA::ShaderGen {
 	// Control flow analysis is partially based on
 	// https://github.com/PabloMK7/citra/blob/d0179559466ff09731d74474322ee880fbb44b00/src/video_core/shader/generator/glsl_shader_decompiler.cpp#L33
 	struct ControlFlow {
+		// A continuous range of addresses
+		struct AddressRange {
+			u32 start, end;
+			AddressRange(u32 start, u32 end) : start(start), end(end) {}
+
+			// Use lexicographic comparison for functions in order to sort them in a set
+			bool operator<(const AddressRange& other) const { return std::tie(start, end) < std::tie(other.start, other.end); }
+		};
+
 		struct Function {
 			using Labels = std::set<u32>;
 
@@ -29,20 +39,22 @@ namespace PICA::ShaderGen {
 			ExitMode exitMode = ExitMode::Unknown;
 
 			explicit Function(u32 start, u32 end) : start(start), end(end) {}
-			// Use lexicographic comparison for functions in order to sort them in a set
+			bool operator<(const Function& other) const { return AddressRange(start, end) < AddressRange(other.start, other.end); }
-			bool operator<(const Function& other) const { return std::tie(start, end) < std::tie(other.start, other.end); }
+
+			std::string getIdentifier() const { return "func_" + std::to_string(start) + "_to_" + std::to_string(end); }
+			std::string getForwardDecl() const { return "void " + getIdentifier() + "();\n"; }
+			std::string getCallStatement() const { return getIdentifier() + "()"; }
 		};
 
 		std::set<Function> functions{};
+		std::map<AddressRange, Function::ExitMode> exitMap{};
 
 		// Tells us whether analysis of the shader we're trying to compile failed, in which case we'll need to fail back to shader emulation
 		// On the CPU
 		bool analysisFailed = false;
 
-		void analyze(const PICAShader& shader, u32 entrypoint);
-
 		// This will recursively add all functions called by the function too, as analyzeFunction will call addFunction on control flow instructions
-		const Function* addFunction(u32 start, u32 end) {
+		const Function* addFunction(const PICAShader& shader, u32 start, u32 end) {
 			auto searchIterator = functions.find(Function(start, end));
 			if (searchIterator != functions.end()) {
 				return &(*searchIterator);
@@ -50,9 +62,9 @@ namespace PICA::ShaderGen {
 
 			// Add this function and analyze it if it doesn't already exist
 			Function function(start, end);
-			function.exitMode = analyzeFunction(start, end, function.outLabels);
+			function.exitMode = analyzeFunction(shader, start, end, function.outLabels);
 			
-			// This function 
+			// This function could not be fully analyzed, report failure
 			if (function.exitMode == Function::ExitMode::Unknown) {
 				analysisFailed = true;
 				return nullptr;
@@ -63,10 +75,14 @@ namespace PICA::ShaderGen {
 			return &(*it);
 		}
 
-		Function::ExitMode analyzeFunction(u32 start, u32 end, Function::Labels& labels);
+		void analyze(const PICAShader& shader, u32 entrypoint);
+		Function::ExitMode analyzeFunction(const PICAShader& shader, u32 start, u32 end, Function::Labels& labels);
 	};
 
 	class ShaderDecompiler {
+		using AddressRange = ControlFlow::AddressRange;
+		using Function = ControlFlow::Function;
+
 		ControlFlow controlFlow{};
 
 		PICAShader& shader;
@@ -74,14 +90,20 @@ namespace PICA::ShaderGen {
 		std::string decompiledShader;
 
 		u32 entrypoint;
-		u32 currentPC;
 
 		API api;
 		Language language;
 
+		void compileInstruction(u32& pc, bool& finished);
+		void compileRange(const AddressRange& range);
+		void callFunction(const Function& function);
+		const Function* findFunction(const AddressRange& range);
+
+		void writeAttributes();
+
 	  public:
 		ShaderDecompiler(PICAShader& shader, EmulatorConfig& config, u32 entrypoint, API api, Language language)
-			: shader(shader), entrypoint(entrypoint), currentPC(entrypoint), config(config), api(api), language(language), decompiledShader("") {}
+			: shader(shader), entrypoint(entrypoint), config(config), api(api), language(language), decompiledShader("") {}
 
 		std::string decompile();
 	};

src/core/PICA/shader_decompiler.cpp

@@ -10,13 +10,75 @@ using ExitMode = Function::ExitMode;
 void ControlFlow::analyze(const PICAShader& shader, u32 entrypoint) {
 	analysisFailed = false;
 
-	const Function* function = addFunction(entrypoint, PICAShader::maxInstructionCount);
+	const Function* function = addFunction(shader, entrypoint, PICAShader::maxInstructionCount);
 	if (function == nullptr) {
 		analysisFailed = true;
 	}
 }
 
-ExitMode analyzeFunction(u32 start, u32 end, Function::Labels& labels) { return ExitMode::Unknown; }
+ExitMode ControlFlow::analyzeFunction(const PICAShader& shader, u32 start, u32 end, Function::Labels& labels) {
+	// Initialize exit mode to unknown by default, in order to detect things like unending loops
+	auto [it, inserted] = exitMap.emplace(AddressRange(start, end), ExitMode::Unknown);
+	// Function has already been analyzed and is in the map so it wasn't added, don't analyze again
+	if (!inserted) {
+		return it->second;
+	}
+
+	// Make sure not to go out of bounds on the shader
+	for (u32 pc = start; pc < PICAShader::maxInstructionCount && pc != end; pc++) {
+		const u32 instruction = shader.loadedShader[pc];
+		const u32 opcode = instruction >> 26;
+
+		switch (opcode) {
+			case ShaderOpcodes::JMPC: Helpers::panic("Unimplemented control flow operation (JMPC)");
+			case ShaderOpcodes::JMPU: Helpers::panic("Unimplemented control flow operation (JMPU)");
+			case ShaderOpcodes::IFU: Helpers::panic("Unimplemented control flow operation (IFU)");
+			case ShaderOpcodes::IFC: Helpers::panic("Unimplemented control flow operation (IFC)");
+			case ShaderOpcodes::CALL: Helpers::panic("Unimplemented control flow operation (CALL)");
+			case ShaderOpcodes::CALLC: Helpers::panic("Unimplemented control flow operation (CALLC)");
+			case ShaderOpcodes::CALLU: Helpers::panic("Unimplemented control flow operation (CALLU)");
+			case ShaderOpcodes::LOOP: Helpers::panic("Unimplemented control flow operation (LOOP)");
+			case ShaderOpcodes::END: it->second = ExitMode::AlwaysEnd; return it->second;
+
+			default: break;
+		}
+	}
+
+	// A function without control flow instructions will always reach its "return point" and return
+	return ExitMode::AlwaysReturn;
+}
+
+void ShaderDecompiler::compileRange(const AddressRange& range) {
+	u32 pc = range.start;
+	const u32 end = range.end >= range.start ? range.end : PICAShader::maxInstructionCount;
+	bool finished = false;
+
+	while (pc < end && !finished) {
+		compileInstruction(pc, finished);
+	}
+}
+
+const Function* ShaderDecompiler::findFunction(const AddressRange& range) {
+	for (const Function& func : controlFlow.functions) {
+		if (range.start == func.start && range.end == func.end) {
+			return &func;
+		}
+	}
+
+	return nullptr;
+}
+
+void ShaderDecompiler::writeAttributes() {
+	decompiledShader += R"(
+		layout(std140) uniform PICAShaderUniforms {
+			vec4 uniform_float[96];
+			uvec4 uniform_int;
+			uint uniform_bool;
+		};
+)";
+
+	decompiledShader += "\n";
+}
 
 std::string ShaderDecompiler::decompile() {
 	controlFlow.analyze(shader, entrypoint);
@@ -28,11 +90,13 @@ std::string ShaderDecompiler::decompile() {
 	decompiledShader = "";
 
 	switch (api) {
-		case API::GL: decompiledShader += "#version 410 core"; break;
+		case API::GL: decompiledShader += "#version 410 core\n"; break;
-		case API::GLES: decompiledShader += "#version 300 es"; break;
+		case API::GLES: decompiledShader += "#version 300 es\n"; break;
 		default: break;
 	}
 
+	writeAttributes();
+
 	if (config.accurateShaderMul) {
 		// Safe multiplication handler from Citra: Handles the PICA's 0 * inf = 0 edge case
 		decompiledShader += R"(
@@ -43,10 +107,46 @@ std::string ShaderDecompiler::decompile() {
 		)";
 	}
 
+	// Forward declare every generated function first so that we can easily call anything from anywhere.
+	for (auto& func : controlFlow.functions) {
+		decompiledShader += func.getForwardDecl();
+	}
+
+	decompiledShader += "void pica_shader_main() {\n";
+	AddressRange mainFunctionRange(entrypoint, PICAShader::maxInstructionCount);
+	callFunction(*findFunction(mainFunctionRange));
+	decompiledShader += "}\n";
+
+	for (auto& func : controlFlow.functions) {
+		if (func.outLabels.size() > 0) {
+			Helpers::panic("Function with out labels");
+		}
+
+		decompiledShader += "void " + func.getIdentifier() + "() {\n";
+		compileRange(AddressRange(func.start, func.end));
+		decompiledShader += "}\n";
+	}
+
 	return decompiledShader;
 }
 
-std::string PICA::ShaderGen::decompileShader(PICAShader& shader, EmulatorConfig& config, u32 entrypoint, API api, Language language) {
+void ShaderDecompiler::compileInstruction(u32& pc, bool& finished) {
+	const u32 instruction = shader.loadedShader[pc];
+	const u32 opcode = instruction >> 26;
+
+	switch (opcode) {
+		case ShaderOpcodes::DP4: decompiledShader += "dp4\n"; break;
+		case ShaderOpcodes::MOV: decompiledShader += "mov\n"; break;
+		case ShaderOpcodes::END: finished = true; return;
+		default: Helpers::warn("GLSL recompiler: Unknown opcode: %X", opcode); break;
+	}
+
+	pc++;
+}
+
+void ShaderDecompiler::callFunction(const Function& function) { decompiledShader += function.getCallStatement() + ";\n"; }
+
+std::string ShaderGen::decompileShader(PICAShader& shader, EmulatorConfig& config, u32 entrypoint, API api, Language language) {
 	ShaderDecompiler decompiler(shader, config, entrypoint, api, language);
 
 	return decompiler.decompile();