Add Go to Source Definition (#3349)

Author: jakebailey

_extension/package.json

@@ -113,6 +113,12 @@
                 "enablement": "typescript.native-preview.serverRunning",
                 "category": "TypeScript Native Preview"
             },
+            {
+                "command": "typescript.native-preview.goToSourceDefinition",
+                "title": "Go to Source Definition",
+                "enablement": "typescript.native-preview.serverRunning",
+                "category": "TypeScript Native Preview"
+            },
             {
                 "title": "Show References of CodeLens",
                 "command": "typescript.native-preview.codeLens.showLocations",
@@ -154,7 +160,22 @@
                 "enablement": "typescript.native-preview.serverRunning",
                 "category": "Developer: TypeScript Native Preview"
             }
-        ]
+        ],
+        "menus": {
+            "commandPalette": [
+                {
+                    "command": "typescript.native-preview.goToSourceDefinition",
+                    "when": "typescript.native-preview.serverRunning && tsSupportsSourceDefinition"
+                }
+            ],
+            "editor/context": [
+                {
+                    "command": "typescript.native-preview.goToSourceDefinition",
+                    "when": "typescript.native-preview.serverRunning && tsSupportsSourceDefinition && (resourceLangId == typescript || resourceLangId == typescriptreact || resourceLangId == javascript || resourceLangId == javascriptreact)",
+                    "group": "navigation@1.41"
+                }
+            ]
+        }
     },
     "main": "./dist/extension.bundle.js",
     "files": [

_extension/src/client.ts

@@ -20,6 +20,7 @@ import {
     configurationMiddleware,
     sendNotificationMiddleware,
 } from "./configurationMiddleware";
+import { registerSourceDefinitionFeature } from "./languageFeatures/sourceDefinition";
 import { registerTagClosingFeature } from "./languageFeatures/tagClosing";
 import * as tr from "./telemetryReporting";
 import {
@@ -205,6 +206,7 @@ export class Client implements vscode.Disposable {
 
         this.disposables.push(
             serverTelemetryListener,
+            registerSourceDefinitionFeature(this.client),
             registerTagClosingFeature("typescript", this.documentSelector, this.client),
             registerTagClosingFeature("javascript", this.documentSelector, this.client),
         );

_extension/src/languageFeatures/sourceDefinition.ts

@@ -0,0 +1,84 @@
+import * as vscode from "vscode";
+import {
+    LanguageClient,
+    Location,
+    LocationLink,
+} from "vscode-languageclient/node";
+
+const sourceDefinitionMethod = "custom/textDocument/sourceDefinition";
+const sourceDefinitionCommand = "typescript.native-preview.goToSourceDefinition";
+const sourceDefinitionContext = "tsSupportsSourceDefinition";
+
+type SourceDefinitionResponse = Location | Location[] | LocationLink[] | null;
+
+export function registerSourceDefinitionFeature(client: LanguageClient): vscode.Disposable {
+    const capabilities = client.initializeResult?.capabilities as { customSourceDefinitionProvider?: boolean; } | undefined;
+    const enabled = !!capabilities?.customSourceDefinitionProvider;
+    void vscode.commands.executeCommand("setContext", sourceDefinitionContext, enabled);
+
+    if (!enabled) {
+        return new vscode.Disposable(() => {
+            void vscode.commands.executeCommand("setContext", sourceDefinitionContext, false);
+        });
+    }
+
+    const disposable = vscode.commands.registerCommand(sourceDefinitionCommand, async () => {
+        const activeEditor = vscode.window.activeTextEditor;
+        if (!activeEditor) {
+            vscode.window.showErrorMessage("Go to Source Definition failed. No editor is active.");
+            return;
+        }
+
+        const { document } = activeEditor;
+        if (!["javascript", "javascriptreact", "typescript", "typescriptreact"].includes(document.languageId)) {
+            vscode.window.showErrorMessage("Go to Source Definition failed. Unsupported file type.");
+            return;
+        }
+
+        const position = activeEditor.selection.active;
+        await vscode.window.withProgress({
+            location: vscode.ProgressLocation.Window,
+            title: "Finding source definitions",
+        }, async (_, token) => {
+            let response: SourceDefinitionResponse;
+            try {
+                response = await client.sendRequest<SourceDefinitionResponse>(
+                    sourceDefinitionMethod,
+                    client.code2ProtocolConverter.asTextDocumentPositionParams(document, position),
+                    token,
+                );
+            }
+            catch {
+                return;
+            }
+
+            if (token.isCancellationRequested) {
+                return;
+            }
+
+            const p2c = client.protocol2CodeConverter;
+            const items = !response ? [] : Array.isArray(response) ? response : [response];
+            const locations = items.map(item =>
+                LocationLink.is(item)
+                    ? p2c.asLocation({ uri: item.targetUri, range: item.targetSelectionRange })
+                    : p2c.asLocation(item)
+            );
+
+            await vscode.commands.executeCommand(
+                "editor.action.goToLocations",
+                document.uri,
+                position,
+                locations,
+                "goto",
+                "No source definitions found.",
+            );
+        });
+    });
+
+    return vscode.Disposable.from(
+        disposable,
+        new vscode.Disposable(() => {
+            void vscode.commands.executeCommand("setContext", sourceDefinitionContext, false);
+        }),
+    );
+}

internal/fourslash/_scripts/convertFourslash.mts

@@ -253,6 +253,7 @@ function parseFourslashStatement(statement: ts.Statement): Cmd[] {
                 case "baselineGetDefinitionAtPosition":
                 case "baselineGoToType":
                 case "baselineGoToImplementation":
+                case "baselineGoToSourceDefinition":
                     // Both `baselineGoToDefinition` and `baselineGetDefinitionAtPosition` take the same
                     // arguments, but differ in that...
                     //  - `verify.baselineGoToDefinition(...)` called getDefinitionAndBoundSpan
@@ -1289,14 +1290,14 @@ function parseBaselineDocumentHighlightsArgs(args: readonly ts.Expression[]): [V
 }
 
 function parseBaselineGoToDefinitionArgs(
-    funcName: "baselineGoToDefinition" | "baselineGoToType" | "baselineGetDefinitionAtPosition" | "baselineGoToImplementation",
+    funcName: "baselineGoToDefinition" | "baselineGoToType" | "baselineGetDefinitionAtPosition" | "baselineGoToImplementation" | "baselineGoToSourceDefinition",
     args: readonly ts.Expression[],
 ): [VerifyBaselineGoToDefinitionCmd] {
     let boundSpan: true | undefined;
     if (funcName === "baselineGoToDefinition") {
         boundSpan = true;
     }
-    let kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType" | "verifyBaselineGoToImplementation";
+    let kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType" | "verifyBaselineGoToImplementation" | "verifyBaselineGoToSourceDefinition";
     switch (funcName) {
         case "baselineGoToDefinition":
         case "baselineGetDefinitionAtPosition":
@@ -1308,6 +1309,9 @@ function parseBaselineGoToDefinitionArgs(
         case "baselineGoToImplementation":
             kind = "verifyBaselineGoToImplementation";
             break;
+        case "baselineGoToSourceDefinition":
+            kind = "verifyBaselineGoToSourceDefinition";
+            break;
     }
     const newArgs = [];
     for (const arg of args) {
@@ -2986,7 +2990,7 @@ interface VerifyBaselineFindAllReferencesCmd {
 }
 
 interface VerifyBaselineGoToDefinitionCmd {
-    kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType" | "verifyBaselineGoToImplementation";
+    kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType" | "verifyBaselineGoToImplementation" | "verifyBaselineGoToSourceDefinition";
     markers: string[];
     boundSpan?: true;
     ranges?: boolean;
@@ -3310,6 +3314,11 @@ function generateBaselineGoToDefinition({ markers, ranges, kind, boundSpan }: Ve
                 return `f.VerifyBaselineGoToImplementation(t)`;
             }
             return `f.VerifyBaselineGoToImplementation(t, ${markers.join(", ")})`;
+        case "verifyBaselineGoToSourceDefinition":
+            if (ranges || markers.length === 0) {
+                return `f.VerifyBaselineGoToSourceDefinition(t)`;
+            }
+            return `f.VerifyBaselineGoToSourceDefinition(t, ${markers.join(", ")})`;
     }
 }
 
@@ -3502,6 +3511,7 @@ function generateCmd(cmd: Cmd): string {
         case "verifyBaselineGoToDefinition":
         case "verifyBaselineGoToType":
         case "verifyBaselineGoToImplementation":
+        case "verifyBaselineGoToSourceDefinition":
             return generateBaselineGoToDefinition(cmd);
         case "verifyBaselineQuickInfo":
             // Quick Info -> Hover

internal/fourslash/_scripts/unparsedTests.txt

@@ -1737,23 +1737,6 @@ getMatchingBracesAdjacentBraces.ts parse error: "Unrecognized fourslash statemen
 getMatchingBracesNegativeCases.ts parse error: "Unrecognized fourslash statement: test.ranges().forEach(...)"
 getNameOrDottedNameSpan.ts parse error: "Unrecognized fourslash statement: verify.nameOrDottedNameSpanTextIs(...)"
 globalCompletionListInsideObjectLiterals.ts parse error: "Expected string literal or object literal for expected completion item, got exact.filter(name => name !== 'p1')"
-goToSource1_localJsBesideDts.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource10_mapFromAtTypes3.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource11_propertyOfAlias.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource12_callbackParam.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource13_nodenext.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource14_unresolvedRequireDestructuring.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource15_bundler.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource16_callbackParamDifferentFile.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource17_AddsFileToProject.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource18_reusedFromDifferentFolder.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource2_nodeModulesWithTypes.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource3_nodeModulesAtTypes.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource5_sameAsGoToDef1.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource6_sameAsGoToDef2.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource7_conditionallyMinified.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource8_mapFromAtTypes.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
-goToSource9_mapFromAtTypes2.ts parse error: "Unrecognized fourslash statement: verify.baselineGoToSourceDefinition(...)"
 identationAfterInterfaceCall.ts parse error: "Unrecognized verify content function: indentationIs"
 impliedNodeFormat.ts parse error: "Expected a single string literal argument in edit.paste, got `\\n\"${\"a\".repeat(256)}\";`"
 importDeclPaste0.ts parse error: "Unrecognized edit function: disableFormatting"

internal/fourslash/baselineutil.go

@@ -28,6 +28,7 @@ const (
 	findAllReferencesCmd        baselineCommand = "findAllReferences"
 	goToDefinitionCmd           baselineCommand = "goToDefinition"
 	goToImplementationCmd       baselineCommand = "goToImplementation"
+	goToSourceDefinitionCmd     baselineCommand = "goToSourceDefinition"
 	goToTypeDefinitionCmd       baselineCommand = "goToType"
 	inlayHintsCmd               baselineCommand = "Inlay Hints"
 	nonSuggestionDiagnosticsCmd baselineCommand = "Syntax and Semantic Diagnostics"
@@ -274,7 +275,7 @@ func (f *FourslashTest) getBaselineOptions(command baselineCommand, testPath str
 				return strings.Join(fixedLines, "\n")
 			},
 		}
-	case goToDefinitionCmd, goToTypeDefinitionCmd, goToImplementationCmd:
+	case goToDefinitionCmd, goToTypeDefinitionCmd, goToImplementationCmd, goToSourceDefinitionCmd:
 		return baseline.Options{
 			Subfolder:   subfolder,
 			IsSubmodule: true,
@@ -515,7 +516,9 @@ type baselineFourslashLocationsOptions struct {
 	endMarkerSuffix   func(span documentSpan) *string
 	getLocationData   func(span documentSpan) string
 
-	additionalSpan *documentSpan
+	additionalSpan      *documentSpan
+	preserveResultOrder bool
+	orderedFiles        []lsproto.DocumentUri
 }
 
 func locationToSpan(loc lsproto.Location) documentSpan {
@@ -534,6 +537,9 @@ func (f *FourslashTest) getBaselineForLocationsWithFileContents(locations []lspr
 
 func (f *FourslashTest) getBaselineForSpansWithFileContents(spans []documentSpan, options baselineFourslashLocationsOptions) string {
 	spansByFile := collections.GroupBy(spans, func(span documentSpan) lsproto.DocumentUri { return span.uri })
+	if options.preserveResultOrder {
+		options.orderedFiles = uniqueFilesInSpanOrder(spans)
+	}
 	return f.getBaselineForGroupedSpansWithFileContents(
 		spansByFile,
 		options,
@@ -549,25 +555,16 @@ func (f *FourslashTest) getBaselineForGroupedSpansWithFileContents(groupedRanges
 	spanToContextId := map[documentSpan]int{}
 
 	baselineEntries := []string{}
-	walkDirFn := func(path string, d vfs.DirEntry, e error) error {
-		if e != nil {
-			return e
-		}
-
-		if !d.Type().IsRegular() {
-			return nil
-		}
-
+	addFileEntry := func(path string) {
 		fileName := lsconv.FileNameToDocumentURI(path)
 		ranges := groupedRanges.Get(fileName)
 		if len(ranges) == 0 {
-			return nil
+			return
 		}
 
 		content, ok := f.textOfFile(path)
 		if !ok {
-			// !!! error?
-			return nil
+			return
 		}
 
 		if options.marker != nil && options.marker.FileName() == path {
@@ -579,17 +576,34 @@ func (f *FourslashTest) getBaselineForGroupedSpansWithFileContents(groupedRanges
 		}
 
 		baselineEntries = append(baselineEntries, f.getBaselineContentForFile(path, content, ranges, spanToContextId, options))
-		return nil
 	}
+	walkDirFn := func(path string, d vfs.DirEntry, e error) error {
+		if e != nil {
+			return e
+		}
+
+		if !d.Type().IsRegular() {
+			return nil
+		}
 
-	err := f.vfs.WalkDir("/", walkDirFn)
-	if err != nil && !errors.Is(err, fs.ErrNotExist) {
-		panic("walkdir error during fourslash baseline: " + err.Error())
+		addFileEntry(path)
+		return nil
 	}
 
-	err = f.vfs.WalkDir("bundled:///", walkDirFn)
-	if err != nil && !errors.Is(err, fs.ErrNotExist) {
-		panic("walkdir error during fourslash baseline: " + err.Error())
+	if options.preserveResultOrder {
+		for _, uri := range options.orderedFiles {
+			addFileEntry(uri.FileName())
+		}
+	} else {
+		err := f.vfs.WalkDir("/", walkDirFn)
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			panic("walkdir error during fourslash baseline: " + err.Error())
+		}
+
+		err = f.vfs.WalkDir("bundled:///", walkDirFn)
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			panic("walkdir error during fourslash baseline: " + err.Error())
+		}
 	}
 
 	// In Strada, there is a bug where we only ever add additional spans to baselines if we haven't
@@ -620,6 +634,22 @@ func (f *FourslashTest) getBaselineForGroupedSpansWithFileContents(groupedRanges
 	return strings.Join(baselineEntries, "\n\n")
 }
 
+func uniqueFilesInSpanOrder(spans []documentSpan) []lsproto.DocumentUri {
+	if len(spans) == 0 {
+		return nil
+	}
+	seen := map[lsproto.DocumentUri]struct{}{}
+	result := make([]lsproto.DocumentUri, 0, len(spans))
+	for _, span := range spans {
+		if _, ok := seen[span.uri]; ok {
+			continue
+		}
+		seen[span.uri] = struct{}{}
+		result = append(result, span.uri)
+	}
+	return result
+}
+
 func (f *FourslashTest) textOfFile(fileName string) (string, bool) {
 	if _, ok := f.openFiles[fileName]; ok {
 		return f.getScriptInfo(fileName).content, true