Implement support for Code Lenses (#2145)

Author: DanielRosenwasser

_extension/package.json

@@ -40,17 +40,23 @@
                         ],
                         "default": "verbose",
                         "description": "Trace TypeScript Go server communication.",
-                        "tags": ["experimental"]
+                        "tags": [
+                            "experimental"
+                        ]
                     },
                     "typescript.native-preview.pprofDir": {
                         "type": "string",
                         "description": "Directory to write pprof profiles to.",
-                        "tags": ["experimental"]
+                        "tags": [
+                            "experimental"
+                        ]
                     },
                     "typescript.native-preview.tsdk": {
                         "type": "string",
                         "description": "Path to the @typescript/native-preview package or tsgo binary directory. If not specified, the extension will look for it in the default location.",
-                        "tags": ["experimental"]
+                        "tags": [
+                            "experimental"
+                        ]
                     }
                 }
             }
@@ -91,6 +97,11 @@
                 "title": "Report Issue",
                 "enablement": "typescript.native-preview.serverRunning",
                 "category": "TypeScript Native Preview"
+            },
+            {
+                "title": "Show References of CodeLens",
+                "command": "typescript.native-preview.codeLens.showLocations",
+                "enablement": "false"
             }
         ]
     },

_extension/src/client.ts

@@ -1,8 +1,11 @@
 import * as vscode from "vscode";
 import {
+    DocumentUri,
     LanguageClient,
     LanguageClientOptions,
+    Location,
     NotebookDocumentFilter,
+    Position,
     ServerOptions,
     TextDocumentFilter,
     TransportKind,
@@ -14,6 +17,8 @@ import {
 } from "./util";
 import { getLanguageForUri } from "./util";
 
+const codeLensShowLocationsCommandName = "typescript.native-preview.codeLens.showLocations";
+
 export class Client {
     private outputChannel: vscode.OutputChannel;
     private traceOutputChannel: vscode.OutputChannel;
@@ -32,6 +37,9 @@ export class Client {
             ],
             outputChannel: this.outputChannel,
             traceOutputChannel: this.traceOutputChannel,
+            initializationOptions: {
+                codeLensShowLocationsCommandName,
+            },
             diagnosticPullOptions: {
                 onChange: true,
                 onSave: true,
@@ -119,10 +127,34 @@ export class Client {
         await this.client.start();
         vscode.commands.executeCommand("setContext", "typescript.native-preview.serverRunning", true);
         this.onStartedCallbacks.forEach(callback => callback());
-        return new vscode.Disposable(() => {
-            if (this.client) {
-                this.client.stop();
+
+        const codeLensLocationsCommand = vscode.commands.registerCommand(codeLensShowLocationsCommandName, (...args: unknown[]) => {
+            if (args.length !== 3) {
+                throw new Error("Unexpected number of arguments.");
             }
+
+            const lspUri = args[0] as DocumentUri;
+            const lspPosition = args[1] as Position;
+            const lspLocations = args[2] as Location[];
+
+            const editorUri = vscode.Uri.parse(lspUri);
+            const editorPosition = new vscode.Position(lspPosition.line, lspPosition.character);
+            const editorLocations = lspLocations.map(loc =>
+                new vscode.Location(
+                    vscode.Uri.parse(loc.uri),
+                    new vscode.Range(
+                        new vscode.Position(loc.range.start.line, loc.range.start.character),
+                        new vscode.Position(loc.range.end.line, loc.range.end.character),
+                    ),
+                )
+            );
+
+            vscode.commands.executeCommand("editor.action.showReferences", editorUri, editorPosition, editorLocations);
+        });
+
+        return new vscode.Disposable(() => {
+            this.client?.stop();
+            codeLensLocationsCommand.dispose();
             vscode.commands.executeCommand("setContext", "typescript.native-preview.serverRunning", false);
         });
     }

internal/ast/ast.go

@@ -2214,11 +2214,11 @@ func IsWriteAccess(node *Node) bool {
 }
 
 func IsWriteAccessForReference(node *Node) bool {
-	decl := getDeclarationFromName(node)
+	decl := GetDeclarationFromName(node)
 	return (decl != nil && declarationIsWriteAccess(decl)) || node.Kind == KindDefaultKeyword || IsWriteAccess(node)
 }
 
-func getDeclarationFromName(name *Node) *Declaration {
+func GetDeclarationFromName(name *Node) *Declaration {
 	if name == nil || name.Parent == nil {
 		return nil
 	}

internal/diagnostics/diagnostics_generated.go

@@ -19,6 +19,22 @@
         "category": "Message",
         "code": 100004
     },
+    "{0} references": {
+        "category": "Message",
+        "code": 100005
+    },
+    "1 reference": {
+        "category": "Message",
+        "code": 100006
+    },
+    "{0} implementations": {
+        "category": "Message",
+        "code": 100007
+    },
+    "1 implementation": {
+        "category": "Message",
+        "code": 100008
+    },
     "Non-relative paths are not allowed. Did you forget a leading './'?": {
         "category": "Error",
         "code": 5090

internal/diagnostics/extraDiagnosticMessages.json

@@ -34,6 +34,7 @@ const (
 	renameCmd                   baselineCommand = "findRenameLocations"
 	signatureHelpCmd            baselineCommand = "SignatureHelp"
 	smartSelectionCmd           baselineCommand = "Smart Selection"
+	codeLensesCmd               baselineCommand = "Code Lenses"
 )
 
 type baselineCommand string

internal/fourslash/baselineutil.go

@@ -115,8 +115,8 @@ func (w *lspWriter) Write(msg *lsproto.Message) error {
 	return nil
 }
 
-func (r *lspWriter) Close() {
-	close(r.c)
+func (w *lspWriter) Close() {
+	close(w.c)
 }
 
 var (
@@ -294,13 +294,16 @@ func (f *FourslashTest) nextID() int32 {
 	return id
 }
 
+const showCodeLensLocationsCommandName = "typescript.showCodeLensLocations"
+
 func (f *FourslashTest) initialize(t *testing.T, capabilities *lsproto.ClientCapabilities) {
 	params := &lsproto.InitializeParams{
 		Locale: ptrTo("en-US"),
 		InitializationOptions: &lsproto.InitializationOptions{
 			// Hack: disable push diagnostics entirely, since the fourslash runner does not
 			// yet gracefully handle non-request messages.
-			DisablePushDiagnostics: ptrTo(true),
+			DisablePushDiagnostics:           ptrTo(true),
+			CodeLensShowLocationsCommandName: ptrTo(showCodeLensLocationsCommandName),
 		},
 	}
 	params.Capabilities = getCapabilitiesWithDefaults(capabilities)
@@ -1319,7 +1322,9 @@ func (f *FourslashTest) VerifyBaselineFindAllReferences(
 				Uri: lsconv.FileNameToDocumentURI(f.activeFilename),
 			},
 			Position: f.currentCaretPosition,
-			Context:  &lsproto.ReferenceContext{},
+			Context: &lsproto.ReferenceContext{
+				IncludeDeclaration: true,
+			},
 		}
 		result := sendRequest(t, f, lsproto.TextDocumentReferencesInfo, params)
 		f.addResultToBaseline(t, findAllReferencesCmd, f.getBaselineForLocationsWithFileContents(*result.Locations, baselineFourslashLocationsOptions{
@@ -1330,6 +1335,61 @@ func (f *FourslashTest) VerifyBaselineFindAllReferences(
 	}
 }
 
+func (f *FourslashTest) VerifyBaselineCodeLens(t *testing.T, preferences *lsutil.UserPreferences) {
+	if preferences != nil {
+		reset := f.ConfigureWithReset(t, preferences)
+		defer reset()
+	}
+
+	foundAtLeastOneCodeLens := false
+	for _, openFile := range slices.Sorted(maps.Keys(f.openFiles)) {
+		params := &lsproto.CodeLensParams{
+			TextDocument: lsproto.TextDocumentIdentifier{
+				Uri: lsconv.FileNameToDocumentURI(openFile),
+			},
+		}
+
+		unresolvedCodeLensList := sendRequest(t, f, lsproto.TextDocumentCodeLensInfo, params)
+		if unresolvedCodeLensList.CodeLenses == nil || len(*unresolvedCodeLensList.CodeLenses) == 0 {
+			continue
+		}
+		foundAtLeastOneCodeLens = true
+
+		for _, unresolvedCodeLens := range *unresolvedCodeLensList.CodeLenses {
+			assert.Assert(t, unresolvedCodeLens != nil)
+			resolvedCodeLens := sendRequest(t, f, lsproto.CodeLensResolveInfo, unresolvedCodeLens)
+			assert.Assert(t, resolvedCodeLens != nil)
+			assert.Assert(t, resolvedCodeLens.Command != nil, "Expected resolved code lens to have a command.")
+			if len(resolvedCodeLens.Command.Command) > 0 {
+				assert.Equal(t, resolvedCodeLens.Command.Command, showCodeLensLocationsCommandName)
+			}
+
+			var locations []lsproto.Location
+			// commandArgs: (DocumentUri, Position, Location[])
+			if commandArgs := resolvedCodeLens.Command.Arguments; commandArgs != nil {
+				locs, err := roundtripThroughJson[[]lsproto.Location]((*commandArgs)[2])
+				if err != nil {
+					t.Fatalf("failed to re-encode code lens locations: %v", err)
+				}
+				locations = locs
+			}
+
+			f.addResultToBaseline(t, codeLensesCmd, f.getBaselineForLocationsWithFileContents(locations, baselineFourslashLocationsOptions{
+				marker: &RangeMarker{
+					fileName: openFile,
+					LSRange:  resolvedCodeLens.Range,
+					Range:    f.converters.FromLSPRange(f.getScriptInfo(openFile), resolvedCodeLens.Range),
+				},
+				markerName: "/*CODELENS: " + resolvedCodeLens.Command.Title + "*/",
+			}))
+		}
+	}
+
+	if !foundAtLeastOneCodeLens {
+		t.Fatalf("Expected at least one code lens in any open file, but got none.")
+	}
+}
+
 func (f *FourslashTest) MarkTestAsStradaServer() {
 	f.isStradaServer = true
 }
@@ -2155,6 +2215,25 @@ func ptrTo[T any](v T) *T {
 	return &v
 }
 
+// This function is intended for spots where a complex
+// value needs to be reinterpreted following some prior JSON deserialization.
+// The default deserializer for `any` properties will give us a map at runtime,
+// but we want to validate against, and use, the types as returned from the the language service.
+//
+// Use this function sparingly. You can treat it as a "map-to-struct" converter,
+// but updating the original types is probably better in most cases.
+func roundtripThroughJson[T any](value any) (T, error) {
+	var result T
+	bytes, err := json.Marshal(value)
+	if err != nil {
+		return result, fmt.Errorf("failed to marshal value to JSON: %w", err)
+	}
+	if err := json.Unmarshal(bytes, &result); err != nil {
+		return result, fmt.Errorf("failed to unmarshal value from JSON: %w", err)
+	}
+	return result, nil
+}
+
 // Insert text at the current caret position.
 func (f *FourslashTest) Insert(t *testing.T, text string) {
 	f.typeText(t, text)

internal/fourslash/fourslash.go

@@ -0,0 +1,53 @@
+package fourslash_test
+
+import (
+	"testing"
+
+	"github.com/microsoft/typescript-go/internal/fourslash"
+	"github.com/microsoft/typescript-go/internal/ls/lsutil"
+	"github.com/microsoft/typescript-go/internal/testutil"
+)
+
+func TestCodeLensFunctionExpressions01(t *testing.T) {
+	t.Parallel()
+	defer testutil.RecoverAndFail(t, "Panic on fourslash test")
+
+	const content = `
+// @filename: anonymousFunctionExpressions.ts
+export let anonFn1 = function () {};
+export const anonFn2 = function () {};
+
+let anonFn3 = function () {};
+const anonFn4 = function () {};
+
+// @filename: arrowFunctions.ts
+export let arrowFn1 = () => {};
+export const arrowFn2 = () => {};
+
+let arrowFn3 = () => {};
+const arrowFn4 = () => {};
+
+// @filename: namedFunctions.ts
+export let namedFn1 = function namedFn1() {
+    namedFn1();
+}
+namedFn1();
+
+export const namedFn2 = function namedFn2() {
+    namedFn2();
+}
+namedFn2();
+
+let namedFn3 = function namedFn3() {};
+const namedFn4 = function namedFn4() {};
+`
+	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	f.VerifyBaselineCodeLens(t, &lsutil.UserPreferences{
+		ReferencesCodeLensEnabled:            true,
+		ReferencesCodeLensShowOnAllFunctions: true,
+
+		ImplementationsCodeLensEnabled:                true,
+		ImplementationsCodeLensShowOnInterfaceMethods: true,
+		ImplementationsCodeLensShowOnAllClassMethods:  true,
+	})
+}