From 112a618581bf69901c660a952ca0312424485f6e Mon Sep 17 00:00:00 2001
From: Paris DOUADY <paris.douady@hotmail.fr>
Date: Sun, 26 Feb 2023 23:15:26 +0100
Subject: [PATCH] enforce discard only in fragment

---
 src/valid/function.rs |  1 +
 tests/wgsl-errors.rs  | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/valid/function.rs b/src/valid/function.rs
index b615895ae3..386937d705 100644
--- a/src/valid/function.rs
+++ b/src/valid/function.rs
@@ -595,6 +595,7 @@ impl super::Validator {
                     finished = true;
                 }
                 S::Kill => {
+                    stages &= super::ShaderStages::FRAGMENT;
                     finished = true;
                 }
                 S::Barrier(_) => {
diff --git a/tests/wgsl-errors.rs b/tests/wgsl-errors.rs
index d76252ced5..d44dd97f37 100644
--- a/tests/wgsl-errors.rs
+++ b/tests/wgsl-errors.rs
@@ -943,6 +943,35 @@ fn invalid_arrays() {
     }
 }
 
+#[test]
+fn discard_in_wrong_stage() {
+    check_validation! {
+        "@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    if global_id.x == 3u {
+        discard;
+    }
+}":
+        Err(naga::valid::ValidationError::EntryPoint {
+            stage: naga::ShaderStage::Compute,
+            source: naga::valid::EntryPointError::ForbiddenStageOperations,
+            ..
+        })
+    }
+
+    check_validation! {
+        "@vertex
+fn main() {
+    discard;
+}":
+        Err(naga::valid::ValidationError::EntryPoint {
+            stage: naga::ShaderStage::Vertex,
+            source: naga::valid::EntryPointError::ForbiddenStageOperations,
+            ..
+        })
+    }
+}
+
 #[test]
 fn invalid_structs() {
     check_validation! {