From 14764896c34b230a5563f48d8b8e565de2f3aa10 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 2 Feb 2024 22:28:02 -0800
Subject: Capability type checking. (#3530)

* Capability type checking.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 .../slang-language-server-semantic-tokens.cpp      | 28 +++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

(limited to 'source/slang/slang-language-server-semantic-tokens.cpp')

diff --git a/source/slang/slang-language-server-semantic-tokens.cpp b/source/slang/slang-language-server-semantic-tokens.cpp
index ae10d62e8..3a40c8e92 100644
--- a/source/slang/slang-language-server-semantic-tokens.cpp
+++ b/source/slang/slang-language-server-semantic-tokens.cpp
@@ -126,7 +126,7 @@ List<SemanticToken> getSemanticTokens(Linkage* linkage, Module* module, UnownedS
         }
         maybeInsertToken(token);
     };
-    iterateAST(
+    iterateASTWithLanguageServerFilter(
         fileName,
         manager,
         module->getModuleDecl(),
@@ -240,8 +240,34 @@ List<SemanticToken> getSemanticTokens(Linkage* linkage, Module* module, UnownedS
                     token.length = (int)attr->originalIdentifierToken.getContentLength();
                     token.type = SemanticTokenType::Type;
                     maybeInsertToken(token);
+
+                    // Insert capability names as enum cases.
+                    if (as<RequireCapabilityAttribute>(attr))
+                    {
+                        for (auto arg : attr->args)
+                        {
+                            if (auto varExpr = as<VarExpr>(arg))
+                            {
+                                if (varExpr->name)
+                                {
+                                    SemanticToken capToken = _createSemanticToken(
+                                        manager, varExpr->loc, nullptr);
+                                    capToken.length = (int)varExpr->name->text.getLength();
+                                    capToken.type = SemanticTokenType::EnumMember;
+                                    maybeInsertToken(capToken);
+                                }
+                            }
+                        }
+                    }
                 }
             }
+            else if (auto targetCase = as<TargetCaseStmt>(node))
+            {
+                SemanticToken token = _createSemanticToken(
+                    manager, targetCase->capabilityToken.loc, targetCase->capabilityToken.getName());
+                token.type = SemanticTokenType::EnumMember;
+                maybeInsertToken(token);
+            }
             else if (auto spirvAsmExpr = as<SPIRVAsmExpr>(node))
             {
                 // Highlight opcodes and enums.
-- 
cgit v1.2.3