Handle find all references for symbol merged with UMD module and global var

Fixes #29093

Author: sheetalkamat

src/harness/fourslash.ts

@@ -939,8 +939,8 @@ namespace FourSlash {
             const startFile = this.activeFile.fileName;
             for (const fileName of files) {
                 const searchFileNames = startFile === fileName ? [startFile] : [startFile, fileName];
-                const highlights = this.getDocumentHighlightsAtCurrentPosition(searchFileNames)!;
-                if (!highlights.every(dh => ts.contains(searchFileNames, dh.fileName))) {
+                const highlights = this.getDocumentHighlightsAtCurrentPosition(searchFileNames);
+                if (highlights && !highlights.every(dh => ts.contains(searchFileNames, dh.fileName))) {
                     this.raiseError(`When asking for document highlights only in files ${searchFileNames}, got document highlights in ${unique(highlights, dh => dh.fileName)}`);
                 }
             }

src/services/findAllReferences.ts

@@ -111,7 +111,7 @@ namespace ts.FindAllReferences {
         return flattenEntries(Core.getReferencedSymbolsForNode(position, node, program, sourceFiles, cancellationToken, options, sourceFilesSet));
     }
 
-    function flattenEntries(referenceSymbols: SymbolAndEntries[] | undefined): ReadonlyArray<Entry> | undefined {
+    function flattenEntries(referenceSymbols: ReadonlyArray<SymbolAndEntries> | undefined): ReadonlyArray<Entry> | undefined {
         return referenceSymbols && flatMap(referenceSymbols, r => r.references);
     }
 
@@ -282,6 +282,11 @@ namespace ts.FindAllReferences {
         return createTextSpanFromBounds(start, end);
     }
 
+    export function getTextSpanOfEntry(entry: Entry) {
+        return entry.kind === EntryKind.Span ? entry.textSpan :
+            getTextSpan(entry.node, entry.node.getSourceFile());
+    }
+
     /** A node is considered a writeAccess iff it is a name of a declaration or a target of an assignment */
     function isWriteAccessForReference(node: Node): boolean {
         const decl = getDeclarationFromName(node);
@@ -353,7 +358,7 @@ namespace ts.FindAllReferences {
 /* @internal */
 namespace ts.FindAllReferences.Core {
     /** Core find-all-references algorithm. Handles special cases before delegating to `getReferencedSymbolsForSymbol`. */
-    export function getReferencedSymbolsForNode(position: number, node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}, sourceFilesSet: ReadonlyMap<true> = arrayToSet(sourceFiles, f => f.fileName)): SymbolAndEntries[] | undefined {
+    export function getReferencedSymbolsForNode(position: number, node: Node, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options = {}, sourceFilesSet: ReadonlyMap<true> = arrayToSet(sourceFiles, f => f.fileName)): ReadonlyArray<SymbolAndEntries> | undefined {
         if (isSourceFile(node)) {
             const reference = GoToDefinition.getReferenceAtPosition(node, position, program);
             const moduleSymbol = reference && program.getTypeChecker().getMergedSymbol(reference.file.symbol);
@@ -368,7 +373,7 @@ namespace ts.FindAllReferences.Core {
         }
 
         const checker = program.getTypeChecker();
-        let symbol = checker.getSymbolAtLocation(node);
+        const symbol = checker.getSymbolAtLocation(node);
 
         // Could not find a symbol e.g. unknown identifier
         if (!symbol) {
@@ -380,23 +385,92 @@ namespace ts.FindAllReferences.Core {
             return getReferencedSymbolsForModule(program, symbol.parent!, /*excludeImportTypeOfExportEquals*/ false, sourceFiles, sourceFilesSet);
         }
 
-        let moduleReferences: SymbolAndEntries[] = emptyArray;
-        const moduleSourceFile = isModuleSymbol(symbol);
-        let referencedNode: Node | undefined = node;
-        if (moduleSourceFile) {
-            const exportEquals = symbol.exports!.get(InternalSymbolName.ExportEquals);
-            // If !!exportEquals, we're about to add references to `import("mod")` anyway, so don't double-count them.
-            moduleReferences = getReferencedSymbolsForModule(program, symbol, !!exportEquals, sourceFiles, sourceFilesSet);
-            if (!exportEquals || !sourceFilesSet.has(moduleSourceFile.fileName)) return moduleReferences;
-            // Continue to get references to 'export ='.
-            symbol = skipAlias(exportEquals, checker);
-            referencedNode = undefined;
+        const moduleReferences = getReferencedSymbolsForModuleIfDeclaredBySourceFile(symbol, program, sourceFiles, cancellationToken, options, sourceFilesSet);
+        if (moduleReferences && !(symbol.flags & SymbolFlags.Transient)) {
+            return moduleReferences;
+        }
+
+        const aliasedSymbol = getMergedAliasedSymbolOfNamespaceExportDeclaration(node, symbol, checker);
+        const moduleReferencesOfExportTarget = aliasedSymbol &&
+            getReferencedSymbolsForModuleIfDeclaredBySourceFile(aliasedSymbol, program, sourceFiles, cancellationToken, options, sourceFilesSet);
+
+        const references = getReferencedSymbolsForSymbol(symbol, node, sourceFiles, sourceFilesSet, checker, cancellationToken, options);
+        return mergeReferences(program, moduleReferences, references, moduleReferencesOfExportTarget);
+    }
+
+    function getMergedAliasedSymbolOfNamespaceExportDeclaration(node: Node, symbol: Symbol, checker: TypeChecker) {
+        if (node.parent && isNamespaceExportDeclaration(node.parent)) {
+            const aliasedSymbol = checker.getAliasedSymbol(symbol);
+            const targetSymbol = checker.getMergedSymbol(aliasedSymbol);
+            if (aliasedSymbol !== targetSymbol) {
+                return targetSymbol;
+            }
+        }
+        return undefined;
+    }
+
+    function getReferencedSymbolsForModuleIfDeclaredBySourceFile(symbol: Symbol, program: Program, sourceFiles: ReadonlyArray<SourceFile>, cancellationToken: CancellationToken, options: Options, sourceFilesSet: ReadonlyMap<true>) {
+        const moduleSourceFile = symbol.flags & SymbolFlags.Module ? find(symbol.declarations, isSourceFile) : undefined;
+        if (!moduleSourceFile) return undefined;
+        const exportEquals = symbol.exports!.get(InternalSymbolName.ExportEquals);
+        // If !!exportEquals, we're about to add references to `import("mod")` anyway, so don't double-count them.
+        const moduleReferences = getReferencedSymbolsForModule(program, symbol, !!exportEquals, sourceFiles, sourceFilesSet);
+        if (!exportEquals || !sourceFilesSet.has(moduleSourceFile.fileName)) return moduleReferences;
+        // Continue to get references to 'export ='.
+        const checker = program.getTypeChecker();
+        symbol = skipAlias(exportEquals, checker);
+        return mergeReferences(program, moduleReferences, getReferencedSymbolsForSymbol(symbol, /*node*/ undefined, sourceFiles, sourceFilesSet, checker, cancellationToken, options));
+    }
+
+    function mergeReferences(program: Program, ...referencesToMerge: (SymbolAndEntries[] | undefined)[]): SymbolAndEntries[] | undefined {
+        let result: SymbolAndEntries[] | undefined;
+        for (const references of referencesToMerge) {
+            if (!references || !references.length) continue;
+            if (!result) {
+                result = references;
+                continue;
+            }
+            for (const entry of references) {
+                if (!entry.definition || entry.definition.type !== DefinitionKind.Symbol) {
+                    result.push(entry);
+                    continue;
+                }
+                const symbol = entry.definition.symbol;
+                const refIndex = findIndex(result, ref => !!ref.definition &&
+                    ref.definition.type === DefinitionKind.Symbol &&
+                    ref.definition.symbol === symbol);
+                if (refIndex === -1) {
+                    result.push(entry);
+                    continue;
+                }
+
+                const reference = result[refIndex];
+                result[refIndex] = {
+                    definition: reference.definition,
+                    references: reference.references.concat(entry.references).sort((entry1, entry2) => {
+                        const entry1File = getSourceFileIndexOfEntry(program, entry1);
+                        const entry2File = getSourceFileIndexOfEntry(program, entry2);
+                        if (entry1File !== entry2File) {
+                            return compareValues(entry1File, entry2File);
+                        }
+
+                        const entry1Span = getTextSpanOfEntry(entry1);
+                        const entry2Span = getTextSpanOfEntry(entry2);
+                        return entry1Span.start !== entry2Span.start ?
+                            compareValues(entry1Span.start, entry2Span.start) :
+                            compareValues(entry1Span.length, entry2Span.length);
+                    })
+                };
+            }
         }
-        return concatenate(moduleReferences, getReferencedSymbolsForSymbol(symbol, referencedNode, sourceFiles, sourceFilesSet, checker, cancellationToken, options));
+        return result;
     }
 
-    function isModuleSymbol(symbol: Symbol): SourceFile | undefined {
-        return symbol.flags & SymbolFlags.Module ? find(symbol.declarations, isSourceFile) : undefined;
+    function getSourceFileIndexOfEntry(program: Program, entry: Entry) {
+        const sourceFile = entry.kind === EntryKind.Span ?
+            program.getSourceFile(entry.fileName)! :
+            entry.node.getSourceFile();
+        return program.getSourceFiles().indexOf(sourceFile);
     }
 
     function getReferencedSymbolsForModule(program: Program, symbol: Symbol, excludeImportTypeOfExportEquals: boolean, sourceFiles: ReadonlyArray<SourceFile>, sourceFilesSet: ReadonlyMap<true>): SymbolAndEntries[] {
@@ -435,7 +509,7 @@ namespace ts.FindAllReferences.Core {
                     break;
                 default:
                     // This may be merged with something.
-                    Debug.fail("Expected a module symbol to be declared by a SourceFile or ModuleDeclaration.");
+                    Debug.assert(!!(symbol.flags & SymbolFlags.Transient), "Expected a module symbol to be declared by a SourceFile or ModuleDeclaration.");
             }
         }
 
@@ -551,6 +625,8 @@ namespace ts.FindAllReferences.Core {
         // If the symbol is declared as part of a declaration like `{ type: "a" } | { type: "b" }`, use the property on the union type to get more references.
         return firstDefined(symbol.declarations, decl => {
             if (!decl.parent) {
+                // Ignore UMD module and global merge
+                if (symbol.flags & SymbolFlags.Transient) return undefined;
                 // Assertions for GH#21814. We should be handling SourceFile symbols in `getReferencedSymbolsForModule` instead of getting here.
                 Debug.fail(`Unexpected symbol at ${Debug.showSyntaxKind(node)}: ${Debug.showSymbol(symbol)}`);
             }
@@ -588,6 +664,12 @@ namespace ts.FindAllReferences.Core {
         Class,
     }
 
+    function getNonModuleSymbolOfMergedModuleSymbol(symbol: Symbol) {
+        if (!(symbol.flags & (SymbolFlags.Module | SymbolFlags.Transient))) return undefined;
+        const decl = symbol.declarations && find(symbol.declarations, d => !isSourceFile(d) && !isModuleDeclaration(d));
+        return decl && decl.symbol;
+    }
+
     /**
      * Holds all state needed for the finding references.
      * Unlike `Search`, there is only one `State`.
@@ -648,7 +730,7 @@ namespace ts.FindAllReferences.Core {
             // The other two forms seem to be handled downstream (e.g. in `skipPastExportOrImportSpecifier`), so special-casing the first form
             // here appears to be intentional).
             const {
-                text = stripQuotes(unescapeLeadingUnderscores((getLocalSymbolForExportDefault(symbol) || symbol).escapedName)),
+                text = stripQuotes(unescapeLeadingUnderscores((getLocalSymbolForExportDefault(symbol) || getNonModuleSymbolOfMergedModuleSymbol(symbol) || symbol).escapedName)),
                 allSearchSymbols = [symbol],
             } = searchOptions;
             const escapedText = escapeLeadingUnderscores(text);
@@ -1573,6 +1655,13 @@ namespace ts.FindAllReferences.Core {
             if (res2) return res2;
         }
 
+        const aliasedSymbol = getMergedAliasedSymbolOfNamespaceExportDeclaration(location, symbol, checker);
+        if (aliasedSymbol) {
+            // In case of UMD module and global merging, search for global as well
+            const res = cbSymbol(aliasedSymbol, /*rootSymbol*/ undefined, /*baseSymbol*/ undefined, EntryKind.Node);
+            if (res) return res;
+        }
+
         const res = fromRoot(symbol);
         if (res) return res;
 

tests/cases/fourslash/findAllReferencesUmdModuleAsGlobalConst.ts

@@ -0,0 +1,43 @@
+/// <reference path='fourslash.ts' />
+
+// @Filename: /node_modules/@types/three/three-core.d.ts
+////export class Vector3 {
+////    constructor(x?: number, y?: number, z?: number);
+////    x: number;
+////    y: number;
+////}
+
+// @Filename: /node_modules/@types/three/index.d.ts
+////export * from "./three-core";
+////export as namespace [|{| "isWriteAccess": true, "isDefinition": true |}THREE|];
+
+// @Filename: /typings/global.d.ts
+////import * as _THREE from '[|three|]';
+////declare global {
+////    const [|{| "isWriteAccess": true, "isDefinition": true |}THREE|]: typeof _THREE;
+////}
+
+// @Filename: /src/index.ts
+////export const a = {};
+////let v = new [|THREE|].Vector2();
+
+// @Filename: /tsconfig.json
+////{
+////    "compilerOptions": {
+////        "esModuleInterop": true,
+////        "outDir": "./build/js/",
+////        "noImplicitAny": true,
+////        "module": "es6",
+////        "target": "es6",
+////        "allowJs": true,
+////        "skipLibCheck": true,
+////        "lib": ["es2016", "dom"],
+////        "typeRoots": ["node_modules/@types/"],
+////        "types": ["three"]
+//// 	},
+////    "files": ["/src/index.ts", "typings/global.d.ts"]
+////}
+
+// TODO:: this should be var THREE: typeof import instead of module name as var but thats existing issue and repros with quickInfo too.
+verify.singleReferenceGroup(`module "/node_modules/@types/three/index"
+var "/node_modules/@types/three/index": typeof import("/node_modules/@types/three/index")`);