(feat) declarationMap support (#1878)

#1732

This also adds an optional getOriginalFilePosition to DocumentMapper for sourcemapping to another file. A new DtsDocumentSnapshot for dts files implements getOriginalFilePosition to sourcemap with declarationMap. Currently, "find references", "get type definition", "go to definition", and "find implementations" are updated to account for the declaration map.

packages/language-server/src/lib/documents/DocumentMapper.ts

@@ -11,19 +11,31 @@ import {
     CodeAction,
     SelectionRange,
     TextEdit,
-    InsertReplaceEdit
+    InsertReplaceEdit,
+    Location
 } from 'vscode-languageserver';
 import { TagInformation, offsetAt, positionAt, getLineOffsets } from './utils';
 import { Logger } from '../../logger';
 import { generatedPositionFor, originalPositionFor, TraceMap } from '@jridgewell/trace-mapping';
 
+export interface FilePosition extends Position {
+    uri?: string;
+}
+
 export interface DocumentMapper {
     /**
      * Map the generated position to the original position
      * @param generatedPosition Position in fragment
      */
     getOriginalPosition(generatedPosition: Position): Position;
 
+    /**
+     * Map the generated position to the original position.
+     * Differs from getOriginalPosition this might maps to different file
+     * @param generatedPosition Position in fragment
+     */
+    getOriginalFilePosition?(generatedPosition: Position): FilePosition;
+
     /**
      * Map the original position to the generated position
      * @param originalPosition Position in parent
@@ -207,7 +219,13 @@ export function mapRangeToOriginal(
         end: fragment.getOriginalPosition(range.end)
     };
 
-    // Range may be mapped one character short - reverse that for "in the same line" cases
+    checkRangeLength(originalRange, range);
+
+    return originalRange;
+}
+
+/** Range may be mapped one character short - reverse that for "in the same line" cases*/
+function checkRangeLength(originalRange: { start: Position; end: Position }, range: Range) {
     if (
         originalRange.start.line === originalRange.end.line &&
         range.start.line === range.end.line &&
@@ -216,8 +234,31 @@ export function mapRangeToOriginal(
     ) {
         originalRange.end.character += 1;
     }
+}
 
-    return originalRange;
+export function mapLocationToOriginal(
+    fragment: Pick<DocumentMapper, 'getOriginalPosition' | 'getURL' | 'getOriginalFilePosition'>,
+    range: Range
+): Location {
+    const map = (
+        fragment.getOriginalFilePosition ??
+        (fragment.getOriginalPosition as (position: Position) => FilePosition)
+    ).bind(fragment);
+
+    const start = map(range.start);
+    const end = map(range.end);
+
+    const originalRange: Range = {
+        start: { line: start.line, character: start.character },
+        end: { line: end.line, character: end.character }
+    };
+
+    checkRangeLength(originalRange, range);
+
+    return {
+        range: originalRange,
+        uri: start.uri ? start.uri : fragment.getURL()
+    };
 }
 
 export function mapRangeToGenerated(fragment: DocumentMapper, range: Range): Range {

packages/language-server/src/plugins/typescript/DocumentSnapshot.ts

@@ -1,4 +1,4 @@
-import { EncodedSourceMap, TraceMap } from '@jridgewell/trace-mapping';
+import { EncodedSourceMap, TraceMap, originalPositionFor } from '@jridgewell/trace-mapping';
 import { walk } from 'svelte/compiler';
 import { TemplateNode } from 'svelte/types/compiler/interfaces';
 import { svelte2tsx, IExportedNames } from 'svelte2tsx';
@@ -13,9 +13,10 @@ import {
     positionAt,
     TagInformation,
     isInTag,
-    getLineOffsets
+    getLineOffsets,
+    FilePosition
 } from '../../lib/documents';
-import { pathToUrl } from '../../utils';
+import { pathToUrl, urlToPath } from '../../utils';
 import { ConsumerDocumentMapper } from './DocumentMapper';
 import { SvelteNode } from './svelte-ast-utils';
 import {
@@ -24,6 +25,9 @@ import {
     isSvelteFilePath,
     getTsCheckComment
 } from './utils';
+import { Logger } from '../../logger';
+import { dirname, join, resolve } from 'path';
+import { URI } from 'vscode-uri';
 
 /**
  * An error which occurred while trying to parse/preprocess the svelte file contents.
@@ -138,6 +142,11 @@ export namespace DocumentSnapshot {
             }
         }
 
+        const declarationExtensions = [ts.Extension.Dcts, ts.Extension.Dts, ts.Extension.Dmts];
+        if (declarationExtensions.some((ext) => normalizedPath.endsWith(ext))) {
+            return new DtsDocumentSnapshot(INITIAL_VERSION, filePath, originalText, tsSystem);
+        }
+
         return new JSOrTSDocumentSnapshot(INITIAL_VERSION, filePath, originalText);
     }
 
@@ -440,10 +449,168 @@ export class JSOrTSDocumentSnapshot extends IdentityMapper implements DocumentSn
         this.lineOffsets = undefined;
     }
 
-    private getLineOffsets() {
+    protected getLineOffsets() {
         if (!this.lineOffsets) {
             this.lineOffsets = getLineOffsets(this.text);
         }
         return this.lineOffsets;
     }
 }
+
+const sourceMapCommentRegExp = /^\/\/[@#] source[M]appingURL=(.+)\r?\n?$/;
+const whitespaceOrMapCommentRegExp = /^\s*(\/\/[@#] .*)?$/;
+const base64UrlRegExp =
+    /^data:(?:application\/json(?:;charset=[uU][tT][fF]-8);base64,([A-Za-z0-9+\/=]+)$)?/;
+
+export class DtsDocumentSnapshot extends JSOrTSDocumentSnapshot implements DocumentMapper {
+    private traceMap: TraceMap | undefined;
+    private mapperInitialized = false;
+
+    constructor(version: number, filePath: string, text: string, private tsSys: ts.System) {
+        super(version, filePath, text);
+    }
+
+    getOriginalFilePosition(generatedPosition: Position): FilePosition {
+        if (!this.mapperInitialized) {
+            this.traceMap = this.initMapper();
+            this.mapperInitialized = true;
+        }
+
+        const mapped = this.traceMap
+            ? originalPositionFor(this.traceMap, {
+                  line: generatedPosition.line + 1,
+                  column: generatedPosition.character
+              })
+            : undefined;
+
+        if (!mapped || mapped.line == null || !mapped.source) {
+            return generatedPosition;
+        }
+
+        const originalFilePath = URI.isUri(mapped.source)
+            ? urlToPath(mapped.source)
+            : this.filePath
+            ? resolve(dirname(this.filePath), mapped.source).toString()
+            : undefined;
+
+        // ex: library publish with declarationMap but npmignore the original files
+        if (!originalFilePath || !this.tsSys.fileExists(originalFilePath)) {
+            return generatedPosition;
+        }
+
+        return {
+            line: mapped.line,
+            character: mapped.column,
+            uri: pathToUrl(originalFilePath)
+        };
+    }
+
+    private initMapper() {
+        const sourceMapUrl = tryGetSourceMappingURL(this.getLineOffsets(), this.getFullText());
+
+        if (!sourceMapUrl) {
+            return;
+        }
+
+        const match = sourceMapUrl.match(base64UrlRegExp);
+        if (match) {
+            const base64Json = match[1];
+            if (!base64Json || !this.tsSys.base64decode) {
+                return;
+            }
+
+            return this.initMapperByRawSourceMap(this.tsSys.base64decode(base64Json));
+        }
+
+        const tryingLocations = new Set([
+            resolve(dirname(this.filePath), sourceMapUrl),
+            this.filePath + '.map'
+        ]);
+
+        for (const mapFilePath of tryingLocations) {
+            if (!this.tsSys.fileExists(mapFilePath)) {
+                continue;
+            }
+
+            const mapFileContent = this.tsSys.readFile(mapFilePath);
+            if (mapFileContent) {
+                return this.initMapperByRawSourceMap(mapFileContent);
+            }
+        }
+
+        this.logFailedToResolveSourceMap("can't find valid sourcemap file");
+    }
+
+    private initMapperByRawSourceMap(input: string) {
+        const map = tryParseRawSourceMap(input);
+
+        // don't support inline sourcemap because
+        // it must be a file that editor can point to
+        if (
+            !map ||
+            !map.mappings ||
+            map.sourcesContent?.some((content) => typeof content === 'string')
+        ) {
+            this.logFailedToResolveSourceMap('invalid or unsupported sourcemap');
+            return;
+        }
+
+        return new TraceMap(map);
+    }
+
+    private logFailedToResolveSourceMap(...errors: any[]) {
+        Logger.debug(`Resolving declaration map for ${this.filePath} failed. `, ...errors);
+    }
+}
+
+// https://github.com/microsoft/TypeScript/blob/1dc5b28b94b4a63f735a42d6497d538434d69b66/src/compiler/sourcemap.ts#L381
+function tryGetSourceMappingURL(lineOffsets: number[], text: string) {
+    for (let index = lineOffsets.length - 1; index >= 0; index--) {
+        const line = text.slice(lineOffsets[index], lineOffsets[index + 1]);
+        const comment = sourceMapCommentRegExp.exec(line);
+        if (comment) {
+            return comment[1].trimEnd();
+        }
+        // If we see a non-whitespace/map comment-like line, break, to avoid scanning up the entire file
+        else if (!line.match(whitespaceOrMapCommentRegExp)) {
+            break;
+        }
+    }
+}
+
+// https://github.com/microsoft/TypeScript/blob/1dc5b28b94b4a63f735a42d6497d538434d69b66/src/compiler/sourcemap.ts#L402
+
+function isRawSourceMap(x: any): x is EncodedSourceMap {
+    return (
+        x !== null &&
+        typeof x === 'object' &&
+        x.version === 3 &&
+        typeof x.file === 'string' &&
+        typeof x.mappings === 'string' &&
+        Array.isArray(x.sources) &&
+        x.sources.every((source: any) => typeof source === 'string') &&
+        (x.sourceRoot === undefined || x.sourceRoot === null || typeof x.sourceRoot === 'string') &&
+        (x.sourcesContent === undefined ||
+            x.sourcesContent === null ||
+            (Array.isArray(x.sourcesContent) &&
+                x.sourcesContent.every(
+                    (content: any) => typeof content === 'string' || content === null
+                ))) &&
+        (x.names === undefined ||
+            x.names === null ||
+            (Array.isArray(x.names) && x.names.every((name: any) => typeof name === 'string')))
+    );
+}
+
+function tryParseRawSourceMap(text: string) {
+    try {
+        const parsed = JSON.parse(text);
+        if (isRawSourceMap(parsed)) {
+            return parsed;
+        }
+    } catch {
+        // empty
+    }
+
+    return undefined;
+}

packages/language-server/src/plugins/typescript/TypeScriptPlugin.ts

@@ -80,6 +80,7 @@ import { LSAndTSDocResolver } from './LSAndTSDocResolver';
 import { ignoredBuildDirectories } from './SnapshotManager';
 import { isAttributeName, isAttributeShorthand, isEventHandler } from './svelte-ast-utils';
 import {
+    convertToLocationForReferenceOrDefinition,
     convertToLocationRange,
     getScriptKindFromFileName,
     isInScript,
@@ -378,10 +379,14 @@ export class TypeScriptPlugin
                     snapshot = await snapshots.retrieve(def.fileName);
                 }
 
+                const defLocation = convertToLocationForReferenceOrDefinition(
+                    snapshot,
+                    def.textSpan
+                );
                 return LocationLink.create(
-                    pathToUrl(def.fileName),
-                    convertToLocationRange(snapshot, def.textSpan),
-                    convertToLocationRange(snapshot, def.textSpan),
+                    defLocation.uri,
+                    defLocation.range,
+                    defLocation.range,
                     convertToLocationRange(tsDoc, defs.textSpan)
                 );
             })

packages/language-server/src/plugins/typescript/features/FindFileReferencesProvider.ts

@@ -1,10 +1,10 @@
 import { Location } from 'vscode-languageserver';
 import { URI } from 'vscode-uri';
-import { pathToUrl } from '../../../utils';
 import { FileReferencesProvider } from '../../interfaces';
 import { LSAndTSDocResolver } from '../LSAndTSDocResolver';
 import { convertToLocationRange, hasNonZeroRange } from '../utils';
 import { SnapshotMap } from './utils';
+import { pathToUrl } from '../../../utils';
 
 export class FindFileReferencesProviderImpl implements FileReferencesProvider {
     constructor(private readonly lsAndTsDocResolver: LSAndTSDocResolver) {}

packages/language-server/src/plugins/typescript/features/FindReferencesProvider.ts

@@ -1,11 +1,11 @@
 import type ts from 'typescript';
 import { Location, Position, ReferenceContext } from 'vscode-languageserver';
 import { Document } from '../../../lib/documents';
-import { flatten, isNotNullOrUndefined, pathToUrl } from '../../../utils';
+import { flatten, isNotNullOrUndefined } from '../../../utils';
 import { FindReferencesProvider } from '../../interfaces';
 import { SvelteDocumentSnapshot } from '../DocumentSnapshot';
 import { LSAndTSDocResolver } from '../LSAndTSDocResolver';
-import { convertToLocationRange, hasNonZeroRange } from '../utils';
+import { convertToLocationForReferenceOrDefinition, hasNonZeroRange } from '../utils';
 import {
     get$storeOffsetOf$storeDeclaration,
     getStoreOffsetOf$storeDeclaration,
@@ -125,10 +125,8 @@ export class FindReferencesProviderImpl implements FindReferencesProvider {
             return null;
         }
 
-        const location = Location.create(
-            pathToUrl(ref.fileName),
-            convertToLocationRange(snapshot, ref.textSpan)
-        );
+        // TODO we should deduplicate if we support finding references from multiple language service
+        const location = convertToLocationForReferenceOrDefinition(snapshot, ref.textSpan);
 
         // Some references are in generated code but not wrapped with explicit ignore comments.
         // These show up as zero-length ranges, so filter them out.

packages/language-server/src/plugins/typescript/features/ImplementationProvider.ts

@@ -1,6 +1,6 @@
 import { Position, Location } from 'vscode-languageserver-protocol';
-import { Document, mapRangeToOriginal } from '../../../lib/documents';
-import { pathToUrl, isNotNullOrUndefined } from '../../../utils';
+import { Document, mapLocationToOriginal } from '../../../lib/documents';
+import { isNotNullOrUndefined } from '../../../utils';
 import { ImplementationProvider } from '../../interfaces';
 import { LSAndTSDocResolver } from '../LSAndTSDocResolver';
 import { convertRange } from '../utils';
@@ -47,13 +47,13 @@ export class ImplementationProviderImpl implements ImplementationProvider {
                     snapshot = await snapshots.retrieve(implementation.fileName);
                 }
 
-                const range = mapRangeToOriginal(
+                const location = mapLocationToOriginal(
                     snapshot,
                     convertRange(snapshot, implementation.textSpan)
                 );
 
-                if (range.start.line >= 0 && range.end.line >= 0) {
-                    return Location.create(pathToUrl(implementation.fileName), range);
+                if (location.range.start.line >= 0 && location.range.end.line >= 0) {
+                    return location;
                 }
             })
         );