Support find-all-references starting from a reference path or reference types comment (#21007)

Author: Andy

src/services/documentHighlights.ts

@@ -2,19 +2,15 @@
 namespace ts.DocumentHighlights {
     export function getDocumentHighlights(program: Program, cancellationToken: CancellationToken, sourceFile: SourceFile, position: number, sourceFilesToSearch: SourceFile[]): DocumentHighlights[] | undefined {
         const node = getTouchingWord(sourceFile, position, /*includeJsDocComment*/ true);
-        // Note that getTouchingWord indicates failure by returning the sourceFile node.
-        if (node === sourceFile) return undefined;
 
-        Debug.assert(node.parent !== undefined);
-
-        if (isJsxOpeningElement(node.parent) && node.parent.tagName === node || isJsxClosingElement(node.parent)) {
+        if (node.parent && (isJsxOpeningElement(node.parent) && node.parent.tagName === node || isJsxClosingElement(node.parent))) {
             // For a JSX element, just highlight the matching tag, not all references.
             const { openingElement, closingElement } = node.parent.parent;
             const highlightSpans = [openingElement, closingElement].map(({ tagName }) => getHighlightSpanForNode(tagName, sourceFile));
             return [{ fileName: sourceFile.fileName, highlightSpans }];
         }
 
-        return getSemanticDocumentHighlights(node, program, cancellationToken, sourceFilesToSearch) || getSyntacticDocumentHighlights(node, sourceFile);
+        return getSemanticDocumentHighlights(position, node, program, cancellationToken, sourceFilesToSearch) || getSyntacticDocumentHighlights(node, sourceFile);
     }
 
     function getHighlightSpanForNode(node: Node, sourceFile: SourceFile): HighlightSpan {
@@ -25,8 +21,8 @@ namespace ts.DocumentHighlights {
         };
     }
 
-    function getSemanticDocumentHighlights(node: Node, program: Program, cancellationToken: CancellationToken, sourceFilesToSearch: SourceFile[]): DocumentHighlights[] {
-        const referenceEntries = FindAllReferences.getReferenceEntriesForNode(node, program, sourceFilesToSearch, cancellationToken);
+    function getSemanticDocumentHighlights(position: number, node: Node, program: Program, cancellationToken: CancellationToken, sourceFilesToSearch: SourceFile[]): DocumentHighlights[] {
+        const referenceEntries = FindAllReferences.getReferenceEntriesForNode(position, node, program, sourceFilesToSearch, cancellationToken);
         return referenceEntries && convertReferencedSymbols(referenceEntries);
     }
 

src/services/findAllReferences.ts

@@ -52,12 +52,12 @@ namespace ts.FindAllReferences {
     export function getImplementationsAtPosition(program: Program, cancellationToken: CancellationToken, sourceFiles: ReadonlyArray<SourceFile>, sourceFile: SourceFile, position: number): ImplementationLocation[] {
         // A node in a JSDoc comment can't have an implementation anyway.
         const node = getTouchingPropertyName(sourceFile, position, /*includeJsDocComment*/ false);
-        const referenceEntries = getImplementationReferenceEntries(program, cancellationToken, sourceFiles, node);
+        const referenceEntries = getImplementationReferenceEntries(program, cancellationToken, sourceFiles, node, position);
         const checker = program.getTypeChecker();
         return map(referenceEntries, entry => toImplementationLocation(entry, checker));
     }
 
-    function getImplementationReferenceEntries(program: Program, cancellationToken: CancellationToken, sourceFiles: ReadonlyArray<SourceFile>, node: Node): Entry[] | undefined {
+    function getImplementationReferenceEntries(program: Program, cancellationToken: CancellationToken, sourceFiles: ReadonlyArray<SourceFile>, node: Node, position: number): Entry[] | undefined {
         if (node.kind === SyntaxKind.SourceFile) {
             return undefined;
         }
@@ -78,7 +78,7 @@ namespace ts.FindAllReferences {
         }
         else {
             // Perform "Find all References" and retrieve only those that are implementations
-            return getReferenceEntriesForNode(node, program, sourceFiles, cancellationToken, { implementations: true });
+            return getReferenceEntriesForNode(position, node, program, sourceFiles, cancellationToken, { implementations: true });
         }
     }
 
@@ -87,13 +87,13 @@ namespace ts.FindAllReferences {
         return map(x, toReferenceEntry);
     }
 
-    export function getReferenceEntriesForNode(node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}): Entry[] | undefined {
-        return flattenEntries(Core.getReferencedSymbolsForNode(node, program, sourceFiles, cancellationToken, options));
+    export function getReferenceEntriesForNode(position: number, node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}): Entry[] | undefined {
+        return flattenEntries(Core.getReferencedSymbolsForNode(position, node, program, sourceFiles, cancellationToken, options));
     }
 
     function findAllReferencedSymbols(program: Program, cancellationToken: CancellationToken, sourceFiles: ReadonlyArray<SourceFile>, sourceFile: SourceFile, position: number, options?: Options): SymbolAndEntries[] | undefined {
         const node = getTouchingPropertyName(sourceFile, position, /*includeJsDocComment*/ true);
-        return Core.getReferencedSymbolsForNode(node, program, sourceFiles, cancellationToken, options);
+        return Core.getReferencedSymbolsForNode(position, node, program, sourceFiles, cancellationToken, options);
     }
 
     function flattenEntries(referenceSymbols: SymbolAndEntries[]): Entry[] {
@@ -242,9 +242,10 @@ namespace ts.FindAllReferences {
 /* @internal */
 namespace ts.FindAllReferences.Core {
     /** Core find-all-references algorithm. Handles special cases before delegating to `getReferencedSymbolsForSymbol`. */
-    export function getReferencedSymbolsForNode(node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}): SymbolAndEntries[] | undefined {
-        if (node.kind === ts.SyntaxKind.SourceFile) {
-            return undefined;
+    export function getReferencedSymbolsForNode(position: number, node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}): SymbolAndEntries[] | undefined {
+        if (isSourceFile(node)) {
+            const reference = GoToDefinition.getReferenceAtPosition(node, position, program);
+            return reference && getReferencedSymbolsForModule(program, program.getTypeChecker().getMergedSymbol(reference.file.symbol), sourceFiles);
         }
 
         if (!options.implementations) {
@@ -260,11 +261,7 @@ namespace ts.FindAllReferences.Core {
         // Could not find a symbol e.g. unknown identifier
         if (!symbol) {
             // String literal might be a property (and thus have a symbol), so do this here rather than in getReferencedSymbolsSpecial.
-            if (!options.implementations && node.kind === SyntaxKind.StringLiteral) {
-                return getReferencesForStringLiteral(<StringLiteral>node, sourceFiles, cancellationToken);
-            }
-            // Can't have references to something that we have no symbol for.
-            return undefined;
+            return !options.implementations && isStringLiteral(node) ? getReferencesForStringLiteral(node, sourceFiles, cancellationToken) : undefined;
         }
 
         if (symbol.flags & SymbolFlags.Module && isModuleReferenceLocation(node)) {

src/services/goToDefinition.ts

@@ -1,22 +1,9 @@
 /* @internal */
 namespace ts.GoToDefinition {
     export function getDefinitionAtPosition(program: Program, sourceFile: SourceFile, position: number): DefinitionInfo[] {
-        /// Triple slash reference comments
-        const comment = findReferenceInPosition(sourceFile.referencedFiles, position);
-        if (comment) {
-            const referenceFile = tryResolveScriptReference(program, sourceFile, comment);
-            if (referenceFile) {
-                return [getDefinitionInfoForFileReference(comment.fileName, referenceFile.fileName)];
-            }
-            // Might still be on jsdoc, so keep looking.
-        }
-
-        // Type reference directives
-        const typeReferenceDirective = findReferenceInPosition(sourceFile.typeReferenceDirectives, position);
-        if (typeReferenceDirective) {
-            const referenceFile = program.getResolvedTypeReferenceDirectives().get(typeReferenceDirective.fileName);
-            return referenceFile && referenceFile.resolvedFileName &&
-                [getDefinitionInfoForFileReference(typeReferenceDirective.fileName, referenceFile.resolvedFileName)];
+        const reference = getReferenceAtPosition(sourceFile, position, program);
+        if (reference) {
+            return [getDefinitionInfoForFileReference(reference.fileName, reference.file.fileName)];
         }
 
         const node = getTouchingPropertyName(sourceFile, position, /*includeJsDocComment*/ true);
@@ -115,6 +102,23 @@ namespace ts.GoToDefinition {
         return getDefinitionFromSymbol(typeChecker, symbol, node);
     }
 
+    export function getReferenceAtPosition(sourceFile: SourceFile, position: number, program: Program): { fileName: string, file: SourceFile } | undefined {
+        const referencePath = findReferenceInPosition(sourceFile.referencedFiles, position);
+        if (referencePath) {
+            const file = tryResolveScriptReference(program, sourceFile, referencePath);
+            return file && { fileName: referencePath.fileName, file };
+        }
+
+        const typeReferenceDirective = findReferenceInPosition(sourceFile.typeReferenceDirectives, position);
+        if (typeReferenceDirective) {
+            const reference = program.getResolvedTypeReferenceDirectives().get(typeReferenceDirective.fileName);
+            const file = reference && program.getSourceFile(reference.resolvedFileName);
+            return file && { fileName: typeReferenceDirective.fileName, file };
+        }
+
+        return undefined;
+    }
+
     /// Goto type
     export function getTypeDefinitionAtPosition(typeChecker: TypeChecker, sourceFile: SourceFile, position: number): DefinitionInfo[] {
         const node = getTouchingPropertyName(sourceFile, position, /*includeJsDocComment*/ true);
@@ -301,7 +305,7 @@ namespace ts.GoToDefinition {
         return createDefinitionInfo(decl, symbolKind, symbolName, containerName);
     }
 
-    function findReferenceInPosition(refs: ReadonlyArray<FileReference>, pos: number): FileReference {
+    export function findReferenceInPosition(refs: ReadonlyArray<FileReference>, pos: number): FileReference {
         for (const ref of refs) {
             if (ref.pos <= pos && pos <= ref.end) {
                 return ref;

tests/cases/fourslash/findAllRefsForModule.ts

@@ -12,12 +12,12 @@
 ////const a = require("[|../a|]");
 
 // @Filename: /d.ts
-//// /// <reference path="[|./a|]" />
+//// /// <reference path="[|./a.ts|]" />
 
 verify.noErrors();
 
 const ranges = test.ranges();
 const [r0, r1, r2] = ranges;
-verify.referenceGroups([r0, r1], [{ definition: 'module "/a"', ranges: [r0, r2, r1] }]);
-// TODO:GH#15736
-verify.referenceGroups(r2, undefined);
+verify.referenceGroups(ranges, [{ definition: 'module "/a"', ranges: [r0, r2, r1] }]);
+// Testing that it works with documentHighlights too
+verify.rangesAreDocumentHighlights();

tests/cases/fourslash/findAllRefsForModuleGlobal.ts

@@ -12,6 +12,4 @@ verify.noErrors();
 
 const ranges = test.ranges();
 const [r0, r1, r2] = ranges;
-verify.referenceGroups([r1, r2], [{ definition: 'module "/node_modules/foo/index"', ranges: [r0, r1, r2] }]);
-// TODO:GH#15736
-verify.referenceGroups(r0, undefined);
+verify.singleReferenceGroup('module "/node_modules/foo/index"');