[LSP] fix go to definition

internal/ls/definitionsourcemap.go

@@ -0,0 +1,285 @@
+package ls
+
+import (
+	"sort"
+	"strings"
+
+	"github.com/go-json-experiment/json"
+	"github.com/microsoft/typescript-go/internal/compiler"
+	"github.com/microsoft/typescript-go/internal/core"
+	"github.com/microsoft/typescript-go/internal/lsp/lsproto"
+	"github.com/microsoft/typescript-go/internal/sourcemap"
+	"github.com/microsoft/typescript-go/internal/tspath"
+)
+
+type DefinitionSourceMapper struct {
+	program *compiler.Program
+}
+
+func NewDefinitionSourceMapper(program *compiler.Program) *DefinitionSourceMapper {
+	return &DefinitionSourceMapper{
+		program: program,
+	}
+}
+
+// FallbackFileReader is an optional interface that VFS implementations can implement
+// to provide fallback file reading for files not tracked by the program.
+type FallbackFileReader interface {
+	ReadFileWithFallback(path string) (string, bool)
+}
+
+// readFileWithFallback tries to read a file from the program's FS first,
+// and falls back to underlying FS for files not tracked by the program
+// (like .d.ts.map files or source files outside the TypeScript program).
+func (dsm *DefinitionSourceMapper) readFileWithFallback(fileName string) (string, bool) {
+	if fallbackFS, ok := dsm.program.Host().FS().(FallbackFileReader); ok {
+		return fallbackFS.ReadFileWithFallback(fileName)
+	}
+	return "", false
+}
+
+func (dsm *DefinitionSourceMapper) MapDefinitionLocations(locations []lsproto.Location) []lsproto.Location {
+	mappedLocations := make([]lsproto.Location, 0, len(locations))
+
+	for _, location := range locations {
+		if mappedLocation := dsm.MapSingleLocation(location); mappedLocation != nil {
+			mappedLocations = core.AppendIfUnique(mappedLocations, *mappedLocation)
+		} else {
+			mappedLocations = core.AppendIfUnique(mappedLocations, location)
+		}
+	}
+
+	return mappedLocations
+}
+
+func (dsm *DefinitionSourceMapper) MapSingleLocation(location lsproto.Location) *lsproto.Location {
+	fileName := location.Uri.FileName()
+
+	if !strings.HasSuffix(fileName, ".d.ts") {
+		return nil
+	}
+
+	if strings.HasPrefix(fileName, "^/bundled/") {
+		return nil
+	}
+
+	return dsm.tryGetSourcePosition(fileName, location)
+}
+
+func (dsm *DefinitionSourceMapper) tryGetSourcePosition(fileName string, location lsproto.Location) *lsproto.Location {
+	content, ok := dsm.readFileWithFallback(fileName)
+	if !ok {
+		return nil
+	}
+
+	mapURL := tryGetSourceMappingURL(content)
+	if mapURL == "" {
+		return nil
+	}
+
+	var mapFileName string
+	if strings.HasPrefix(mapURL, "/") {
+		mapFileName = mapURL
+	} else {
+		dir := tspath.GetDirectoryPath(fileName)
+		mapFileName = tspath.CombinePaths(dir, mapURL)
+	}
+
+	mapContent, ok := dsm.readFileWithFallback(mapFileName)
+	if !ok {
+		return nil
+	}
+
+	var sourceMap SourceMap
+	if err := json.Unmarshal([]byte(mapContent), &sourceMap); err != nil {
+		return nil
+	}
+
+	if len(sourceMap.Sources) == 0 {
+		return nil
+	}
+
+	mapDir := tspath.GetDirectoryPath(mapFileName)
+	declLineStarts := core.ComputeLineStarts(content)
+	declPos := int(declLineStarts[location.Range.Start.Line]) + int(location.Range.Start.Character)
+
+	sourceRange := dsm.getSourceRangeFromMappings(DocumentPosition{FileName: fileName, Pos: declPos}, sourceMap, mapDir)
+	if sourceRange == nil {
+		return nil
+	}
+
+	sourceFileContent, ok := dsm.readFileWithFallback(sourceRange.FileName)
+	if !ok {
+		return nil
+	}
+
+	sourceLineStarts := core.ComputeLineStarts(sourceFileContent)
+	sourceStartLine, sourceStartChar := core.PositionToLineAndCharacter(sourceRange.Start, sourceLineStarts)
+	sourceEndLine, sourceEndChar := core.PositionToLineAndCharacter(sourceRange.End, sourceLineStarts)
+
+	return &lsproto.Location{
+		Uri: FileNameToDocumentURI(sourceRange.FileName),
+		Range: lsproto.Range{
+			Start: lsproto.Position{Line: uint32(sourceStartLine), Character: uint32(sourceStartChar)},
+			End:   lsproto.Position{Line: uint32(sourceEndLine), Character: uint32(sourceEndChar)},
+		},
+	}
+}
+
+type SourceRange struct {
+	FileName string
+	Start    int
+	End      int
+}
+
+func (dsm *DefinitionSourceMapper) getSourceRangeFromMappings(loc DocumentPosition, sourceMap SourceMap, mapDir string) *SourceRange {
+	if len(sourceMap.Sources) == 0 || sourceMap.Mappings == "" {
+		return nil
+	}
+
+	declContent, ok := dsm.readFileWithFallback(loc.FileName)
+	if !ok {
+		return nil
+	}
+	declLineStarts := core.ComputeLineStarts(declContent)
+
+	mappings := decodeMappings(sourceMap.Mappings, declLineStarts)
+	if len(mappings) == 0 {
+		return nil
+	}
+
+	targetIndex := sort.Search(len(mappings), func(i int) bool {
+		return mappings[i].GeneratedPosition > loc.Pos
+	}) - 1
+
+	if targetIndex < 0 {
+		targetIndex = 0
+	}
+	if targetIndex >= len(mappings) {
+		targetIndex = len(mappings) - 1
+	}
+
+	currentMapping := mappings[targetIndex]
+	if currentMapping.SourceIndex < 0 || currentMapping.SourceIndex >= len(sourceMap.Sources) {
+		return nil
+	}
+
+	var endMapping *MappedPosition
+	if targetIndex+1 < len(mappings) {
+		nextMapping := mappings[targetIndex+1]
+		if nextMapping.SourceIndex == currentMapping.SourceIndex {
+			endMapping = &nextMapping
+		}
+	}
+
+	var sourcePath string
+	source := sourceMap.Sources[currentMapping.SourceIndex]
+	if strings.HasPrefix(source, "/") {
+		sourcePath = tspath.NormalizePath(source)
+	} else {
+		sourcePath = tspath.ResolvePath(mapDir, source)
+	}
+
+	startLine := currentMapping.SourcePosition / 10000
+	startColumn := currentMapping.SourcePosition % 10000
+
+	sourceContent, ok := dsm.readFileWithFallback(sourcePath)
+	if !ok {
+		return &SourceRange{
+			FileName: sourcePath,
+			Start:    0,
+			End:      0,
+		}
+	}
+
+	sourceLineStarts := core.ComputeLineStarts(sourceContent)
+	sourceStartPos := int(sourceLineStarts[startLine]) + startColumn
+
+	var sourceEndPos int
+	if endMapping != nil {
+		endLine := endMapping.SourcePosition / 10000
+		endColumn := endMapping.SourcePosition % 10000
+		sourceEndPos = int(sourceLineStarts[endLine]) + endColumn
+	} else {
+		declNextPos := len(declContent)
+		if targetIndex+1 < len(mappings) {
+			declNextPos = mappings[targetIndex+1].GeneratedPosition
+		}
+		rangeLen := declNextPos - currentMapping.GeneratedPosition
+
+		sourceEndPos = sourceStartPos + rangeLen
+		if sourceEndPos > len(sourceContent) {
+			sourceEndPos = len(sourceContent)
+		}
+		if sourceEndPos <= sourceStartPos {
+			sourceEndPos = sourceStartPos + 1
+		}
+	}
+
+	return &SourceRange{
+		FileName: sourcePath,
+		Start:    sourceStartPos,
+		End:      sourceEndPos,
+	}
+}
+
+func tryGetSourceMappingURL(content string) string {
+	lines := strings.Split(content, "\n")
+	for i := len(lines) - 1; i >= 0; i-- {
+		line := strings.TrimSpace(lines[i])
+		if strings.HasPrefix(line, "//# sourceMappingURL=") {
+			return strings.TrimPrefix(line, "//# sourceMappingURL=")
+		}
+		if line != "" && !strings.HasPrefix(line, "//") {
+			break
+		}
+	}
+	return ""
+}
+
+type SourceMap struct {
+	Sources    []string `json:"sources"`
+	Mappings   string   `json:"mappings"`
+	Names      []string `json:"names"`
+	File       string   `json:"file"`
+	SourceRoot string   `json:"sourceRoot"`
+}
+
+type DocumentPosition struct {
+	FileName string
+	Pos      int
+}
+
+type MappedPosition struct {
+	GeneratedPosition int
+	SourceIndex       int
+	SourcePosition    int
+}
+
+func decodeMappings(mappings string, declLineStarts []core.TextPos) []MappedPosition {
+	var result []MappedPosition
+
+	decoder := sourcemap.DecodeMappings(mappings)
+
+	for mapping, done := decoder.Next(); !done; mapping, done = decoder.Next() {
+		if mapping.IsSourceMapping() {
+			generatedPos := int(declLineStarts[mapping.GeneratedLine]) + mapping.GeneratedCharacter
+
+			result = append(result, MappedPosition{
+				GeneratedPosition: generatedPos,
+				SourceIndex:       int(mapping.SourceIndex),
+				SourcePosition:    mapping.SourceLine*10000 + mapping.SourceCharacter,
+			})
+		}
+	}
+
+	if err := decoder.Error(); err != nil {
+		return nil
+	}
+
+	sort.Slice(result, func(i, j int) bool {
+		return result[i].GeneratedPosition < result[j].GeneratedPosition
+	})
+
+	return result
+}

internal/lsp/server.go

@@ -10,6 +10,7 @@ import (
 	"os/signal"
 	"runtime/debug"
 	"slices"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"syscall"
@@ -710,7 +711,62 @@ func (s *Server) handleSignatureHelp(ctx context.Context, languageService *ls.La
 }
 
 func (s *Server) handleDefinition(ctx context.Context, ls *ls.LanguageService, params *lsproto.DefinitionParams) (lsproto.DefinitionResponse, error) {
-	return ls.ProvideDefinition(ctx, params.TextDocument.Uri, params.Position)
+	rawResponse, err := ls.ProvideDefinition(ctx, params.TextDocument.Uri, params.Position)
+	if err != nil {
+		return rawResponse, err
+	}
+
+	if locations := rawResponse.Locations; locations != nil {
+		mappedLocations := s.mapDefinitionLocationsForProject(*locations, params.TextDocument.Uri)
+		return lsproto.LocationOrLocationsOrDefinitionLinksOrNull{
+			Locations: &mappedLocations,
+		}, nil
+	}
+
+	return rawResponse, nil
+}
+
+func (s *Server) mapDefinitionLocationsForProject(locations []lsproto.Location, requestingFileUri lsproto.DocumentUri) []lsproto.Location {
+	mappedLocations := make([]lsproto.Location, 0, len(locations))
+
+	snapshot, release := s.session.Snapshot()
+	defer release()
+
+	requestingProject := snapshot.GetDefaultProject(requestingFileUri)
+	if requestingProject == nil {
+		return locations
+	}
+
+	program := requestingProject.GetProgram()
+	if program == nil {
+		return locations
+	}
+
+	sourceMapper := ls.NewDefinitionSourceMapper(program)
+
+	for _, location := range locations {
+		// Skip bundled library URIs that have invalid URL encoding
+		// These URIs contain "%3A" which causes panics in FileName()
+		uriString := string(location.Uri)
+		if strings.Contains(uriString, "bundled%3A") {
+			mappedLocations = append(mappedLocations, location)
+			continue
+		}
+
+		fileName := location.Uri.FileName()
+		if !strings.HasSuffix(fileName, ".d.ts") {
+			mappedLocations = append(mappedLocations, location)
+			continue
+		}
+
+		if mappedLocation := sourceMapper.MapSingleLocation(location); mappedLocation != nil {
+			mappedLocations = append(mappedLocations, *mappedLocation)
+		} else {
+			mappedLocations = append(mappedLocations, location)
+		}
+	}
+
+	return mappedLocations
 }
 
 func (s *Server) handleTypeDefinition(ctx context.Context, ls *ls.LanguageService, params *lsproto.TypeDefinitionParams) (lsproto.TypeDefinitionResponse, error) {

internal/project/compilerhost.go

@@ -147,6 +147,15 @@ func (fs *compilerFS) ReadFile(path string) (contents string, ok bool) {
 	return "", false
 }
 
+// ReadFileWithFallback implements vfs.FS.
+// This reads the file from the underlying FS if it's not tracked by the program.
+func (fs *compilerFS) ReadFileWithFallback(path string) (contents string, ok bool) {
+	if fh := fs.source.GetFile(path); fh != nil {
+		return fh.Content(), true
+	}
+	return fs.source.FS().ReadFile(path)
+}
+
 // Realpath implements vfs.FS.
 func (fs *compilerFS) Realpath(path string) string {
 	return fs.source.FS().Realpath(path)

testdata/baselines/reference/fourslash/goToDef/DeclarationMapGoToDefinition.baseline.jsonc

@@ -1,12 +1,12 @@
 // === goToDefinition ===
-// === /indexdef.d.ts ===
+// === /index.ts ===
 
-// export declare class Foo {
+// export class Foo {
 //     member: string;
-//     [|methodName|](propName: SomeType): void;
-//     otherMethod(): {
-//         x: number;
-//         y?: undefined;
+//     [|methodName|](propName: SomeType): void {}
+//     otherMethod() {
+//         if (Math.random() > 0.5) {
+//             return {x: 42};
 // // --- (line: 7) skipped ---