Fix completion details auto-imports crashes

src/harness/fourslashImpl.ts

@@ -882,18 +882,28 @@ export class TestState {
 
     public verifyCompletions(options: FourSlashInterface.VerifyCompletionsOptions) {
         if (options.marker === undefined) {
-            this.verifyCompletionsWorker(options);
+            return this.verifyCompletionsWorker(options);
         }
         else {
-            for (const marker of toArray(options.marker)) {
-                this.goToMarker(marker);
-                this.verifyCompletionsWorker({ ...options, marker });
+            if (ts.isArray(options.marker)) {
+                for (const marker of options.marker) {
+                    this.goToMarker(marker);
+                    this.verifyCompletionsWorker({ ...options, marker });
+                }
+                return {
+                    andApplyCodeAction: () => {
+                        this.raiseError(`Cannot apply code action when multiple markers are specified.`);
+                    }
+                };
             }
+            this.goToMarker(options.marker);
+            return this.verifyCompletionsWorker({ ...options, marker: options.marker });
         }
     }
 
-    private verifyCompletionsWorker(options: FourSlashInterface.VerifyCompletionsOptions): void {
-        const actualCompletions = this.getCompletionListAtCaret({ ...options.preferences, triggerCharacter: options.triggerCharacter })!;
+    private verifyCompletionsWorker(options: FourSlashInterface.VerifyCompletionsOptions) {
+        const preferences = options.preferences;
+        const actualCompletions = this.getCompletionListAtCaret({ ...preferences, triggerCharacter: options.triggerCharacter })!;
         if (!actualCompletions) {
             if (ts.hasProperty(options, "exact") && (options.exact === undefined || ts.isArray(options.exact) && !options.exact.length)) {
                 return;
@@ -917,6 +927,7 @@ export class TestState {
         }
 
         const nameToEntries = new Map<string, ts.CompletionEntry[]>();
+        const nameAndSourceToData = new Map<string, ts.CompletionEntryData | false>();
         for (const entry of actualCompletions.entries) {
             const entries = nameToEntries.get(entry.name);
             if (!entries) {
@@ -934,6 +945,15 @@ export class TestState {
                 }
                 entries.push(entry);
             }
+            if (entry.data && entry.source) {
+                const key = `${entry.name}|${entry.source}`;
+                if (nameAndSourceToData.has(key)) {
+                    nameAndSourceToData.set(key, false);
+                }
+                else {
+                    nameAndSourceToData.set(key, entry.data);
+                }
+            }
         }
 
         if (ts.hasProperty(options, "exact")) {
@@ -977,6 +997,26 @@ export class TestState {
                 }
             }
         }
+
+        return {
+            andApplyCodeAction: (options: {
+                name: string,
+                source: string,
+                description: string,
+                newFileContent?: string,
+                newRangeContent?: string,
+            }) => {
+                const { name, source, description, newFileContent, newRangeContent } = options;
+                const data = nameAndSourceToData.get(`${options.name}|${options.source}`);
+                if (data === false) {
+                    this.raiseError(`Multiple completion entries found for '${options.name}' from '${options.source}'. This API cannot be used. Use 'verify.applyCodeActionFromCompletion' instead.`);
+                }
+                if (data === undefined) {
+                    this.raiseError(`No completion entry found for '${options.name}' from '${options.source}'`);
+                }
+                this.applyCodeActionFromCompletion(/*markerName*/ undefined, { name, source, data, description, newFileContent, newRangeContent, preferences });
+            }
+        };
     }
 
     private verifyCompletionEntry(actual: ts.CompletionEntry, expected: FourSlashInterface.ExpectedCompletionEntry) {

src/harness/fourslashInterfaceImpl.ts

@@ -252,9 +252,17 @@ export class Verify extends VerifyNegatable {
     }
 
     public completions(...optionsArray: VerifyCompletionsOptions[]) {
+        if (optionsArray.length === 1) {
+            return this.state.verifyCompletions(optionsArray[0]);
+        }
         for (const options of optionsArray) {
             this.state.verifyCompletions(options);
         }
+        return {
+            andApplyCodeAction: () => {
+                throw new Error("Cannot call andApplyCodeAction on multiple completions requests");
+            },
+        };
     }
 
     public getInlayHints(expected: readonly VerifyInlayHintsOptions[], span: ts.TextSpan, preference?: ts.UserPreferences) {

src/services/codefixes/importFixes.ts

@@ -230,9 +230,9 @@ function createImportAdderWorker(sourceFile: SourceFile, program: Program, useAu
         const symbolName = getNameForExportedSymbol(exportedSymbol, getEmitScriptTarget(compilerOptions));
         const checker = program.getTypeChecker();
         const symbol = checker.getMergedSymbol(skipAlias(exportedSymbol, checker));
-        const exportInfo = getAllExportInfoForSymbol(sourceFile, symbol, symbolName, /*isJsxTagName*/ false, program, host, preferences, cancellationToken);
+        const exportInfo = getAllExportInfoForSymbol(sourceFile, symbol, symbolName, moduleSymbol, /*isJsxTagName*/ false, program, host, preferences, cancellationToken);
         const useRequire = shouldUseRequire(sourceFile, program);
-        const fix = getImportFixForSymbol(sourceFile, Debug.checkDefined(exportInfo), moduleSymbol, program, /*position*/ undefined, !!isValidTypeOnlyUseSite, useRequire, host, preferences);
+        const fix = getImportFixForSymbol(sourceFile, Debug.checkDefined(exportInfo), program, /*position*/ undefined, !!isValidTypeOnlyUseSite, useRequire, host, preferences);
         if (fix) {
             addImport({ fix, symbolName, errorIdentifierText: undefined });
         }
@@ -490,6 +490,7 @@ interface FixAddToExistingImportInfo {
 export function getImportCompletionAction(
     targetSymbol: Symbol,
     moduleSymbol: Symbol,
+    exportMapKey: string | undefined,
     sourceFile: SourceFile,
     symbolName: string,
     isJsxTagName: boolean,
@@ -501,15 +502,25 @@ export function getImportCompletionAction(
     cancellationToken: CancellationToken,
 ): { readonly moduleSpecifier: string, readonly codeAction: CodeAction } {
     const compilerOptions = program.getCompilerOptions();
+    let exportInfos;
 
-    const exportInfos = pathIsBareSpecifier(stripQuotes(moduleSymbol.name))
-        ? [getSingleExportInfoForSymbol(targetSymbol, moduleSymbol, program, host)]
-        : getAllExportInfoForSymbol(sourceFile, targetSymbol, symbolName, isJsxTagName, program, host, preferences, cancellationToken);
+    if (exportMapKey) {
+        // The new way: `exportMapKey` should be in the `data` of each auto-import completion entry and
+        // sent back when asking for details.
+        exportInfos = getExportInfoMap(sourceFile, host, program, preferences, cancellationToken).get(sourceFile.path, exportMapKey);
+        Debug.assertIsDefined(exportInfos, "Some exportInfo should match the specified exportMapKey");
+    }
+    else {
+        // The old way, kept alive for super old editors that don't give us `data` back.
+        exportInfos = pathIsBareSpecifier(stripQuotes(moduleSymbol.name))
+            ? [getSingleExportInfoForSymbol(targetSymbol, symbolName, moduleSymbol, program, host)]
+            : getAllExportInfoForSymbol(sourceFile, targetSymbol, symbolName, moduleSymbol, isJsxTagName, program, host, preferences, cancellationToken);
+        Debug.assertIsDefined(exportInfos, "Some exportInfo should match the specified symbol / moduleSymbol");
+    }
 
-    Debug.assertIsDefined(exportInfos);
     const useRequire = shouldUseRequire(sourceFile, program);
     const isValidTypeOnlyUseSite = isValidTypeOnlyAliasUseSite(getTokenAtPosition(sourceFile, position));
-    const fix = Debug.checkDefined(getImportFixForSymbol(sourceFile, exportInfos, moduleSymbol, program, position, isValidTypeOnlyUseSite, useRequire, host, preferences));
+    const fix = Debug.checkDefined(getImportFixForSymbol(sourceFile, exportInfos, program, position, isValidTypeOnlyUseSite, useRequire, host, preferences));
     return {
         moduleSpecifier: fix.moduleSpecifier,
         codeAction: codeFixActionToCodeAction(codeActionForFix(
@@ -532,8 +543,7 @@ export function getPromoteTypeOnlyCompletionAction(sourceFile: SourceFile, symbo
     return fix && codeFixActionToCodeAction(codeActionForFix({ host, formatContext, preferences }, sourceFile, symbolName, fix, includeSymbolNameInDescription, compilerOptions, preferences));
 }
 
-function getImportFixForSymbol(sourceFile: SourceFile, exportInfos: readonly SymbolExportInfo[], moduleSymbol: Symbol, program: Program, position: number | undefined, isValidTypeOnlyUseSite: boolean, useRequire: boolean, host: LanguageServiceHost, preferences: UserPreferences) {
-    Debug.assert(exportInfos.some(info => info.moduleSymbol === moduleSymbol || info.symbol.parent === moduleSymbol), "Some exportInfo should match the specified moduleSymbol");
+function getImportFixForSymbol(sourceFile: SourceFile, exportInfos: readonly SymbolExportInfo[], program: Program, position: number | undefined, isValidTypeOnlyUseSite: boolean, useRequire: boolean, host: LanguageServiceHost, preferences: UserPreferences) {
     const packageJsonImportFilter = createPackageJsonImportFilter(sourceFile, preferences, host);
     return getBestFix(getImportFixes(exportInfos, position, isValidTypeOnlyUseSite, useRequire, program, sourceFile, host, preferences).fixes, sourceFile, program, packageJsonImportFilter, host);
 }
@@ -542,17 +552,17 @@ function codeFixActionToCodeAction({ description, changes, commands }: CodeFixAc
     return { description, changes, commands };
 }
 
-function getAllExportInfoForSymbol(importingFile: SourceFile, symbol: Symbol, symbolName: string, preferCapitalized: boolean, program: Program, host: LanguageServiceHost, preferences: UserPreferences, cancellationToken: CancellationToken | undefined): readonly SymbolExportInfo[] | undefined {
+function getAllExportInfoForSymbol(importingFile: SourceFile, symbol: Symbol, symbolName: string, moduleSymbol: Symbol, preferCapitalized: boolean, program: Program, host: LanguageServiceHost, preferences: UserPreferences, cancellationToken: CancellationToken | undefined): readonly SymbolExportInfo[] | undefined {
     const getChecker = createGetChecker(program, host);
     return getExportInfoMap(importingFile, host, program, preferences, cancellationToken)
         .search(importingFile.path, preferCapitalized, name => name === symbolName, info => {
-            if (skipAlias(info[0].symbol, getChecker(info[0].isFromPackageJson)) === symbol) {
+            if (skipAlias(info[0].symbol, getChecker(info[0].isFromPackageJson)) === symbol && info.some(i => i.moduleSymbol === moduleSymbol || i.symbol.parent === moduleSymbol)) {
                 return info;
             }
         });
 }
 
-function getSingleExportInfoForSymbol(symbol: Symbol, moduleSymbol: Symbol, program: Program, host: LanguageServiceHost): SymbolExportInfo {
+function getSingleExportInfoForSymbol(symbol: Symbol, symbolName: string, moduleSymbol: Symbol, program: Program, host: LanguageServiceHost): SymbolExportInfo {
     const compilerOptions = program.getCompilerOptions();
     const mainProgramInfo = getInfoWithChecker(program.getTypeChecker(), /*isFromPackageJson*/ false);
     if (mainProgramInfo) {
@@ -566,7 +576,7 @@ function getSingleExportInfoForSymbol(symbol: Symbol, moduleSymbol: Symbol, prog
         if (defaultInfo && skipAlias(defaultInfo.symbol, checker) === symbol) {
             return { symbol: defaultInfo.symbol, moduleSymbol, moduleFileName: undefined, exportKind: defaultInfo.exportKind, targetFlags: skipAlias(symbol, checker).flags, isFromPackageJson };
         }
-        const named = checker.tryGetMemberInModuleExportsAndProperties(symbol.name, moduleSymbol);
+        const named = checker.tryGetMemberInModuleExportsAndProperties(symbolName, moduleSymbol);
         if (named && skipAlias(named, checker) === symbol) {
             return { symbol: named, moduleSymbol, moduleFileName: undefined, exportKind: ExportKind.Named, targetFlags: skipAlias(symbol, checker).flags, isFromPackageJson };
         }

src/services/completions.ts

@@ -461,6 +461,7 @@ interface SymbolOriginInfoResolvedExport extends SymbolOriginInfo {
     symbolName: string;
     moduleSymbol: Symbol;
     exportName: string;
+    exportMapKey?: string;
     moduleSpecifier: string;
 }
 
@@ -1924,6 +1925,7 @@ function originToCompletionEntryData(origin: SymbolOriginInfoExport | SymbolOrig
     if (originIsResolvedExport(origin)) {
         const resolvedData: CompletionEntryDataResolved = {
             exportName: origin.exportName,
+            exportMapKey: origin.exportMapKey,
             moduleSpecifier: origin.moduleSpecifier,
             ambientModuleName,
             fileName: origin.fileName,
@@ -1948,6 +1950,7 @@ function completionEntryDataToSymbolOriginInfo(data: CompletionEntryData, comple
         const resolvedOrigin: SymbolOriginInfoResolvedExport = {
             kind: SymbolOriginInfoKind.ResolvedExport,
             exportName: data.exportName,
+            exportMapKey: data.exportMapKey,
             moduleSpecifier: data.moduleSpecifier,
             symbolName: completionName,
             fileName: data.fileName,
@@ -2454,6 +2457,7 @@ function getCompletionEntryCodeActionsAndSourceDisplay(
     const { moduleSpecifier, codeAction } = codefix.getImportCompletionAction(
         targetSymbol,
         moduleSymbol,
+        data?.exportMapKey,
         sourceFile,
         name,
         isJsxOpeningTagName,
@@ -3439,8 +3443,8 @@ function getCompletionData(
                         // module resolution modes, getting past this point guarantees that we'll be
                         // able to generate a suitable module specifier, so we can safely show a completion,
                         // even if we defer computing the module specifier.
-                        const firstImportableExportInfo = find(info, isImportableExportInfo);
-                        if (!firstImportableExportInfo) {
+                        info = filter(info, isImportableExportInfo);
+                        if (!info.length) {
                             return;
                         }
 
@@ -3457,9 +3461,9 @@ function getCompletionData(
                         // it should be identical regardless of which one is used. During the subsequent
                         // `CompletionEntryDetails` request, we'll get all the ExportInfos again and pick
                         // the best one based on the module specifier it produces.
-                        let exportInfo = firstImportableExportInfo, moduleSpecifier;
+                        let exportInfo = info[0], moduleSpecifier;
                         if (result !== "skipped") {
-                            ({ exportInfo = firstImportableExportInfo, moduleSpecifier } = result);
+                            ({ exportInfo = info[0], moduleSpecifier } = result);
                         }
 
                         const isDefaultExport = exportInfo.exportKind === ExportKind.Default;

src/services/types.ts

@@ -1337,6 +1337,7 @@ export interface CompletionEntryDataAutoImport {
      * in the case of InternalSymbolName.ExportEquals and InternalSymbolName.Default.
      */
     exportName: string;
+    exportMapKey?: string;
     moduleSpecifier?: string;
     /** The file name declaring the export's module symbol, if it was an external module */
     fileName?: string;
@@ -1347,7 +1348,6 @@ export interface CompletionEntryDataAutoImport {
 }
 
 export interface CompletionEntryDataUnresolved extends CompletionEntryDataAutoImport {
-    /** The key in the `ExportMapCache` where the completion entry's `SymbolExportInfo[]` is found */
     exportMapKey: string;
 }
 

tests/baselines/reference/api/tsserverlibrary.d.ts

@@ -10590,6 +10590,7 @@ declare namespace ts {
          * in the case of InternalSymbolName.ExportEquals and InternalSymbolName.Default.
          */
         exportName: string;
+        exportMapKey?: string;
         moduleSpecifier?: string;
         /** The file name declaring the export's module symbol, if it was an external module */
         fileName?: string;
@@ -10599,7 +10600,6 @@ declare namespace ts {
         isPackageJsonImport?: true;
     }
     interface CompletionEntryDataUnresolved extends CompletionEntryDataAutoImport {
-        /** The key in the `ExportMapCache` where the completion entry's `SymbolExportInfo[]` is found */
         exportMapKey: string;
     }
     interface CompletionEntryDataResolved extends CompletionEntryDataAutoImport {

tests/baselines/reference/api/typescript.d.ts

@@ -6688,6 +6688,7 @@ declare namespace ts {
          * in the case of InternalSymbolName.ExportEquals and InternalSymbolName.Default.
          */
         exportName: string;
+        exportMapKey?: string;
         moduleSpecifier?: string;
         /** The file name declaring the export's module symbol, if it was an external module */
         fileName?: string;
@@ -6697,7 +6698,6 @@ declare namespace ts {
         isPackageJsonImport?: true;
     }
     interface CompletionEntryDataUnresolved extends CompletionEntryDataAutoImport {
-        /** The key in the `ExportMapCache` where the completion entry's `SymbolExportInfo[]` is found */
         exportMapKey: string;
     }
     interface CompletionEntryDataResolved extends CompletionEntryDataAutoImport {

tests/baselines/reference/importStatementCompletions3.baseline

@@ -49,6 +49,7 @@
           "isImportStatementCompletion": true,
           "data": {
             "exportName": "foo",
+            "exportMapKey": "foo|19|",
             "moduleSpecifier": "./$foo",
             "fileName": "/tests/cases/fourslash/$foo.ts"
           },

tests/baselines/reference/tsserver/moduleSpecifierCache/caches-module-specifiers-within-a-file.js

@@ -961,6 +961,7 @@ Info 67   [00:02:07.000] response:
             "isImportStatementCompletion": true,
             "data": {
               "exportName": "foo",
+              "exportMapKey": "foo|*|",
               "moduleSpecifier": "./a",
               "fileName": "/src/a.ts"
             }
@@ -993,6 +994,7 @@ Info 67   [00:02:07.000] response:
             "isImportStatementCompletion": true,
             "data": {
               "exportName": "observable",
+              "exportMapKey": "observable|*|",
               "moduleSpecifier": "mobx",
               "fileName": "/node_modules/mobx/index.d.ts",
               "isPackageJsonImport": true

tests/baselines/reference/tsserver/moduleSpecifierCache/invalidates-module-specifiers-when-changes-happen-in-contained-node_modules-directories.js

@@ -961,6 +961,7 @@ Info 67   [00:02:07.000] response:
             "isImportStatementCompletion": true,
             "data": {
               "exportName": "foo",
+              "exportMapKey": "foo|*|",
               "moduleSpecifier": "./a",
               "fileName": "/src/a.ts"
             }
@@ -993,6 +994,7 @@ Info 67   [00:02:07.000] response:
             "isImportStatementCompletion": true,
             "data": {
               "exportName": "observable",
+              "exportMapKey": "observable|*|",
               "moduleSpecifier": "mobx",
               "fileName": "/node_modules/mobx/index.d.ts",
               "isPackageJsonImport": true

tests/cases/fourslash/completionsImport_named_didNotExistBefore.ts

@@ -31,9 +31,7 @@ verify.completions({
         },
     ], { noLib: true }),
     preferences: { includeCompletionsForModuleExports: true },
-});
-
-verify.applyCodeActionFromCompletion("", {
+}).andApplyCodeAction({
     name: "Test1",
     source: "/a",
     description: `Update import from "./a"`,

tests/cases/fourslash/fourslash.ts

@@ -284,7 +284,13 @@ declare namespace FourSlashInterface {
     class verify extends verifyNegatable {
         assertHasRanges(ranges: Range[]): void;
         caretAtMarker(markerName?: string): void;
-        completions(...options: CompletionsOptions[]): void;
+        completions(...options: CompletionsOptions[]): { andApplyCodeAction(options: {
+            name: string,
+            source: string,
+            description: string,
+            newFileContent?: string,
+            newRangeContent?: string,
+        }): void };
         applyCodeActionFromCompletion(markerName: string | undefined, options: {
             name: string,
             source?: string,