From 353bb96bd371bb600b8885e3a75c4c237778cb50 Mon Sep 17 00:00:00 2001
From: XAMPPRocky <4464295+XAMPPRocky@users.noreply.github.com>
Date: Mon, 25 Jan 2021 10:39:47 +0000
Subject: [PATCH] Add gpu_only proc macro (#392)

---
 Cargo.lock                         |  4 ++++
 crates/spirv-std-macros/Cargo.toml |  4 ++++
 crates/spirv-std-macros/src/lib.rs | 25 +++++++++++++++++++++++++
 crates/spirv-std/src/derivative.rs |  4 +---
 crates/spirv-std/src/lib.rs        |  4 ++--
 crates/spirv-std/src/textures.rs   | 23 +++--------------------
 6 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5db134b668..f44ff49179 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2081,6 +2081,10 @@ dependencies = [
 [[package]]
 name = "spirv-std-macros"
 version = "0.1.0"
+dependencies = [
+ "quote",
+ "syn",
+]
 
 [[package]]
 name = "spirv-tools"
diff --git a/crates/spirv-std-macros/Cargo.toml b/crates/spirv-std-macros/Cargo.toml
index 61a0dd3b39..249d46926c 100644
--- a/crates/spirv-std-macros/Cargo.toml
+++ b/crates/spirv-std-macros/Cargo.toml
@@ -8,3 +8,7 @@ repository = "https://github.com/EmbarkStudios/rust-gpu"
 
 [lib]
 proc-macro = true
+
+[dependencies]
+quote = "1.0.8"
+syn = { version = "1.0.58", features=["full"] }
diff --git a/crates/spirv-std-macros/src/lib.rs b/crates/spirv-std-macros/src/lib.rs
index c10135faab..ed9228eb09 100644
--- a/crates/spirv-std-macros/src/lib.rs
+++ b/crates/spirv-std-macros/src/lib.rs
@@ -29,3 +29,28 @@ pub fn spirv(_attr: TokenStream, item: TokenStream) -> TokenStream {
     }
     tokens.into_iter().collect()
 }
+
+#[proc_macro_attribute]
+pub fn gpu_only(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let syn::ItemFn {
+        attrs,
+        vis,
+        sig,
+        block,
+    } = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let fn_name = sig.ident.clone();
+
+    let output = quote::quote! {
+        // Don't warn on unused arguments on the CPU side.
+        #[cfg_attr(not(target_arch = "spirv"), allow(unused_variables))]
+        #(#attrs)* #vis #sig {
+            #[cfg(target_arch="spirv")] { #block }
+            #[cfg(not(target_arch="spirv"))] {
+                unimplemented!(concat!("`", stringify!(#fn_name), "` is only available on SPIR-V platforms."))
+            }
+        }
+    };
+
+    output.into()
+}
diff --git a/crates/spirv-std/src/derivative.rs b/crates/spirv-std/src/derivative.rs
index 6b9a6e7d58..ac56cdd758 100644
--- a/crates/spirv-std/src/derivative.rs
+++ b/crates/spirv-std/src/derivative.rs
@@ -56,10 +56,8 @@ macro_rules! deriv_caps {
 
 macro_rules! deriv_fn {
     ($name:ident, $inst:ident, $needs_caps:tt) => {
+        #[spirv_std_macros::gpu_only]
         fn $name(self) -> Self {
-            #[cfg(not(target_arch = "spirv"))]
-            panic!(concat!(stringify!($name), " is not supported on the CPU"));
-            #[cfg(target_arch = "spirv")]
             unsafe {
                 let mut o = Default::default();
                 deriv_caps!($needs_caps);
diff --git a/crates/spirv-std/src/lib.rs b/crates/spirv-std/src/lib.rs
index 98f4339b9e..6bba2ba861 100644
--- a/crates/spirv-std/src/lib.rs
+++ b/crates/spirv-std/src/lib.rs
@@ -52,8 +52,8 @@ pub use num_traits;
 pub use textures::*;
 
 /// Calls the `OpDemoteToHelperInvocationEXT` instruction, which corresponds to discard() in HLSL
+#[spirv_std_macros::gpu_only]
 pub fn demote_to_helper_invocation() {
-    #[cfg(target_arch = "spirv")]
     unsafe {
         asm!(
             "OpExtension \"SPV_EXT_demote_to_helper_invocation\"",
@@ -64,8 +64,8 @@ pub fn demote_to_helper_invocation() {
 }
 
 /// Calls the `OpKill` instruction, which corresponds to discard() in GLSL
+#[spirv_std_macros::gpu_only]
 pub fn discard() {
-    #[cfg(target_arch = "spirv")]
     unsafe {
         asm!("OpKill", "%unused = OpLabel");
     }
diff --git a/crates/spirv-std/src/textures.rs b/crates/spirv-std/src/textures.rs
index 6198538e09..aa80dce2e9 100644
--- a/crates/spirv-std/src/textures.rs
+++ b/crates/spirv-std/src/textures.rs
@@ -23,14 +23,8 @@ pub struct Image2d {
 }
 
 impl Image2d {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, sampler: Sampler, coord: Vec2) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = sampler;
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(
@@ -66,14 +60,8 @@ pub struct Image2dArray {
 }
 
 impl Image2dArray {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, sampler: Sampler, coord: Vec3A) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = sampler;
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(
@@ -101,13 +89,8 @@ pub struct SampledImage<I> {
 }
 
 impl SampledImage<Image2d> {
+    #[spirv_std_macros::gpu_only]
     pub fn sample(&self, coord: Vec2) -> Vec4 {
-        #[cfg(not(target_arch = "spirv"))]
-        {
-            let _ = coord;
-            panic!("Image sampling not supported on CPU");
-        }
-        #[cfg(target_arch = "spirv")]
         unsafe {
             let mut result = Default::default();
             asm!(