Implement support for Code Lenses

_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"
             }
         ]
     },
@@ -115,4 +126,4 @@
         "@vscode/vsce": "^3.7.0",
         "esbuild": "^0.27.0"
     }
-}
+}

_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/extraDiagnosticMessages.json

@@ -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/fourslash/fourslash.go

@@ -1235,7 +1235,9 @@ func (f *FourslashTest) VerifyBaselineFindAllReferences(
 				Uri: lsconv.FileNameToDocumentURI(f.activeFilename),
 			},
 			Position: f.currentCaretPosition,
-			Context:  &lsproto.ReferenceContext{},
+			Context: &lsproto.ReferenceContext{
+				IncludeDeclaration: true,
+			},
 		}
 		resMsg, result, resultOk := sendRequest(t, f, lsproto.TextDocumentReferencesInfo, params)
 		if resMsg == nil {

internal/ls/codelens.go

@@ -0,0 +1,196 @@
+package ls
+
+import (
+	"context"
+
+	"github.com/microsoft/typescript-go/internal/ast"
+	"github.com/microsoft/typescript-go/internal/core"
+	"github.com/microsoft/typescript-go/internal/diagnostics"
+	"github.com/microsoft/typescript-go/internal/ls/lsutil"
+	"github.com/microsoft/typescript-go/internal/lsp/lsproto"
+	"github.com/microsoft/typescript-go/internal/scanner"
+)
+
+func (l *LanguageService) ProvideCodeLenses(ctx context.Context, documentURI lsproto.DocumentUri) (lsproto.CodeLensResponse, error) {
+	_, file := l.getProgramAndFile(documentURI)
+
+	userPrefs := l.UserPreferences()
+	if !userPrefs.ReferencesCodeLensEnabled && !userPrefs.ImplementationsCodeLensEnabled {
+		return lsproto.CodeLensResponse{}, nil
+	}
+
+	var lenses []*lsproto.CodeLens
+	var visit func(node *ast.Node) bool
+	visit = func(node *ast.Node) bool {
+		if ctx.Err() != nil {
+			return true
+		}
+
+		if userPrefs.ReferencesCodeLensEnabled && isValidReferenceLensNode(node, userPrefs) {
+			lenses = append(lenses, l.newCodeLensForNode(documentURI, file, node, lsproto.CodeLensKindReferences))
+		}
+
+		if userPrefs.ImplementationsCodeLensEnabled && isValidImplementationsCodeLensNode(node, userPrefs) {
+			lenses = append(lenses, l.newCodeLensForNode(documentURI, file, node, lsproto.CodeLensKindImplementations))
+		}
+
+		node.ForEachChild(visit)
+		return false
+	}
+
+	visit(file.AsNode())
+
+	return lsproto.CodeLensResponse{
+		CodeLenses: &lenses,
+	}, nil
+}
+
+func (l *LanguageService) ResolveCodeLens(ctx context.Context, codeLens *lsproto.CodeLens, showLocationsCommandName *string) (*lsproto.CodeLens, error) {
+	uri := codeLens.Data.Uri
+	textDoc := lsproto.TextDocumentIdentifier{
+		Uri: uri,
+	}
+
+	var locs []lsproto.Location
+	var lensTitle string
+	switch codeLens.Data.Kind {
+	case lsproto.CodeLensKindReferences:
+		references, err := l.ProvideReferences(ctx, &lsproto.ReferenceParams{
+			TextDocument: textDoc,
+			Position:     codeLens.Range.Start,
+			Context: &lsproto.ReferenceContext{
+				// Don't include the declaration in the references count.
+				IncludeDeclaration: false,
+			},
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		if references.Locations != nil {
+			locs = *references.Locations
+		}
+
+		if len(locs) == 1 {
+			lensTitle = diagnostics.X_1_reference.Message()
+		} else {
+			lensTitle = diagnostics.X_0_references.Format(len(locs))
+		}
+	case lsproto.CodeLensKindImplementations:
+		// "Force" link support to be false so that we only get `Locations` back,
+		// and don't include the "current" node in the results.
+		findImplsOptions := provideImplementationsOpts{
+			requireLocationsResult: true,
+			dropOriginNodes:        true,
+		}
+		implementations, err := l.provideImplementationsEx(ctx, &lsproto.ImplementationParams{
+			TextDocument: textDoc,
+			Position:     codeLens.Range.Start,
+		}, findImplsOptions)
+		if err != nil {
+			return nil, err
+		}
+
+		if implementations.Locations != nil {
+			locs = *implementations.Locations
+		}
+
+		if len(locs) == 1 {
+			lensTitle = diagnostics.X_1_implementation.Message()
+		} else {
+			lensTitle = diagnostics.X_0_implementations.Format(len(locs))
+		}
+	}
+
+	cmd := &lsproto.Command{
+		Title: lensTitle,
+	}
+	if len(locs) > 0 && showLocationsCommandName != nil {
+		cmd.Command = *showLocationsCommandName
+		cmd.Arguments = &[]any{
+			uri,
+			codeLens.Range.Start,
+			locs,
+		}
+	}
+
+	codeLens.Command = cmd
+	return codeLens, nil
+}
+
+func (l *LanguageService) newCodeLensForNode(fileUri lsproto.DocumentUri, file *ast.SourceFile, node *ast.Node, kind lsproto.CodeLensKind) *lsproto.CodeLens {
+	nodeForRange := node
+	nodeName := node.Name()
+	if nodeName != nil {
+		nodeForRange = nodeName
+	}
+	pos := scanner.SkipTrivia(file.Text(), nodeForRange.Pos())
+
+	return &lsproto.CodeLens{
+		Range: lsproto.Range{
+			Start: l.converters.PositionToLineAndCharacter(file, core.TextPos(pos)),
+			End:   l.converters.PositionToLineAndCharacter(file, core.TextPos(node.End())),
+		},
+		Data: &lsproto.CodeLensData{
+			Kind: kind,
+			Uri:  fileUri,
+		},
+	}
+}
+
+func isValidImplementationsCodeLensNode(node *ast.Node, userPrefs *lsutil.UserPreferences) bool {
+	switch node.Kind {
+	// Always show on interfaces
+	case ast.KindInterfaceDeclaration:
+		// TODO: ast.KindTypeAliasDeclaration?
+		return true
+
+	// If configured, show on interface methods
+	case ast.KindMethodSignature:
+		return userPrefs.ImplementationsCodeLensShowOnInterfaceMethods && node.Parent.Kind == ast.KindInterfaceDeclaration
+
+	// If configured, show on all class methods - but not private ones.
+	case ast.KindMethodDeclaration:
+		if userPrefs.ImplementationsCodeLensShowOnAllClassMethods && node.Parent.Kind == ast.KindClassDeclaration {
+			return !ast.HasModifier(node, ast.ModifierFlagsPrivate) && node.Name().Kind != ast.KindPrivateIdentifier
+		}
+		fallthrough
+
+	// Always show on abstract classes/properties/methods
+	case ast.KindClassDeclaration, ast.KindConstructor,
+		ast.KindGetAccessor, ast.KindSetAccessor, ast.KindPropertyDeclaration:
+		return ast.HasModifier(node, ast.ModifierFlagsAbstract)
+	}
+
+	return false
+}
+
+func isValidReferenceLensNode(node *ast.Node, userPrefs *lsutil.UserPreferences) bool {
+	switch node.Kind {
+	case ast.KindFunctionDeclaration, ast.KindFunctionExpression:
+		if userPrefs.ReferencesCodeLensShowOnAllFunctions {
+			return true
+		}
+		fallthrough
+
+	case ast.KindVariableDeclaration:
+		return ast.GetCombinedModifierFlags(node)&ast.ModifierFlagsExport != 0
+
+	case ast.KindClassDeclaration, ast.KindInterfaceDeclaration, ast.KindTypeAliasDeclaration, ast.KindEnumDeclaration, ast.KindEnumMember:
+		return true
+
+	case ast.KindMethodDeclaration, ast.KindMethodSignature, ast.KindConstructor,
+		ast.KindGetAccessor, ast.KindSetAccessor,
+		ast.KindPropertyDeclaration, ast.KindPropertySignature:
+		// Don't show if child and parent have same start
+		// For https://github.com/microsoft/vscode/issues/90396
+		// !!!
+
+		switch node.Parent.Kind {
+		case ast.KindClassDeclaration, ast.KindInterfaceDeclaration, ast.KindTypeLiteral:
+			return true
+		}
+	}
+
+	return false
+}