Fix missing semantic highlighting in attributes and ExtractExitentialValueExpr. (#2541)

* Fix missing semantic highlighting in attributes and ExtractExitentialValueExpr.

* Fix regression on partially specialized generic expr highlighting.

* Add regression test.

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

10 files changed, 55 insertions, 7 deletions

diff --git a/source/slang/slang-ast-expr.h b/source/slang/slang-ast-expr.h
index a0268cce9..9fdc10807 100644
--- a/source/slang/slang-ast-expr.h
+++ b/source/slang/slang-ast-expr.h
@@ -425,6 +425,7 @@ class ExtractExistentialValueExpr: public Expr
 {
     SLANG_AST_CLASS(ExtractExistentialValueExpr)
     DeclRef<VarDeclBase> declRef;
+    Expr* originalExpr;
 };
 
 class OpenRefExpr : public Expr
diff --git a/source/slang/slang-ast-iterator.h b/source/slang/slang-ast-iterator.h
index 79aade1ee..b23f7c6ca 100644
--- a/source/slang/slang-ast-iterator.h
+++ b/source/slang/slang-ast-iterator.h
@@ -193,6 +193,7 @@ struct ASTIterator
         void visitExtractExistentialValueExpr(ExtractExistentialValueExpr* expr)
         {
             iterator->maybeDispatchCallback(expr);
+            dispatchIfNotNull(expr->originalExpr);
         }
 
         void visitDeclRefExpr(DeclRefExpr* expr)
@@ -438,6 +439,15 @@ void ASTIterator<CallbackFunc>::visitDecl(DeclBase* decl)
             visitDecl(member);
         }
     }
+    for (auto modifier : decl->modifiers)
+    {
+        if (auto attr = as<Attribute>(modifier))
+        {
+            maybeDispatchCallback(attr);
+            for (auto arg : attr->args)
+                visitExpr(arg);
+        }
+    }
 }
 template <typename CallbackFunc>
 void ASTIterator<CallbackFunc>::visitExpr(Expr* expr)
diff --git a/source/slang/slang-ast-support-types.cpp b/source/slang/slang-ast-support-types.cpp
index b550eaa68..aa0513569 100644
--- a/source/slang/slang-ast-support-types.cpp
+++ b/source/slang/slang-ast-support-types.cpp
@@ -57,9 +57,9 @@ Expr* getInnerMostExprFromHigherOrderExpr(Expr* expr)
 UnownedStringSlice getHigherOrderOperatorName(HigherOrderInvokeExpr* expr)
 {
     if (as<ForwardDifferentiateExpr>(expr))
-        return UnownedStringSlice("fwd_diff");
+        return UnownedStringSlice("__fwd_diff");
     else if (as<BackwardDifferentiateExpr>(expr))
-        return UnownedStringSlice("bwd_diff");
+        return UnownedStringSlice("__bwd_diff");
     return UnownedStringSlice();
 }
 }
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index 4b2d490b7..7297ca282 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -160,6 +160,7 @@ namespace Slang
             ExtractExistentialValueExpr* openedValue = m_astBuilder->create<ExtractExistentialValueExpr>();
             openedValue->declRef = varDeclRef;
             openedValue->type = QualType(openedType);
+            openedValue->originalExpr = expr;
 
             // The result of opening an existential is an l-value
             // if the original existential is an l-value.
diff --git a/source/slang/slang-check-overload.cpp b/source/slang/slang-check-overload.cpp
index 911587fd5..09ddde2de 100644
--- a/source/slang/slang-check-overload.cpp
+++ b/source/slang/slang-check-overload.cpp
@@ -647,9 +647,6 @@ namespace Slang
             goto error;
 
         {
-            auto originalAppExpr = as<AppExprBase>(context.originalExpr);
-
-
             Expr* baseExpr;
             switch(candidate.flavor)
             {
@@ -659,7 +656,7 @@ namespace Slang
                     candidate.item,
                     context.baseExpr,
                     context.funcLoc,
-                    originalAppExpr ? originalAppExpr->functionExpr : nullptr);
+                    context.originalExpr);
                 break;
             case OverloadCandidate::Flavor::Expr:
             default:
diff --git a/source/slang/slang-language-server-auto-format.cpp b/source/slang/slang-language-server-auto-format.cpp
index 5bc9031ee..9e41401b2 100644
--- a/source/slang/slang-language-server-auto-format.cpp
+++ b/source/slang/slang-language-server-auto-format.cpp
@@ -263,6 +263,8 @@ List<Edit> formatSource(UnownedStringSlice text, Index lineStart, Index lineEnd,
             }
             else if (!originalHasLineBreak && newHasLineBreak)
             {
+                if (edt.offset < text.getLength() && edt.offset >= 0 && text[edt.offset] == '}')
+                    continue;
                 edt.text = " ";
             }
             else
diff --git a/source/slang/slang-language-server-semantic-tokens.cpp b/source/slang/slang-language-server-semantic-tokens.cpp
index 485dd7a44..d52f631bd 100644
--- a/source/slang/slang-language-server-semantic-tokens.cpp
+++ b/source/slang/slang-language-server-semantic-tokens.cpp
@@ -194,6 +194,16 @@ List<SemanticToken> getSemanticTokens(Linkage* linkage, Module* module, UnownedS
                     maybeInsertToken(token);
                 }
             }
+            else if (auto attr = as<Attribute>(node))
+            {
+                if (attr->getKeywordName())
+                {
+                    SemanticToken token = _createSemanticToken(
+                        manager, attr->getKeywordNameAndLoc().loc, attr->getKeywordName());
+                    token.type = SemanticTokenType::Type;
+                    maybeInsertToken(token);
+                }
+            }
         });
     // Insert macro tokens.
     auto& preprocessorInfo = linkage->contentAssistInfo.preprocessorInfo;
diff --git a/source/slang/slang-workspace-version.cpp b/source/slang/slang-workspace-version.cpp
index d5fc62e79..309d4e51d 100644
--- a/source/slang/slang-workspace-version.cpp
+++ b/source/slang/slang-workspace-version.cpp
@@ -51,7 +51,11 @@ void Workspace::changeDoc(const String& path, LanguageServerProtocol::Range rang
         auto endOffset = doc->getOffset(line, col);
         auto originalText = doc->getText().getUnownedSlice();
         StringBuilder newText;
-        newText << originalText.head(startOffset) << text << originalText.tail(endOffset);
+        if (startOffset != -1)
+            newText << originalText.head(startOffset);
+        newText << text;
+        if (endOffset != -1)
+            newText << originalText.tail(endOffset);
         changeDoc(doc.Ptr(), newText.ProduceString());
     }
 }
diff --git a/tests/language-server/partially-specialized-generic.slang b/tests/language-server/partially-specialized-generic.slang
new file mode 100644
index 000000000..120a33f9e
--- /dev/null
+++ b/tests/language-server/partially-specialized-generic.slang
@@ -0,0 +1,10 @@
+//TEST:LANG_SERVER:
+//HOVER:9,29
+
+struct MyType {}
+
+void main()
+{
+    MyType a;
+    let b = reinterpret<MyType>(a);
+}
diff --git a/tests/language-server/partially-specialized-generic.slang.expected.txt b/tests/language-server/partially-specialized-generic.slang.expected.txt
new file mode 100644
index 000000000..d1a87d8aa
--- /dev/null
+++ b/tests/language-server/partially-specialized-generic.slang.expected.txt
@@ -0,0 +1,13 @@
+--------
+range: 8,24 - 8,30
+content:
+```
+struct MyType
+```
+
+TEST:LANG_SERVER:  
+HOVER:9,29  
+
+{REDACTED}.slang(4)
+
+