From 06ca9accd701bf35428202c076c1af46facaf9e8 Mon Sep 17 00:00:00 2001
From: Hans-Kristian Arntzen <post@arntzen-software.no>
Date: Fri, 22 Jul 2022 12:04:33 +0200
Subject: [PATCH] HLSL: Add option to emit entry point name 1:1 instead of
 main().

MSL backend supports emitting custom name, and there's no reason for
HLSL to not support that as well, but we have to make it an option to
not break existing users.
---
 main.cpp       |  4 ++++
 spirv_hlsl.cpp | 51 ++++++++++++++++++++++++++++++++++----------------
 spirv_hlsl.hpp |  5 +++++
 3 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/main.cpp b/main.cpp
index b430ec81..81db89ce 100644
--- a/main.cpp
+++ b/main.cpp
@@ -1348,6 +1348,10 @@ static string compile_iteration(const CLIArguments &args, std::vector<uint32_t>
 			build_dummy_sampler = true;
 		}
 
+		// If we're explicitly renaming, we probably want that name to be output.
+		if (!args.entry_point_rename.empty())
+			hlsl_opts.use_entry_point_name = true;
+
 		hlsl_opts.support_nonzero_base_vertex_base_instance = args.hlsl_support_nonzero_base;
 		hlsl_opts.force_storage_buffer_as_uav = args.hlsl_force_storage_buffer_as_uav;
 		hlsl_opts.nonwritable_uav_texture_as_srv = args.hlsl_nonwritable_uav_texture_as_srv;
diff --git a/spirv_hlsl.cpp b/spirv_hlsl.cpp
index c5a37d2c..8c15416c 100644
--- a/spirv_hlsl.cpp
+++ b/spirv_hlsl.cpp
@@ -2404,12 +2404,32 @@ string CompilerHLSL::to_func_call_arg(const SPIRFunction::Parameter &arg, uint32
 	return arg_str;
 }
 
+string CompilerHLSL::get_inner_entry_point_name() const
+{
+	auto &execution = get_entry_point();
+
+	if (hlsl_options.use_entry_point_name)
+	{
+		auto name = join(execution.name, "_inner");
+		ParsedIR::sanitize_underscores(name);
+		return name;
+	}
+
+	if (execution.model == ExecutionModelVertex)
+		return "vert_main";
+	else if (execution.model == ExecutionModelFragment)
+		return "frag_main";
+	else if (execution.model == ExecutionModelGLCompute)
+		return "comp_main";
+	else
+		SPIRV_CROSS_THROW("Unsupported execution model.");
+}
+
 void CompilerHLSL::emit_function_prototype(SPIRFunction &func, const Bitset &return_flags)
 {
 	if (func.self != ir.default_entry_point)
 		add_function_overload(func);
 
-	auto &execution = get_entry_point();
 	// Avoid shadow declarations.
 	local_variable_names = resource_names;
 
@@ -2430,14 +2450,7 @@ void CompilerHLSL::emit_function_prototype(SPIRFunction &func, const Bitset &ret
 
 	if (func.self == ir.default_entry_point)
 	{
-		if (execution.model == ExecutionModelVertex)
-			decl += "vert_main";
-		else if (execution.model == ExecutionModelFragment)
-			decl += "frag_main";
-		else if (execution.model == ExecutionModelGLCompute)
-			decl += "comp_main";
-		else
-			SPIRV_CROSS_THROW("Unsupported execution model.");
+		decl += get_inner_entry_point_name();
 		processing_entry_point = true;
 	}
 	else
@@ -2555,7 +2568,13 @@ void CompilerHLSL::emit_hlsl_entry_point()
 		break;
 	}
 
-	statement(require_output ? "SPIRV_Cross_Output " : "void ", "main(", merge(arguments), ")");
+	const char *entry_point_name;
+	if (hlsl_options.use_entry_point_name)
+		entry_point_name = get_entry_point().name.c_str();
+	else
+		entry_point_name = "main";
+
+	statement(require_output ? "SPIRV_Cross_Output " : "void ", entry_point_name, "(", merge(arguments), ")");
 	begin_scope();
 	bool legacy = hlsl_options.shader_model <= 30;
 
@@ -2728,12 +2747,12 @@ void CompilerHLSL::emit_hlsl_entry_point()
 	});
 
 	// Run the shader.
-	if (execution.model == ExecutionModelVertex)
-		statement("vert_main();");
-	else if (execution.model == ExecutionModelFragment)
-		statement("frag_main();");
-	else if (execution.model == ExecutionModelGLCompute)
-		statement("comp_main();");
+	if (execution.model == ExecutionModelVertex ||
+	    execution.model == ExecutionModelFragment ||
+	    execution.model == ExecutionModelGLCompute)
+	{
+		statement(get_inner_entry_point_name(), "();");
+	}
 	else
 		SPIRV_CROSS_THROW("Unsupported shader stage.");
 
diff --git a/spirv_hlsl.hpp b/spirv_hlsl.hpp
index 23b97f11..f01bcf96 100644
--- a/spirv_hlsl.hpp
+++ b/spirv_hlsl.hpp
@@ -137,6 +137,9 @@ public:
 		// If add_vertex_attribute_remap is used and this feature is used,
 		// the semantic name will be queried once per active location.
 		bool flatten_matrix_vertex_input_semantics = false;
+
+		// Rather than emitting main() for the entry point, use the name in SPIR-V.
+		bool use_entry_point_name = false;
 	};
 
 	explicit CompilerHLSL(std::vector<uint32_t> spirv_)
@@ -374,6 +377,8 @@ private:
 	bool builtin_translates_to_nonarray(spv::BuiltIn builtin) const override;
 
 	std::vector<TypeID> composite_selection_workaround_types;
+
+	std::string get_inner_entry_point_name() const;
 };
 } // namespace SPIRV_CROSS_NAMESPACE