Add intellisense support for spirv_asm blocks. (#3264)

Co-authored-by: Yong He <yhe@nvidia.com>

6 files changed, 149 insertions, 11 deletions

diff --git a/source/slang/slang-language-server-auto-format.cpp b/source/slang/slang-language-server-auto-format.cpp
index 5b94712a0..154122769 100644
--- a/source/slang/slang-language-server-auto-format.cpp
+++ b/source/slang/slang-language-server-auto-format.cpp
@@ -58,6 +58,80 @@ String findClangFormatTool()
     return String();
 }
 
+List<TextRange> extractFormattingExclusionRanges(UnownedStringSlice text)
+{
+    List<TextRange> resultRanges;
+    // Search the document and find all the spirv-asm blocks to exclude them from formatting.
+    Lexer lexer;
+    SourceManager manager;
+    manager.initialize(nullptr, nullptr);
+    auto sourceFile = manager.createSourceFileWithString(PathInfo(), text);
+    auto sourceView = manager.createSourceView(sourceFile, nullptr, SourceLoc());
+    DiagnosticSink sink;
+    RootNamePool rootPool;
+    NamePool namePool;
+    namePool.setRootNamePool(&rootPool);
+    MemoryArena memory;
+    memory.init(1 << 16);
+    lexer.initialize(sourceView, &sink, &namePool, &memory);
+    for (;;)
+    {
+        auto headerToken = lexer.lexToken();
+        if (headerToken.type == TokenType::EndOfFile)
+            break;
+        if (headerToken.type != TokenType::Identifier || headerToken.getContent() != toSlice("spirv_asm"))
+        {
+            continue;
+        }
+        
+        // We have found a spirv-asm block, now find the end of it.
+        int braceCounter = 0;
+        Token endToken;
+        Token startToken;
+        for (;;)
+        {
+            auto token = lexer.lexToken();
+            if (token.type == TokenType::EndOfFile)
+                break;
+            switch (token.type)
+            {
+            case TokenType::RBrace:
+                braceCounter--;
+                if (braceCounter == 0)
+                {
+                    endToken = token;
+                    goto breakLabel;
+                }
+            case TokenType::LBrace:
+                {
+                    if (braceCounter == 0)
+                        startToken = token;
+                    braceCounter++;
+                }
+                break;
+            case TokenType::WhiteSpace:
+            case TokenType::LineComment:
+            case TokenType::BlockComment:
+            case TokenType::NewLine:
+                break;
+            default:
+                if (braceCounter == 0)
+                    goto breakLabel;
+                break;
+            }
+        }
+        breakLabel:;
+        if (endToken.type == TokenType::RBrace)
+        {
+            TextRange range;
+            range.offsetStart = startToken.getLoc().getRaw();
+            range.offsetEnd = endToken.getLoc().getRaw();
+            resultRanges.add(range);
+        }
+    }
+    return resultRanges;
+}
+
 void translateXmlEscape(StringBuilder& sb, UnownedStringSlice text)
 {
     if (text.getLength() == 0)
@@ -152,7 +226,7 @@ bool shouldUseFallbackStyle(const FormatOptions& options)
     return true;
 }
 
-List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd, Index cursorOffset, const FormatOptions& options)
+List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd, Index cursorOffset, const List<TextRange>& exclusionRanges, const FormatOptions& options)
 {
     List<Edit> edits;
 
@@ -232,6 +306,16 @@ List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd,
         pos = line.indexOf(lengthStr);
         if (pos == -1)
             continue;
+        Index exclusionRangeId = exclusionRanges.binarySearch(edt.offset, [](TextRange range, Index offset)
+            {
+                if (range.offsetEnd <= offset)
+                    return -1;
+                if (range.offsetStart > offset)
+                    return 1;
+                return 0;
+            });
+        if (exclusionRangeId != -1)
+            continue;
         pos += lengthStr.getLength();
         if (pos < line.getLength() && line[pos] == '\'')
             pos++;
diff --git a/source/slang/slang-language-server-auto-format.h b/source/slang/slang-language-server-auto-format.h
index a2668dbe1..e2cf89a25 100644
--- a/source/slang/slang-language-server-auto-format.h
+++ b/source/slang/slang-language-server-auto-format.h
@@ -13,6 +13,12 @@ struct Edit
     String text;
 };
 
+struct TextRange
+{
+    Index offsetStart;
+    Index offsetEnd;
+};
+
 enum class FormatBehavior
 {
     Standard,
@@ -31,6 +37,9 @@ struct FormatOptions
 };
 
 String findClangFormatTool();
-List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd, Index cursorOffset, const FormatOptions& options);
+
+List<TextRange> extractFormattingExclusionRanges(UnownedStringSlice text);
+
+List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd, Index cursorOffset, const List<TextRange>& exclusionRanges, const FormatOptions& options);
 
 } // namespace Slang
diff --git a/source/slang/slang-language-server-semantic-tokens.cpp b/source/slang/slang-language-server-semantic-tokens.cpp
index 159b91d01..707414c5e 100644
--- a/source/slang/slang-language-server-semantic-tokens.cpp
+++ b/source/slang/slang-language-server-semantic-tokens.cpp
@@ -17,7 +17,8 @@ const char* kSemanticTokenTypes[] = {
     "property",
     "namespace",
     "keyword",
-    "macro"
+    "macro",
+    "string"
 };
 
 static_assert(SLANG_COUNT_OF(kSemanticTokenTypes) == (int)SemanticTokenType::NormalText, "kSemanticTokenTypes must match SemanticTokenType");
@@ -213,6 +214,47 @@ List<SemanticToken> getSemanticTokens(Linkage* linkage, Module* module, UnownedS
                     maybeInsertToken(token);
                 }
             }
+            else if (auto spirvAsmExpr = as<SPIRVAsmExpr>(node))
+            {
+                // Highlight opcodes and enums.
+                for (auto inst : spirvAsmExpr->insts)
+                {
+                    // OpCode
+                    {
+                        SemanticToken token = _createSemanticToken(
+                            manager, inst.opcode.token.loc, inst.opcode.token.getName());
+                        token.type = SemanticTokenType::Function;
+                        maybeInsertToken(token);
+                    }
+                    // Other operands
+                    for (auto operand : inst.operands)
+                    {
+                        SemanticTokenType operandTokenType = SemanticTokenType::NormalText;
+                        switch (operand.flavor)
+                        {
+                        case SPIRVAsmOperand::Flavor::NamedValue:
+                        case SPIRVAsmOperand::Flavor::BuiltinVar:
+                            operandTokenType = SemanticTokenType::EnumMember;
+                            break;
+                        case SPIRVAsmOperand::Flavor::ResultMarker:
+                        case SPIRVAsmOperand::Flavor::TruncateMarker:
+                        case SPIRVAsmOperand::Flavor::GLSL450Set:
+                            operandTokenType = SemanticTokenType::Macro;
+                            break;
+                        case SPIRVAsmOperand::Flavor::Id:
+                            operandTokenType = SemanticTokenType::String;
+                            break;
+                        }
+                        if (operandTokenType != SemanticTokenType::NormalText)
+                        {
+                            SemanticToken token = _createSemanticToken(
+                                manager, operand.token.loc, operand.token.getName());
+                            token.type = operandTokenType;
+                            maybeInsertToken(token);
+                        }
+                    }
+                }
+            }
         });
     // Insert macro tokens.
     auto& preprocessorInfo = linkage->contentAssistInfo.preprocessorInfo;
diff --git a/source/slang/slang-language-server-semantic-tokens.h b/source/slang/slang-language-server-semantic-tokens.h
index 4a7368fcc..8a736345f 100644
--- a/source/slang/slang-language-server-semantic-tokens.h
+++ b/source/slang/slang-language-server-semantic-tokens.h
@@ -11,7 +11,7 @@ namespace Slang
 {
 enum class SemanticTokenType
 {
-    Type, EnumMember, Variable, Parameter, Function, Property, Namespace, Keyword, Macro, NormalText
+    Type, EnumMember, Variable, Parameter, Function, Property, Namespace, Keyword, Macro, String, NormalText
 };
 extern const char* kSemanticTokenTypes[(int)SemanticTokenType::NormalText];
 
diff --git a/source/slang/slang-language-server.cpp b/source/slang/slang-language-server.cpp
index ad88ab5ac..3a21fa278 100644
--- a/source/slang/slang-language-server.cpp
+++ b/source/slang/slang-language-server.cpp
@@ -1470,7 +1470,8 @@ SlangResult LanguageServer::formatting(const LanguageServerProtocol::DocumentFor
     if (m_formatOptions.clangFormatLocation.getLength() == 0)
         m_formatOptions.clangFormatLocation = findClangFormatTool();
     m_formatOptions.fileName = canonicalPath;
-    auto edits = formatSource(doc->getText().getUnownedSlice(), -1, -1, -1, m_formatOptions);
+    List<TextRange> exclusionRange = extractFormattingExclusionRanges(doc->getText().getUnownedSlice());
+    auto edits = formatSource(doc->getText().getUnownedSlice(), -1, -1, -1, exclusionRange, m_formatOptions);
     auto textEdits = translateTextEdits(doc, edits);
     m_connection->sendResult(&textEdits, responseId);
     return SLANG_OK;
@@ -1493,7 +1494,8 @@ SlangResult LanguageServer::rangeFormatting(const LanguageServerProtocol::Docume
     auto options = m_formatOptions;
     if (!m_formatOptions.allowLineBreakInRangeFormatting)
         options.behavior = FormatBehavior::PreserveLineBreak;
-    auto edits = formatSource(doc->getText().getUnownedSlice(), args.range.start.line, args.range.end.line, endOffset, options);
+    List<TextRange> exclusionRange = extractFormattingExclusionRanges(doc->getText().getUnownedSlice());
+    auto edits = formatSource(doc->getText().getUnownedSlice(), args.range.start.line, args.range.end.line, endOffset, exclusionRange, options);
     auto textEdits = translateTextEdits(doc, edits);
     m_connection->sendResult(&textEdits, responseId);
     return SLANG_OK;
@@ -1521,7 +1523,8 @@ SlangResult LanguageServer::onTypeFormatting(const LanguageServerProtocol::Docum
     auto options = m_formatOptions;
     if (!m_formatOptions.allowLineBreakInOnTypeFormatting)
         options.behavior = FormatBehavior::PreserveLineBreak;
-    auto edits = formatSource(doc->getText().getUnownedSlice(), args.position.line, args.position.line, cursorOffset, options);
+    List<TextRange> exclusionRange = extractFormattingExclusionRanges(doc->getText().getUnownedSlice());
+    auto edits = formatSource(doc->getText().getUnownedSlice(), args.position.line, args.position.line, cursorOffset, exclusionRange, options);
     auto textEdits = translateTextEdits(doc, edits);
     m_connection->sendResult(&textEdits, responseId);
     return SLANG_OK;
diff --git a/source/slang/slang-parser.cpp b/source/slang/slang-parser.cpp
index 517ab9e10..54b93a1ef 100644
--- a/source/slang/slang-parser.cpp
+++ b/source/slang/slang-parser.cpp
@@ -6468,7 +6468,6 @@ namespace Slang
                 Diagnostics::spirvInstructionWithoutResultTypeId,
                 ret.opcode.token
             );
-            return std::nullopt;
         }
 
         //
@@ -6519,9 +6518,10 @@ namespace Slang
                 ret.operands.add(*operand);
             }
             else
-                return std::nullopt;
+            {
+                break;
+            }
         }
-
         return ret;
     }
 
@@ -6552,7 +6552,7 @@ namespace Slang
         }
         parser->ReadMatchingToken(TokenType::RBrace);
 
-        return failed ? nullptr : asmExpr;
+        return asmExpr;
     }
 
     static Expr* parsePrefixExpr(Parser* parser)