From bae87afb20f95f9f27c64c4955bbc4464c576509 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 6 Feb 2025 22:02:43 -0800
Subject: Support stage_switch. (#6311)

* Support stage_switch.

* Update proposal status.

* Fix gl_InstanceID.

* Fix.
---
 source/slang/slang-check-stmt.cpp | 44 ++++++++++++++++++++++++++++++---------
 1 file changed, 34 insertions(+), 10 deletions(-)

(limited to 'source/slang/slang-check-stmt.cpp')

diff --git a/source/slang/slang-check-stmt.cpp b/source/slang/slang-check-stmt.cpp
index 151c9324a..db6f00d23 100644
--- a/source/slang/slang-check-stmt.cpp
+++ b/source/slang/slang-check-stmt.cpp
@@ -367,6 +367,40 @@ void SemanticsStmtVisitor::visitTargetSwitchStmt(TargetSwitchStmt* stmt)
     HashSet<Stmt*> checkedStmt;
     for (auto caseStmt : stmt->targetCases)
     {
+        CapabilitySet set((CapabilityName)caseStmt->capability);
+
+        CapabilityName canonicalStage = CapabilityName::Invalid;
+        bool isStage = isStageAtom((CapabilityName)caseStmt->capability, canonicalStage);
+        if (as<StageSwitchStmt>(stmt))
+        {
+            if (!isStage && caseStmt->capability != 0)
+            {
+                getSink()->diagnose(
+                    caseStmt->capabilityToken.loc,
+                    Diagnostics::unknownStageName,
+                    caseStmt->capabilityToken);
+            }
+            caseStmt->capability = (int)canonicalStage;
+        }
+        else
+        {
+            if (isStage)
+            {
+                getSink()->diagnose(
+                    caseStmt->capabilityToken.loc,
+                    Diagnostics::targetSwitchCaseCannotBeAStage);
+            }
+            else if (
+                caseStmt->capabilityToken.getContentLength() != 0 &&
+                (set.getCapabilityTargetSets().getCount() != 1 || set.isInvalid() || set.isEmpty()))
+            {
+                getSink()->diagnose(
+                    caseStmt->capabilityToken.loc,
+                    Diagnostics::invalidTargetSwitchCase,
+                    capabilityNameToString((CapabilityName)caseStmt->capability));
+            }
+        }
+
         if (checkedStmt.contains(caseStmt->body))
             continue;
         subContext.checkStmt(caseStmt);
@@ -377,7 +411,6 @@ void SemanticsStmtVisitor::visitTargetSwitchStmt(TargetSwitchStmt* stmt)
 void SemanticsStmtVisitor::visitTargetCaseStmt(TargetCaseStmt* stmt)
 {
     auto switchStmt = FindOuterStmt<TargetSwitchStmt>();
-    CapabilitySet set((CapabilityName)stmt->capability);
     if (getShared()->isInLanguageServer() &&
         getShared()->getSession()->getCompletionRequestTokenName() ==
             stmt->capabilityToken.getName())
@@ -385,15 +418,6 @@ void SemanticsStmtVisitor::visitTargetCaseStmt(TargetCaseStmt* stmt)
         getShared()->getLinkage()->contentAssistInfo.completionSuggestions.scopeKind =
             CompletionSuggestions::ScopeKind::Capabilities;
     }
-
-    if (stmt->capabilityToken.getContentLength() != 0 &&
-        (set.getCapabilityTargetSets().getCount() != 1 || set.isInvalid() || set.isEmpty()))
-    {
-        getSink()->diagnose(
-            stmt->capabilityToken.loc,
-            Diagnostics::invalidTargetSwitchCase,
-            capabilityNameToString((CapabilityName)stmt->capability));
-    }
     if (!switchStmt)
     {
         getSink()->diagnose(stmt, Diagnostics::caseOutsideSwitch);
-- 
cgit v1.2.3