Use LocationLink in go to definition (#1884)

Author: gabritto

internal/fourslash/_scripts/convertFourslash.mts

@@ -200,7 +200,6 @@ function parseFourslashStatement(statement: ts.Statement): Cmd[] | undefined {
                     //  - `verify.baselineGoToDefinition(...)` called getDefinitionAndBoundSpan
                     //  - `verify.baselineGetDefinitionAtPosition(...)` called getDefinitionAtPosition
                     // LSP doesn't have two separate commands though.
-                    // It's unclear how we would model bound spans though.
                     return parseBaselineGoToDefinitionArgs(func.text, callExpression.arguments);
                 case "baselineRename":
                 case "baselineRenameAtRangesWithText":
@@ -1081,6 +1080,10 @@ function parseBaselineGoToDefinitionArgs(
     funcName: "baselineGoToDefinition" | "baselineGoToType" | "baselineGetDefinitionAtPosition",
     args: readonly ts.Expression[],
 ): [VerifyBaselineGoToDefinitionCmd] | undefined {
+    let boundSpan: true | undefined;
+    if (funcName === "baselineGoToDefinition") {
+        boundSpan = true;
+    }
     let kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType";
     switch (funcName) {
         case "baselineGoToDefinition":
@@ -1105,6 +1108,7 @@ function parseBaselineGoToDefinitionArgs(
                 kind,
                 markers: [],
                 ranges: true,
+                boundSpan,
             }];
         }
         else {
@@ -1116,6 +1120,7 @@ function parseBaselineGoToDefinitionArgs(
     return [{
         kind,
         markers: newArgs,
+        boundSpan,
     }];
 }
 
@@ -1592,6 +1597,7 @@ interface VerifyBaselineFindAllReferencesCmd {
 interface VerifyBaselineGoToDefinitionCmd {
     kind: "verifyBaselineGoToDefinition" | "verifyBaselineGoToType";
     markers: string[];
+    boundSpan?: true;
     ranges?: boolean;
 }
 
@@ -1710,20 +1716,20 @@ function generateBaselineDocumentHighlights({ args, preferences }: VerifyBaselin
     return `f.VerifyBaselineDocumentHighlights(t, ${preferences}, ${args.join(", ")})`;
 }
 
-function generateBaselineGoToDefinition({ markers, ranges, kind }: VerifyBaselineGoToDefinitionCmd): string {
-    let goFunc;
+function generateBaselineGoToDefinition({ markers, ranges, kind, boundSpan }: VerifyBaselineGoToDefinitionCmd): string {
+    const originalSelectionRange = boundSpan ? "true" : "false";
     switch (kind) {
         case "verifyBaselineGoToDefinition":
-            goFunc = "VerifyBaselineGoToDefinition";
-            break;
+            if (ranges || markers.length === 0) {
+                return `f.VerifyBaselineGoToDefinition(t, ${originalSelectionRange})`;
+            }
+            return `f.VerifyBaselineGoToDefinition(t, ${originalSelectionRange}, ${markers.join(", ")})`;
         case "verifyBaselineGoToType":
-            goFunc = "VerifyBaselineGoToTypeDefinition";
-            break;
-    }
-    if (ranges || markers.length === 0) {
-        return `f.${goFunc}(t)`;
+            if (ranges || markers.length === 0) {
+                return `f.VerifyBaselineGoToTypeDefinition(t)`;
+            }
+            return `f.VerifyBaselineGoToTypeDefinition(t, ${markers.join(", ")})`;
     }
-    return `f.${goFunc}(t, ${markers.join(", ")})`;
 }
 
 function generateGoToCommand({ funcName, args }: GoToCmd): string {

internal/fourslash/baselineutil.go

@@ -131,6 +131,8 @@ type baselineFourslashLocationsOptions struct {
 
 	startMarkerPrefix func(span lsproto.Location) *string
 	endMarkerSuffix   func(span lsproto.Location) *string
+
+	additionalLocation *lsproto.Location
 }
 
 func (f *FourslashTest) getBaselineForLocationsWithFileContents(spans []lsproto.Location, options baselineFourslashLocationsOptions) string {
@@ -152,6 +154,7 @@ func (f *FourslashTest) getBaselineForGroupedLocationsWithFileContents(groupedRa
 	// but don't want to print it twice at the end if it already
 	// found in a file with ranges.
 	foundMarker := false
+	foundAdditionalLocation := false
 
 	baselineEntries := []string{}
 	err := f.vfs.WalkDir("/", func(path string, d vfs.DirEntry, e error) error {
@@ -179,6 +182,10 @@ func (f *FourslashTest) getBaselineForGroupedLocationsWithFileContents(groupedRa
 			foundMarker = true
 		}
 
+		if options.additionalLocation != nil && options.additionalLocation.Uri == fileName {
+			foundAdditionalLocation = true
+		}
+
 		baselineEntries = append(baselineEntries, f.getBaselineContentForFile(path, content, ranges, nil, options))
 		return nil
 	})
@@ -187,6 +194,21 @@ func (f *FourslashTest) getBaselineForGroupedLocationsWithFileContents(groupedRa
 		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
+	// already added the file to the baseline.
+	if options.additionalLocation != nil && !foundAdditionalLocation {
+		fileName := options.additionalLocation.Uri.FileName()
+		if content, ok := f.vfs.ReadFile(fileName); ok {
+			baselineEntries = append(
+				baselineEntries,
+				f.getBaselineContentForFile(fileName, content, []lsproto.Range{options.additionalLocation.Range}, nil, options),
+			)
+			if options.marker != nil && options.marker.FileName() == fileName {
+				foundMarker = true
+			}
+		}
+	}
+
 	if !foundMarker && options.marker != nil {
 		// If we didn't find the marker in any file, we need to add it.
 		markerFileName := options.marker.FileName()
@@ -195,7 +217,6 @@ func (f *FourslashTest) getBaselineForGroupedLocationsWithFileContents(groupedRa
 		}
 	}
 
-	// !!! foundAdditionalSpan
 	// !!! skipDocumentContainingOnlyMarker
 
 	return strings.Join(baselineEntries, "\n\n")

internal/fourslash/fourslash.go

@@ -245,6 +245,12 @@ var (
 			ItemDefaults: &[]string{"commitCharacters", "editRange"},
 		},
 	}
+	defaultDefinitionCapabilities = &lsproto.DefinitionClientCapabilities{
+		LinkSupport: ptrTrue,
+	}
+	defaultTypeDefinitionCapabilities = &lsproto.TypeDefinitionClientCapabilities{
+		LinkSupport: ptrTrue,
+	}
 )
 
 func getCapabilitiesWithDefaults(capabilities *lsproto.ClientCapabilities) *lsproto.ClientCapabilities {
@@ -278,6 +284,12 @@ func getCapabilitiesWithDefaults(capabilities *lsproto.ClientCapabilities) *lspr
 	if capabilitiesWithDefaults.Workspace.Configuration == nil {
 		capabilitiesWithDefaults.Workspace.Configuration = ptrTrue
 	}
+	if capabilitiesWithDefaults.TextDocument.Definition == nil {
+		capabilitiesWithDefaults.TextDocument.Definition = defaultDefinitionCapabilities
+	}
+	if capabilitiesWithDefaults.TextDocument.TypeDefinition == nil {
+		capabilitiesWithDefaults.TextDocument.TypeDefinition = defaultTypeDefinitionCapabilities
+	}
 	return &capabilitiesWithDefaults
 }
 
@@ -1032,6 +1044,7 @@ func (f *FourslashTest) VerifyBaselineFindAllReferences(
 
 func (f *FourslashTest) VerifyBaselineGoToDefinition(
 	t *testing.T,
+	includeOriginalSelectionRange bool,
 	markers ...string,
 ) {
 	referenceLocations := f.lookupMarkersOrGetRanges(t, markers)
@@ -1064,17 +1077,35 @@ func (f *FourslashTest) VerifyBaselineGoToDefinition(
 		}
 
 		var resultAsLocations []lsproto.Location
+		var additionalLocation *lsproto.Location
 		if result.Locations != nil {
 			resultAsLocations = *result.Locations
 		} else if result.Location != nil {
 			resultAsLocations = []lsproto.Location{*result.Location}
 		} else if result.DefinitionLinks != nil {
-			t.Fatalf("Unexpected definition response type at marker '%s': %T", *f.lastKnownMarkerName, result.DefinitionLinks)
+			var originRange *lsproto.Range
+			resultAsLocations = core.Map(*result.DefinitionLinks, func(link *lsproto.LocationLink) lsproto.Location {
+				if originRange != nil && originRange != link.OriginSelectionRange {
+					panic("multiple different origin ranges in definition links")
+				}
+				originRange = link.OriginSelectionRange
+				return lsproto.Location{
+					Uri:   link.TargetUri,
+					Range: link.TargetSelectionRange,
+				}
+			})
+			if originRange != nil && includeOriginalSelectionRange {
+				additionalLocation = &lsproto.Location{
+					Uri:   lsconv.FileNameToDocumentURI(f.activeFilename),
+					Range: *originRange,
+				}
+			}
 		}
 
 		f.addResultToBaseline(t, "goToDefinition", f.getBaselineForLocationsWithFileContents(resultAsLocations, baselineFourslashLocationsOptions{
-			marker:     markerOrRange,
-			markerName: "/*GOTO DEF*/",
+			marker:             markerOrRange,
+			markerName:         "/*GOTO DEF*/",
+			additionalLocation: additionalLocation,
 		}))
 	}
 }
@@ -1118,7 +1149,12 @@ func (f *FourslashTest) VerifyBaselineGoToTypeDefinition(
 		} else if result.Location != nil {
 			resultAsLocations = []lsproto.Location{*result.Location}
 		} else if result.DefinitionLinks != nil {
-			t.Fatalf("Unexpected type definition response type at marker '%s': %T", *f.lastKnownMarkerName, result.DefinitionLinks)
+			resultAsLocations = core.Map(*result.DefinitionLinks, func(link *lsproto.LocationLink) lsproto.Location {
+				return lsproto.Location{
+					Uri:   link.TargetUri,
+					Range: link.TargetSelectionRange,
+				}
+			})
 		}
 
 		f.addResultToBaseline(t, "goToType", f.getBaselineForLocationsWithFileContents(resultAsLocations, baselineFourslashLocationsOptions{

internal/fourslash/tests/gen/ambientShorthandGotoDefinition_test.go

@@ -24,5 +24,5 @@ import [|/*importBang*/bang|] = require("jquery");
 	f.VerifyQuickInfoAt(t, "useBar", "(alias) module \"jquery\"\nimport bar", "")
 	f.VerifyQuickInfoAt(t, "useBaz", "(alias) module \"jquery\"\nimport baz", "")
 	f.VerifyQuickInfoAt(t, "useBang", "(alias) module \"jquery\"\nimport bang = require(\"jquery\")", "")
-	f.VerifyBaselineGoToDefinition(t, "useFoo", "importFoo", "useBar", "useBaz", "importBaz", "useBang", "importBang")
+	f.VerifyBaselineGoToDefinition(t, true, "useFoo", "importFoo", "useBar", "useBaz", "importBaz", "useBang", "importBang")
 }

internal/fourslash/tests/gen/declarationMapGoToDefinition_test.go

@@ -49,5 +49,5 @@ import * as mod from "./indexdef";
 const instance = new mod.Foo();
 instance.[|/*1*/methodName|]({member: 12});`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
-	f.VerifyBaselineGoToDefinition(t, "1")
+	f.VerifyBaselineGoToDefinition(t, true, "1")
 }

internal/fourslash/tests/gen/declarationMapsGoToDefinitionRelativeSourceRoot_test.go

@@ -49,5 +49,5 @@ import * as mod from "./out/indexdef";
 const instance = new mod.Foo();
 instance.[|/*1*/methodName|]({member: 12});`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
-	f.VerifyBaselineGoToDefinition(t, "1")
+	f.VerifyBaselineGoToDefinition(t, true, "1")
 }

internal/fourslash/tests/gen/declarationMapsGoToDefinitionSameNameDifferentDirectory_test.go

@@ -66,5 +66,5 @@ class Button extends [|/*1*/Control|] {
     }
 }`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
-	f.VerifyBaselineGoToDefinition(t, "1", "3")
+	f.VerifyBaselineGoToDefinition(t, true, "1", "3")
 }

internal/fourslash/tests/gen/declarationMapsOutOfDateMapping_test.go

@@ -34,5 +34,5 @@ export class /*2*/Foo {
 import { Foo/*1*/ } from "a";`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
 	f.GoToFile(t, "/home/src/workspaces/project/index.ts")
-	f.VerifyBaselineGoToDefinition(t, "1")
+	f.VerifyBaselineGoToDefinition(t, false, "1")
 }

internal/fourslash/tests/gen/definition01_test.go

@@ -17,5 +17,5 @@ var x = new n.Foo();
 // @Filename: a.ts
  /*2*/export class Foo {}`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
-	f.VerifyBaselineGoToDefinition(t, "1")
+	f.VerifyBaselineGoToDefinition(t, true, "1")
 }

internal/fourslash/tests/gen/definitionNameOnEnumMember_test.go

@@ -18,5 +18,5 @@ func TestDefinitionNameOnEnumMember(t *testing.T) {
 }
 var enumMember = e.[|/*1*/thirdMember|];`
 	f := fourslash.NewFourslash(t, nil /*capabilities*/, content)
-	f.VerifyBaselineGoToDefinition(t, "1")
+	f.VerifyBaselineGoToDefinition(t, false, "1")
 }