From 675e6d480db6366c312e4279da0e3bca5b1be04f Mon Sep 17 00:00:00 2001
From: Timothee Guerin <tiguerin@microsoft.com>
Date: Thu, 8 Aug 2024 14:42:27 -0700
Subject: [PATCH] Fix type argument tmlanguage (#4089)

Fix #3524
---
 ...pe-argument-tmlanguage-2024-7-5-16-16-4.md |  8 +++
 grammars/typespec.json                        |  2 +-
 packages/compiler/src/server/tmlanguage.ts    |  2 +-
 .../compiler/test/server/colorization.test.ts | 72 +++++++++++++------
 4 files changed, 62 insertions(+), 22 deletions(-)
 create mode 100644 .chronus/changes/fix-type-argument-tmlanguage-2024-7-5-16-16-4.md

diff --git a/.chronus/changes/fix-type-argument-tmlanguage-2024-7-5-16-16-4.md b/.chronus/changes/fix-type-argument-tmlanguage-2024-7-5-16-16-4.md
new file mode 100644
index 000000000..e60217f34
--- /dev/null
+++ b/.chronus/changes/fix-type-argument-tmlanguage-2024-7-5-16-16-4.md
@@ -0,0 +1,8 @@
+---
+# Change versionKind to one of: internal, fix, dependencies, feature, deprecation, breaking
+changeKind: fix
+packages:
+  - "@typespec/compiler"
+---
+
+Fix tmlanguage for named type argument in type reference.
diff --git a/grammars/typespec.json b/grammars/typespec.json
index 23e1e6780..5855819f5 100644
--- a/grammars/typespec.json
+++ b/grammars/typespec.json
@@ -1398,7 +1398,7 @@
           "name": "keyword.operator.assignment.tsp"
         }
       },
-      "end": "=",
+      "end": "(?=>)|(?=,|;|@|\\)|\\}|\\b(?:extern)\\b|\\b(?:namespace|model|op|using|import|enum|alias|union|interface|dec|fn)\\b)",
       "endCaptures": {
         "0": {
           "name": "keyword.operator.assignment.tsp"
diff --git a/packages/compiler/src/server/tmlanguage.ts b/packages/compiler/src/server/tmlanguage.ts
index 78911310b..4e25826ea 100644
--- a/packages/compiler/src/server/tmlanguage.ts
+++ b/packages/compiler/src/server/tmlanguage.ts
@@ -343,7 +343,7 @@ const typeArgument: BeginEndRule = {
     "1": { scope: "entity.name.type.tsp" },
     "2": { scope: "keyword.operator.assignment.tsp" },
   },
-  end: `=`,
+  end: `(?=>)|${universalEnd}`,
   endCaptures: {
     "0": { scope: "keyword.operator.assignment.tsp" },
   },
diff --git a/packages/compiler/test/server/colorization.test.ts b/packages/compiler/test/server/colorization.test.ts
index 5f6588c60..f9b34677e 100644
--- a/packages/compiler/test/server/colorization.test.ts
+++ b/packages/compiler/test/server/colorization.test.ts
@@ -827,26 +827,58 @@ function testColorization(description: string, tokenize: Tokenize) {
       });
     });
 
-    it("named template argument list", async () => {
-      const tokens = await tokenize("alias X = Foo<boolean, T = string, U = int32>;");
-      deepStrictEqual(tokens, [
-        Token.keywords.alias,
-        Token.identifiers.type("X"),
-        Token.operators.assignment,
-        Token.identifiers.type("Foo"),
-        Token.punctuation.typeParameters.begin,
-        Token.identifiers.type("boolean"),
-        Token.punctuation.comma,
-        Token.identifiers.type("T"),
-        Token.operators.assignment,
-        Token.identifiers.type("string"),
-        Token.punctuation.comma,
-        Token.identifiers.type("U"),
-        Token.operators.assignment,
-        Token.identifiers.type("int32"),
-        Token.punctuation.typeParameters.end,
-        Token.punctuation.semicolon,
-      ]);
+    describe("template argument", () => {
+      it("multiple named arguments", async () => {
+        const tokens = await tokenize("alias X = Foo<boolean, T = string, U = int32>;");
+        deepStrictEqual(tokens, [
+          Token.keywords.alias,
+          Token.identifiers.type("X"),
+          Token.operators.assignment,
+          Token.identifiers.type("Foo"),
+          Token.punctuation.typeParameters.begin,
+          Token.identifiers.type("boolean"),
+          Token.punctuation.comma,
+          Token.identifiers.type("T"),
+          Token.operators.assignment,
+          Token.identifiers.type("string"),
+          Token.punctuation.comma,
+          Token.identifiers.type("U"),
+          Token.operators.assignment,
+          Token.identifiers.type("int32"),
+          Token.punctuation.typeParameters.end,
+          Token.punctuation.semicolon,
+        ]);
+      });
+
+      it("multiple references", async () => {
+        const tokens = await tokenize(`
+          alias A = Foo<Parameters=string>;
+          alias B = Foo<Parameters=string>;  
+        `);
+        deepStrictEqual(tokens, [
+          Token.keywords.alias,
+          Token.identifiers.type("A"),
+          Token.operators.assignment,
+          Token.identifiers.type("Foo"),
+          Token.punctuation.typeParameters.begin,
+          Token.identifiers.type("Parameters"),
+          Token.operators.assignment,
+          Token.identifiers.type("string"),
+          Token.punctuation.typeParameters.end,
+          Token.punctuation.semicolon,
+          // --
+          Token.keywords.alias,
+          Token.identifiers.type("B"),
+          Token.operators.assignment,
+          Token.identifiers.type("Foo"),
+          Token.punctuation.typeParameters.begin,
+          Token.identifiers.type("Parameters"),
+          Token.operators.assignment,
+          Token.identifiers.type("string"),
+          Token.punctuation.typeParameters.end,
+          Token.punctuation.semicolon,
+        ]);
+      });
     });
 
     describe("enums", () => {