Cherry-pick PR #49442 into release-4.7 (#49447)

TypeScript Bot · andrewbranch · web-flow · commit 4d2983a7362f · 2022-06-09T11:05:04.000-07:00
Component commits: d8f251b Avoid repeating codefix work when resolving auto-import specifiers for completions Co-authored-by: Andrew Branch <andrew@wheream.io>
diff --git a/src/services/codefixes/importFixes.ts b/src/services/codefixes/importFixes.ts
@@ -40,6 +40,9 @@ namespace ts.codefix {
         },
     });
 
+    /**
+     * Computes multiple import additions to a file and writes them to a ChangeTracker.
+     */
     export interface ImportAdder {
         hasFixes(): boolean;
         addImportFromDiagnostic: (diagnostic: DiagnosticWithLocation, context: CodeFixContextBase) => void;
@@ -235,6 +238,47 @@ namespace ts.codefix {
         }
     }
 
+    /**
+     * Computes module specifiers for multiple import additions to a file.
+     */
+     export interface ImportSpecifierResolver {
+        getModuleSpecifierForBestExportInfo(
+            exportInfo: readonly SymbolExportInfo[],
+            symbolName: string,
+            position: number,
+            isValidTypeOnlyUseSite: boolean,
+            fromCacheOnly?: boolean
+        ): { exportInfo?: SymbolExportInfo, moduleSpecifier: string, computedWithoutCacheCount: number } | undefined;
+    }
+
+    export function createImportSpecifierResolver(importingFile: SourceFile, program: Program, host: LanguageServiceHost, preferences: UserPreferences): ImportSpecifierResolver {
+        const packageJsonImportFilter = createPackageJsonImportFilter(importingFile, preferences, host);
+        const importMap = createExistingImportMap(program.getTypeChecker(), importingFile, program.getCompilerOptions());
+        return { getModuleSpecifierForBestExportInfo };
+
+        function getModuleSpecifierForBestExportInfo(
+            exportInfo: readonly SymbolExportInfo[],
+            symbolName: string,
+            position: number,
+            isValidTypeOnlyUseSite: boolean,
+            fromCacheOnly?: boolean,
+        ): { exportInfo?: SymbolExportInfo, moduleSpecifier: string, computedWithoutCacheCount: number } | undefined {
+            const { fixes, computedWithoutCacheCount } = getImportFixes(
+                exportInfo,
+                { symbolName, position },
+                isValidTypeOnlyUseSite,
+                /*useRequire*/ false,
+                program,
+                importingFile,
+                host,
+                preferences,
+                importMap,
+                fromCacheOnly);
+            const result = getBestFix(fixes, importingFile, program, packageJsonImportFilter, host);
+            return result && { ...result, computedWithoutCacheCount };
+        }
+    }
+
     // Sorted with the preferred fix coming first.
     const enum ImportFixKind { UseNamespace, JsdocTypeImport, AddToExisting, AddNew, PromoteTypeOnly }
     // These should not be combined as bitflags, but are given powers of 2 values to
@@ -394,32 +438,6 @@ namespace ts.codefix {
         }
     }
 
-    export function getModuleSpecifierForBestExportInfo(
-        exportInfo: readonly SymbolExportInfo[],
-        symbolName: string,
-        position: number,
-        isValidTypeOnlyUseSite: boolean,
-        importingFile: SourceFile,
-        program: Program,
-        host: LanguageServiceHost,
-        preferences: UserPreferences,
-        packageJsonImportFilter?: PackageJsonImportFilter,
-        fromCacheOnly?: boolean,
-    ): { exportInfo?: SymbolExportInfo, moduleSpecifier: string, computedWithoutCacheCount: number } | undefined {
-        const { fixes, computedWithoutCacheCount } = getImportFixes(
-            exportInfo,
-            { symbolName, position },
-            isValidTypeOnlyUseSite,
-            /*useRequire*/ false,
-            program,
-            importingFile,
-            host,
-            preferences,
-            fromCacheOnly);
-        const result = getBestFix(fixes, importingFile, program, packageJsonImportFilter || createPackageJsonImportFilter(importingFile, preferences, host), host);
-        return result && { ...result, computedWithoutCacheCount };
-    }
-
     function getImportFixes(
         exportInfos: readonly SymbolExportInfo[],
         useNamespaceInfo: {
@@ -433,10 +451,11 @@ namespace ts.codefix {
         sourceFile: SourceFile,
         host: LanguageServiceHost,
         preferences: UserPreferences,
+        importMap = createExistingImportMap(program.getTypeChecker(), sourceFile, program.getCompilerOptions()),
         fromCacheOnly?: boolean,
     ): { computedWithoutCacheCount: number, fixes: readonly ImportFixWithModuleSpecifier[] } {
         const checker = program.getTypeChecker();
-        const existingImports = flatMap(exportInfos, info => getExistingImportDeclarations(info, checker, sourceFile, program.getCompilerOptions()));
+        const existingImports = flatMap(exportInfos, importMap.getImportsForExportInfo);
         const useNamespace = useNamespaceInfo && tryUseExistingNamespaceImport(existingImports, useNamespaceInfo.symbolName, useNamespaceInfo.position, checker);
         const addToExisting = tryAddToExistingImport(existingImports, isValidTypeOnlyUseSite, checker, program.getCompilerOptions());
         if (addToExisting) {
@@ -587,19 +606,34 @@ namespace ts.codefix {
         });
     }
 
-    function getExistingImportDeclarations({ moduleSymbol, exportKind, targetFlags, symbol }: SymbolExportInfo, checker: TypeChecker, importingFile: SourceFile, compilerOptions: CompilerOptions): readonly FixAddToExistingImportInfo[] {
-        // Can't use an es6 import for a type in JS.
-        if (!(targetFlags & SymbolFlags.Value) && isSourceFileJS(importingFile)) return emptyArray;
-        const importKind = getImportKind(importingFile, exportKind, compilerOptions);
-        return mapDefined(importingFile.imports, (moduleSpecifier): FixAddToExistingImportInfo | undefined => {
+    function createExistingImportMap(checker: TypeChecker, importingFile: SourceFile, compilerOptions: CompilerOptions) {
+        let importMap: MultiMap<SymbolId, AnyImportOrRequire> | undefined;
+        for (const moduleSpecifier of importingFile.imports) {
             const i = importFromModuleSpecifier(moduleSpecifier);
             if (isVariableDeclarationInitializedToRequire(i.parent)) {
-                return checker.resolveExternalModuleName(moduleSpecifier) === moduleSymbol ? { declaration: i.parent, importKind, symbol, targetFlags } : undefined;
+                const moduleSymbol = checker.resolveExternalModuleName(moduleSpecifier);
+                if (moduleSymbol) {
+                    (importMap ||= createMultiMap()).add(getSymbolId(moduleSymbol), i.parent);
+                }
             }
-            if (i.kind === SyntaxKind.ImportDeclaration || i.kind === SyntaxKind.ImportEqualsDeclaration) {
-                return checker.getSymbolAtLocation(moduleSpecifier) === moduleSymbol ? { declaration: i, importKind, symbol, targetFlags } : undefined;
+            else if (i.kind === SyntaxKind.ImportDeclaration || i.kind === SyntaxKind.ImportEqualsDeclaration) {
+                const moduleSymbol = checker.getSymbolAtLocation(moduleSpecifier);
+                if (moduleSymbol) {
+                    (importMap ||= createMultiMap()).add(getSymbolId(moduleSymbol), i);
+                }
             }
-        });
+        }
+
+        return {
+            getImportsForExportInfo: ({ moduleSymbol, exportKind, targetFlags, symbol }: SymbolExportInfo): readonly FixAddToExistingImportInfo[] => {
+                // Can't use an es6 import for a type in JS.
+                if (!(targetFlags & SymbolFlags.Value) && isSourceFileJS(importingFile)) return emptyArray;
+                const matchingDeclarations = importMap?.get(getSymbolId(moduleSymbol));
+                if (!matchingDeclarations) return emptyArray;
+                const importKind = getImportKind(importingFile, exportKind, compilerOptions);
+                return matchingDeclarations.map(declaration => ({ declaration, importKind, symbol, targetFlags }));
+            }
+        };
     }
 
     function shouldUseRequire(sourceFile: SourceFile, program: Program): boolean {
diff --git a/src/services/completions.ts b/src/services/completions.ts
@@ -184,16 +184,15 @@ namespace ts.Completions {
     function resolvingModuleSpecifiers<TReturn>(
         logPrefix: string,
         host: LanguageServiceHost,
+        resolver: codefix.ImportSpecifierResolver,
         program: Program,
-        sourceFile: SourceFile,
         position: number,
         preferences: UserPreferences,
         isForImportStatementCompletion: boolean,
         isValidTypeOnlyUseSite: boolean,
         cb: (context: ModuleSpecifierResolutioContext) => TReturn,
     ): TReturn {
         const start = timestamp();
-        const packageJsonImportFilter = createPackageJsonImportFilter(sourceFile, preferences, host);
         // Under `--moduleResolution nodenext`, we have to resolve module specifiers up front, because
         // package.json exports can mean we *can't* resolve a module specifier (that doesn't include a
         // relative path into node_modules), and we want to filter those completions out entirely.
@@ -221,7 +220,7 @@ namespace ts.Completions {
 
         function tryResolve(exportInfo: readonly SymbolExportInfo[], symbolName: string, isFromAmbientModule: boolean): ModuleSpecifierResolutionResult {
             if (isFromAmbientModule) {
-                const result = codefix.getModuleSpecifierForBestExportInfo(exportInfo, symbolName, position, isValidTypeOnlyUseSite, sourceFile, program, host, preferences);
+                const result = resolver.getModuleSpecifierForBestExportInfo(exportInfo, symbolName, position, isValidTypeOnlyUseSite);
                 if (result) {
                     ambientCount++;
                 }
@@ -230,7 +229,7 @@ namespace ts.Completions {
             const shouldResolveModuleSpecifier = needsFullResolution || preferences.allowIncompleteCompletions && resolvedCount < moduleSpecifierResolutionLimit;
             const shouldGetModuleSpecifierFromCache = !shouldResolveModuleSpecifier && preferences.allowIncompleteCompletions && cacheAttemptCount < moduleSpecifierResolutionCacheAttemptLimit;
             const result = (shouldResolveModuleSpecifier || shouldGetModuleSpecifierFromCache)
-                ? codefix.getModuleSpecifierForBestExportInfo(exportInfo, symbolName, position, isValidTypeOnlyUseSite, sourceFile, program, host, preferences, packageJsonImportFilter, shouldGetModuleSpecifierFromCache)
+                ? resolver.getModuleSpecifierForBestExportInfo(exportInfo, symbolName, position, isValidTypeOnlyUseSite, shouldGetModuleSpecifierFromCache)
                 : undefined;
 
             if (!shouldResolveModuleSpecifier && !shouldGetModuleSpecifierFromCache || shouldGetModuleSpecifierFromCache && !result) {
@@ -371,8 +370,8 @@ namespace ts.Completions {
         const newEntries = resolvingModuleSpecifiers(
             "continuePreviousIncompleteResponse",
             host,
+            codefix.createImportSpecifierResolver(file, program, host, preferences),
             program,
-            file,
             location.getStart(),
             preferences,
             /*isForImportStatementCompletion*/ false,
@@ -2206,6 +2205,7 @@ namespace ts.Completions {
         let hasUnresolvedAutoImports = false;
         // This also gets mutated in nested-functions after the return
         let symbols: Symbol[] = [];
+        let importSpecifierResolver: codefix.ImportSpecifierResolver | undefined;
         const symbolToOriginInfoMap: SymbolOriginInfoMap = [];
         const symbolToSortTextMap: SymbolSortTextMap = [];
         const seenPropertySymbols = new Map<SymbolId, true>();
@@ -2448,14 +2448,14 @@ namespace ts.Completions {
                     }
                     else {
                         const fileName = isExternalModuleNameRelative(stripQuotes(moduleSymbol.name)) ? getSourceFileOfModule(moduleSymbol)?.fileName : undefined;
-                        const { moduleSpecifier } = codefix.getModuleSpecifierForBestExportInfo([{
+                        const { moduleSpecifier } = (importSpecifierResolver ||= codefix.createImportSpecifierResolver(sourceFile, program, host, preferences)).getModuleSpecifierForBestExportInfo([{
                             exportKind: ExportKind.Named,
                             moduleFileName: fileName,
                             isFromPackageJson: false,
                             moduleSymbol,
                             symbol: firstAccessibleSymbol,
                             targetFlags: skipAlias(firstAccessibleSymbol, typeChecker).flags,
-                        }], firstAccessibleSymbol.name, position, isValidTypeOnlyAliasUseSite(location), sourceFile, program, host, preferences) || {};
+                        }], firstAccessibleSymbol.name, position, isValidTypeOnlyAliasUseSite(location)) || {};
 
                         if (moduleSpecifier) {
                             const origin: SymbolOriginInfoResolvedExport = {
@@ -2733,8 +2733,8 @@ namespace ts.Completions {
             resolvingModuleSpecifiers(
                 "collectAutoImports",
                 host,
+                importSpecifierResolver ||= codefix.createImportSpecifierResolver(sourceFile, program, host, preferences),
                 program,
-                sourceFile,
                 position,
                 preferences,
                 !!importCompletionNode,
