From c250aed31012a594d84d2ac0cf6cb9e2f7ee0a4e Mon Sep 17 00:00:00 2001
From: Zzzen <callmetsing@gmail.com>
Date: Wed, 29 Nov 2023 05:34:33 +0800
Subject: [PATCH] fix inlay hints for template literal type (#56309)

---
 src/compiler/utilities.ts                     |   3 +-
 src/services/inlayHints.ts                    |  52 +++++++-
 ...tsInteractiveTemplateLiteralTypes.baseline | 116 ++++++++++++++++++
 ...layHintsInteractiveTemplateLiteralTypes.ts |  16 +++
 4 files changed, 181 insertions(+), 6 deletions(-)
 create mode 100644 tests/baselines/reference/inlayHintsInteractiveTemplateLiteralTypes.baseline
 create mode 100644 tests/cases/fourslash/inlayHintsInteractiveTemplateLiteralTypes.ts

diff --git a/src/compiler/utilities.ts b/src/compiler/utilities.ts
index 4a1fe284640..64c9b4ee3d0 100644
--- a/src/compiler/utilities.ts
+++ b/src/compiler/utilities.ts
@@ -5823,7 +5823,8 @@ export function createDiagnosticCollection(): DiagnosticCollection {
 }
 
 const templateSubstitutionRegExp = /\$\{/g;
-function escapeTemplateSubstitution(str: string): string {
+/** @internal */
+export function escapeTemplateSubstitution(str: string): string {
     return str.replace(templateSubstitutionRegExp, "\\${");
 }
 
diff --git a/src/services/inlayHints.ts b/src/services/inlayHints.ts
index c12b76381f9..76d9e6552b9 100644
--- a/src/services/inlayHints.ts
+++ b/src/services/inlayHints.ts
@@ -11,6 +11,7 @@ import {
     EnumMember,
     equateStringsCaseInsensitive,
     escapeString,
+    escapeTemplateSubstitution,
     Expression,
     findChildOfKind,
     findIndex,
@@ -76,7 +77,11 @@ import {
     isQualifiedName,
     isRestTypeNode,
     isSpreadElement,
-    isStringLiteral,
+    isTemplateHead,
+    isTemplateLiteralTypeNode,
+    isTemplateLiteralTypeSpan,
+    isTemplateMiddle,
+    isTemplateTail,
     isThisTypeNode,
     isTupleTypeNode,
     isTypeLiteralNode,
@@ -89,7 +94,7 @@ import {
     isUnionTypeNode,
     isVarConst,
     isVariableDeclaration,
-    LiteralExpression,
+    LiteralLikeNode,
     MethodDeclaration,
     NewExpression,
     Node,
@@ -106,6 +111,7 @@ import {
     Symbol,
     SymbolFlags,
     SyntaxKind,
+    TemplateLiteralLikeNode,
     textSpanIntersectsWith,
     tokenToString,
     TupleTypeReference,
@@ -793,6 +799,28 @@ export function provideInlayHints(context: InlayHintsContext): InlayHint[] {
                     parts.push({ text: tokenToString(node.operator) });
                     visitForDisplayParts(node.operand);
                     break;
+                case SyntaxKind.TemplateLiteralType:
+                    Debug.assertNode(node, isTemplateLiteralTypeNode);
+                    visitForDisplayParts(node.head);
+                    node.templateSpans.forEach(visitForDisplayParts);
+                    break;
+                case SyntaxKind.TemplateHead:
+                    Debug.assertNode(node, isTemplateHead);
+                    parts.push({ text: getLiteralText(node) });
+                    break;
+                case SyntaxKind.TemplateLiteralTypeSpan:
+                    Debug.assertNode(node, isTemplateLiteralTypeSpan);
+                    visitForDisplayParts(node.type);
+                    visitForDisplayParts(node.literal);
+                    break;
+                case SyntaxKind.TemplateMiddle:
+                    Debug.assertNode(node, isTemplateMiddle);
+                    parts.push({ text: getLiteralText(node) });
+                    break;
+                case SyntaxKind.TemplateTail:
+                    Debug.assertNode(node, isTemplateTail);
+                    parts.push({ text: getLiteralText(node) });
+                    break;
                 case SyntaxKind.ThisType:
                     Debug.assertNode(node, isThisTypeNode);
                     parts.push({ text: "this" });
@@ -828,9 +856,23 @@ export function provideInlayHints(context: InlayHintsContext): InlayHint[] {
             });
         }
 
-        function getLiteralText(node: LiteralExpression) {
-            if (isStringLiteral(node)) {
-                return quotePreference === QuotePreference.Single ? `'${escapeString(node.text, CharacterCodes.singleQuote)}'` : `"${escapeString(node.text, CharacterCodes.doubleQuote)}"`;
+        function getLiteralText(node: LiteralLikeNode) {
+            switch (node.kind) {
+                case SyntaxKind.StringLiteral:
+                    return quotePreference === QuotePreference.Single ? `'${escapeString(node.text, CharacterCodes.singleQuote)}'` : `"${escapeString(node.text, CharacterCodes.doubleQuote)}"`;
+                case SyntaxKind.TemplateHead:
+                case SyntaxKind.TemplateMiddle:
+                case SyntaxKind.TemplateTail: {
+                    const rawText = (node as TemplateLiteralLikeNode).rawText ?? escapeTemplateSubstitution(escapeString(node.text, CharacterCodes.backtick));
+                    switch (node.kind) {
+                        case SyntaxKind.TemplateHead:
+                            return "`" + rawText + "${";
+                        case SyntaxKind.TemplateMiddle:
+                            return "}" + rawText + "${";
+                        case SyntaxKind.TemplateTail:
+                            return "}" + rawText + "`";
+                    }
+                }
             }
             return node.text;
         }
diff --git a/tests/baselines/reference/inlayHintsInteractiveTemplateLiteralTypes.baseline b/tests/baselines/reference/inlayHintsInteractiveTemplateLiteralTypes.baseline
new file mode 100644
index 00000000000..51e7aa59b93
--- /dev/null
+++ b/tests/baselines/reference/inlayHintsInteractiveTemplateLiteralTypes.baseline
@@ -0,0 +1,116 @@
+// === Inlay Hints ===
+const lit1 = getTemplateLiteral1();
+          ^
+{
+  "text": "",
+  "displayParts": [
+    {
+      "text": ": "
+    },
+    {
+      "text": "`${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "},${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}`"
+    }
+  ],
+  "position": 73,
+  "kind": "Type",
+  "whitespaceBefore": true
+}
+
+const lit2 = getTemplateLiteral2();
+          ^
+{
+  "text": "",
+  "displayParts": [
+    {
+      "text": ": "
+    },
+    {
+      "text": "`\\${${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "},${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}`"
+    }
+  ],
+  "position": 175,
+  "kind": "Type",
+  "whitespaceBefore": true
+}
+
+const lit3 = getTemplateLiteral3();
+          ^
+{
+  "text": "",
+  "displayParts": [
+    {
+      "text": ": "
+    },
+    {
+      "text": "`start${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}\\${,$${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}end`"
+    }
+  ],
+  "position": 286,
+  "kind": "Type",
+  "whitespaceBefore": true
+}
+
+const lit4 = getTemplateLiteral4();
+          ^
+{
+  "text": "",
+  "displayParts": [
+    {
+      "text": ": "
+    },
+    {
+      "text": "`${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}\\`,${"
+    },
+    {
+      "text": "string"
+    },
+    {
+      "text": "}`"
+    }
+  ],
+  "position": 387,
+  "kind": "Type",
+  "whitespaceBefore": true
+}
\ No newline at end of file
diff --git a/tests/cases/fourslash/inlayHintsInteractiveTemplateLiteralTypes.ts b/tests/cases/fourslash/inlayHintsInteractiveTemplateLiteralTypes.ts
new file mode 100644
index 00000000000..008099f618b
--- /dev/null
+++ b/tests/cases/fourslash/inlayHintsInteractiveTemplateLiteralTypes.ts
@@ -0,0 +1,16 @@
+/// <reference path="fourslash.ts" />
+
+//// declare function getTemplateLiteral1(): `${string},${string}`;
+//// const lit1 = getTemplateLiteral1();
+//// declare function getTemplateLiteral2(): `\${${string},${string}`;
+//// const lit2 = getTemplateLiteral2();
+//// declare function getTemplateLiteral3(): `start${string}\${,$${string}end`;
+//// const lit3 = getTemplateLiteral3();
+//// declare function getTemplateLiteral4(): `${string}\`,${string}`;
+//// const lit4 = getTemplateLiteral4();
+
+
+verify.baselineInlayHints(undefined, {
+    includeInlayVariableTypeHints: true,
+    interactiveInlayHints: true
+});