From 6678f9a49d72e23710308aaa8d3f1ea948bc5345 Mon Sep 17 00:00:00 2001
From: XAMPPRocky <4464295+XAMPPRocky@users.noreply.github.com>
Date: Tue, 2 Mar 2021 14:30:08 +0100
Subject: [PATCH] Add entry point attribute configuration. (#458)

* Add entrypoint attribute configuration.

* Update with feedback

* Add docs

* fmt

* rm usused_attributes

* other changes

* fmt
---
 .../src/codegen_cx/declare.rs                 |  8 ++++++--
 crates/rustc_codegen_spirv/src/symbols.rs     | 20 +++++++++++++++++++
 crates/spirv-builder/src/test/basic.rs        | 20 +++++++++++++++++++
 docs/src/attributes.md                        |  4 ++++
 4 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs b/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs
index 2f001e14e7..847c22f28a 100644
--- a/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs
+++ b/crates/rustc_codegen_spirv/src/codegen_cx/declare.rs
@@ -114,8 +114,12 @@ impl<'tcx> CodegenCx<'tcx> {
         for attr in parse_attrs(self, self.tcx.get_attrs(instance.def_id())) {
             match attr {
                 SpirvAttribute::Entry(entry) => {
-                    let crate_relative_name = instance.to_string();
-                    self.entry_stub(&instance, &fn_abi, declared, crate_relative_name, entry)
+                    let entry_name = entry
+                        .name
+                        .as_ref()
+                        .map(ToString::to_string)
+                        .unwrap_or_else(|| instance.to_string());
+                    self.entry_stub(&instance, &fn_abi, declared, entry_name, entry)
                 }
                 SpirvAttribute::UnrollLoops => {
                     self.unroll_loops_decorations
diff --git a/crates/rustc_codegen_spirv/src/symbols.rs b/crates/rustc_codegen_spirv/src/symbols.rs
index f0e12bd95b..f8e7cdcc33 100644
--- a/crates/rustc_codegen_spirv/src/symbols.rs
+++ b/crates/rustc_codegen_spirv/src/symbols.rs
@@ -33,6 +33,7 @@ pub struct Symbols {
     pub spirv13: Symbol,
     pub spirv14: Symbol,
     pub spirv15: Symbol,
+    pub entry_point_name: Symbol,
     descriptor_set: Symbol,
     binding: Symbol,
     image_type: Symbol,
@@ -370,6 +371,7 @@ impl Symbols {
         Self {
             fmt_decimal: Symbol::intern("fmt_decimal"),
 
+            entry_point_name: Symbol::intern("entry_point_name"),
             spirv: Symbol::intern("spirv"),
             spirv_std: Symbol::intern("spirv_std"),
             libm: Symbol::intern("libm"),
@@ -437,6 +439,7 @@ impl AsRef<[u32]> for ExecutionModeExtra {
 pub struct Entry {
     pub execution_model: ExecutionModel,
     pub execution_modes: Vec<(ExecutionMode, ExecutionModeExtra)>,
+    pub name: Option<Symbol>,
 }
 
 impl From<ExecutionModel> for Entry {
@@ -444,6 +447,7 @@ impl From<ExecutionModel> for Entry {
         Self {
             execution_model,
             execution_modes: Vec::new(),
+            name: None,
         }
     }
 }
@@ -791,6 +795,22 @@ fn parse_entry_attrs(
                             }
                         }
                     }
+                } else if attr_name.name == sym.entry_point_name {
+                    match attr.value_str() {
+                        Some(sym) => {
+                            entry.name = Some(sym);
+                        }
+                        None => {
+                            return Err((
+                                attr_name.span,
+                                format!(
+                                    "#[spirv({}(..))] unknown attribute argument {}",
+                                    name.name.to_ident_string(),
+                                    attr_name.name.to_ident_string()
+                                ),
+                            ))
+                        }
+                    }
                 } else {
                     return Err((
                         attr_name.span,
diff --git a/crates/spirv-builder/src/test/basic.rs b/crates/spirv-builder/src/test/basic.rs
index 7ea104302e..92abe5642a 100644
--- a/crates/spirv-builder/src/test/basic.rs
+++ b/crates/spirv-builder/src/test/basic.rs
@@ -28,6 +28,26 @@ pub fn main() {
 "#);
 }
 
+#[test]
+fn custom_entry_point() {
+    dis_globals(
+        r#"
+#[spirv(fragment(entry_point_name="hello_world"))]
+pub fn main() { }
+"#,
+        r#"OpCapability Shader
+OpCapability VulkanMemoryModel
+OpCapability VariablePointers
+OpExtension "SPV_KHR_vulkan_memory_model"
+OpMemoryModel Logical Vulkan
+OpEntryPoint Fragment %1 "hello_world"
+OpExecutionMode %1 OriginUpperLeft
+OpName %2 "test_project::main"
+%3 = OpTypeVoid
+%4 = OpTypeFunction %3"#,
+    );
+}
+
 #[test]
 // blocked on: https://github.com/EmbarkStudios/rust-gpu/issues/69
 #[ignore]
diff --git a/docs/src/attributes.md b/docs/src/attributes.md
index ef83eb5742..d8b0b115e5 100644
--- a/docs/src/attributes.md
+++ b/docs/src/attributes.md
@@ -17,6 +17,10 @@ fn main() { }
 
 Common values are `#[spirv(fragment)]` and `#[spirv(vertex)]`. A list of all supported names can be found in [spirv_headers](https://docs.rs/spirv_headers/1.5.0/spirv_headers/enum.ExecutionModel.html) - convert the enum name to snake_case for the rust-gpu attribute name.
 
+### Override entry point name
+
+You can override the default `OpEntryPoint` name for any entry point with the `entry_point_name` sub-attribute on any of the execution model attributes. (e.g. `#[spirv(vertex(entry_point_name="foo"))]`)
+
 ## Builtins
 
 When declaring inputs and outputs, sometimes you want to declare it as a "builtin". This means many things, but one example is `gl_Position` from glsl - the GPU assigns inherent meaning to the variable and uses it for placing the vertex in clip space. The equivalent in rust-gpu is called `position`.