Auto-import perf: Do symbol name/flags filtering before cache rehydration (#47678)

* Do symbol name filtering before cache rehydration

* Fix typo

Co-authored-by: Nathan Shively-Sanders <[email protected]>

* Update test

* Fix variable clobbered in merge conflict

Co-authored-by: Nathan Shively-Sanders <[email protected]>

Author: andrewbranch

src/services/completions.ts

@@ -2495,19 +2495,26 @@ namespace ts.Completions {
                 preferences,
                 !!importCompletionNode,
                 context => {
-                    exportInfo.forEach(sourceFile.path, (info, getSymbolName, isFromAmbientModule, exportMapKey) => {
-                        const symbolName = getSymbolName(/*preferCapitalized*/ isRightOfOpenTag);
-                        if (!isIdentifierText(symbolName, getEmitScriptTarget(host.getCompilationSettings()))) return;
-                        if (!detailsEntryId && isStringANonContextualKeyword(symbolName)) return;
-                        // `targetFlags` should be the same for each `info`
-                        if (!isTypeOnlyLocation && !importCompletionNode && !(info[0].targetFlags & SymbolFlags.Value)) return;
-                        if (isTypeOnlyLocation && !(info[0].targetFlags & (SymbolFlags.Module | SymbolFlags.Type))) return;
-                        // Do not try to auto-import something with a lowercase first letter for a JSX tag
-                        const firstChar = symbolName.charCodeAt(0);
-                        if (isRightOfOpenTag && (firstChar < CharacterCodes.A || firstChar > CharacterCodes.Z)) return;
-
-                        const isCompletionDetailsMatch = detailsEntryId && some(info, i => detailsEntryId.source === stripQuotes(i.moduleSymbol.name));
-                        if (isCompletionDetailsMatch || !detailsEntryId && charactersFuzzyMatchInString(symbolName, lowerCaseTokenText)) {
+                    exportInfo.search(
+                        sourceFile.path,
+                        /*preferCapitalized*/ isRightOfOpenTag,
+                        (symbolName, targetFlags) => {
+                            if (!isIdentifierText(symbolName, getEmitScriptTarget(host.getCompilationSettings()))) return false;
+                            if (!detailsEntryId && isStringANonContextualKeyword(symbolName)) return false;
+                            if (!isTypeOnlyLocation && !importCompletionNode && !(targetFlags & SymbolFlags.Value)) return false;
+                            if (isTypeOnlyLocation && !(targetFlags & (SymbolFlags.Module | SymbolFlags.Type))) return false;
+                            // Do not try to auto-import something with a lowercase first letter for a JSX tag
+                            const firstChar = symbolName.charCodeAt(0);
+                            if (isRightOfOpenTag && (firstChar < CharacterCodes.A || firstChar > CharacterCodes.Z)) return false;
+
+                            if (detailsEntryId) return true;
+                            return charactersFuzzyMatchInString(symbolName, lowerCaseTokenText);
+                        },
+                        (info, symbolName, isFromAmbientModule, exportMapKey) => {
+                            if (detailsEntryId && !some(info, i => detailsEntryId.source === stripQuotes(i.moduleSymbol.name))) {
+                                return;
+                            }
+
                             const defaultExportInfo = find(info, isImportableExportInfo);
                             if (!defaultExportInfo) {
                                 return;
@@ -2531,7 +2538,7 @@ namespace ts.Completions {
                                 isFromPackageJson: exportInfo.isFromPackageJson,
                             });
                         }
-                    });
+                    );
 
                     hasUnresolvedAutoImports = context.resolutionLimitExceeded();
                 }

src/services/exportInfoMap.ts

@@ -29,6 +29,7 @@ namespace ts {
         // Used to rehydrate `symbol` and `moduleSymbol` when transient
         id: number;
         symbolName: string;
+        capitalizedSymbolName: string | undefined;
         symbolTableKey: __String;
         moduleName: string;
         moduleFile: SourceFile | undefined;
@@ -48,7 +49,7 @@ namespace ts {
         clear(): void;
         add(importingFile: Path, symbol: Symbol, key: __String, moduleSymbol: Symbol, moduleFile: SourceFile | undefined, exportKind: ExportKind, isFromPackageJson: boolean, checker: TypeChecker): void;
         get(importingFile: Path, key: string): readonly SymbolExportInfo[] | undefined;
-        forEach(importingFile: Path, action: (info: readonly SymbolExportInfo[], getSymbolName: (preferCapitalized?: boolean) => string, isFromAmbientModule: boolean, key: string) => void): void;
+        search(importingFile: Path, preferCapitalized: boolean, matches: (name: string, targetFlags: SymbolFlags) => boolean, action: (info: readonly SymbolExportInfo[], symbolName: string, isFromAmbientModule: boolean, key: string) => void): void;
         releaseSymbols(): void;
         isEmpty(): boolean;
         /** @returns Whether the change resulted in the cache being cleared */
@@ -121,9 +122,12 @@ namespace ts {
                 // 3. Otherwise, we have a default/namespace import that can be imported by any name, and
                 //    `symbolTableKey` will be something undesirable like `export=` or `default`, so we try to
                 //    get a better name.
-                const importedName = exportKind === ExportKind.Named || isExternalModuleSymbol(namedSymbol)
+                const names = exportKind === ExportKind.Named || isExternalModuleSymbol(namedSymbol)
                     ? unescapeLeadingUnderscores(symbolTableKey)
-                    : getNameForExportedSymbol(namedSymbol, /*scriptTarget*/ undefined);
+                    : getNamesForExportedSymbol(namedSymbol, /*scriptTarget*/ undefined);
+
+                const symbolName = typeof names === "string" ? names : names[0];
+                const capitalizedSymbolName = typeof names === "string" ? undefined : names[1];
 
                 const moduleName = stripQuotes(moduleSymbol.name);
                 const id = exportInfoId++;
@@ -132,10 +136,11 @@ namespace ts {
                 const storedModuleSymbol = moduleSymbol.flags & SymbolFlags.Transient ? undefined : moduleSymbol;
                 if (!storedSymbol || !storedModuleSymbol) symbols.set(id, [symbol, moduleSymbol]);
 
-                exportInfo.add(key(importedName, symbol, isExternalModuleNameRelative(moduleName) ? undefined : moduleName, checker), {
+                exportInfo.add(key(symbolName, symbol, isExternalModuleNameRelative(moduleName) ? undefined : moduleName, checker), {
                     id,
                     symbolTableKey,
-                    symbolName: importedName,
+                    symbolName,
+                    capitalizedSymbolName,
                     moduleName,
                     moduleFile,
                     moduleFileName: moduleFile?.fileName,
@@ -152,24 +157,17 @@ namespace ts {
                 const result = exportInfo.get(key);
                 return result?.map(rehydrateCachedInfo);
             },
-            forEach: (importingFile, action) => {
+            search: (importingFile, preferCapitalized, matches, action) => {
                 if (importingFile !== usableByFileName) return;
                 exportInfo.forEach((info, key) => {
                     const { symbolName, ambientModuleName } = parseKey(key);
-                    const rehydrated = info.map(rehydrateCachedInfo);
-                    const filtered = rehydrated.filter((r, i) => isNotShadowedByDeeperNodeModulesPackage(r, info[i].packageName));
-                    if (filtered.length) {
-                        action(
-                            filtered,
-                            preferCapitalized => {
-                                const { symbol, exportKind } = rehydrated[0];
-                                const namedSymbol = exportKind === ExportKind.Default && getLocalSymbolForExportDefault(symbol) || symbol;
-                                return preferCapitalized
-                                    ? getNameForExportedSymbol(namedSymbol, /*scriptTarget*/ undefined, /*preferCapitalized*/ true)
-                                    : symbolName;
-                            },
-                            !!ambientModuleName,
-                            key);
+                    const name = preferCapitalized && info[0].capitalizedSymbolName || symbolName;
+                    if (matches(name, info[0].targetFlags)) {
+                        const rehydrated = info.map(rehydrateCachedInfo);
+                        const filtered = rehydrated.filter((r, i) => isNotShadowedByDeeperNodeModulesPackage(r, info[i].packageName));
+                        if (filtered.length) {
+                            action(filtered, name, !!ambientModuleName, key);
+                        }
                     }
                 });
             },

src/services/utilities.ts

@@ -3210,13 +3210,34 @@ namespace ts {
         return isArray(valueOrArray) ? first(valueOrArray) : valueOrArray;
     }
 
+    export function getNamesForExportedSymbol(symbol: Symbol, scriptTarget: ScriptTarget | undefined): string | [lowercase: string, capitalized: string] {
+        if (needsNameFromDeclaration(symbol)) {
+            const fromDeclaration = getDefaultLikeExportNameFromDeclaration(symbol);
+            if (fromDeclaration) return fromDeclaration;
+            const fileNameCase = codefix.moduleSymbolToValidIdentifier(getSymbolParentOrFail(symbol), scriptTarget, /*preferCapitalized*/ false);
+            const capitalized = codefix.moduleSymbolToValidIdentifier(getSymbolParentOrFail(symbol), scriptTarget, /*preferCapitalized*/ true);
+            if (fileNameCase === capitalized) return fileNameCase;
+            return [fileNameCase, capitalized];
+        }
+        return symbol.name;
+    }
+
     export function getNameForExportedSymbol(symbol: Symbol, scriptTarget: ScriptTarget | undefined, preferCapitalized?: boolean) {
-        if (!(symbol.flags & SymbolFlags.Transient) && (symbol.escapedName === InternalSymbolName.ExportEquals || symbol.escapedName === InternalSymbolName.Default)) {
+        if (needsNameFromDeclaration(symbol)) {
             // Name of "export default foo;" is "foo". Name of "export default 0" is the filename converted to camelCase.
-            return firstDefined(symbol.declarations, d => isExportAssignment(d) ? tryCast(skipOuterExpressions(d.expression), isIdentifier)?.text : undefined)
+            return getDefaultLikeExportNameFromDeclaration(symbol)
                 || codefix.moduleSymbolToValidIdentifier(getSymbolParentOrFail(symbol), scriptTarget, !!preferCapitalized);
         }
         return symbol.name;
+
+    }
+
+    function needsNameFromDeclaration(symbol: Symbol) {
+        return !(symbol.flags & SymbolFlags.Transient) && (symbol.escapedName === InternalSymbolName.ExportEquals || symbol.escapedName === InternalSymbolName.Default);
+    }
+
+    function getDefaultLikeExportNameFromDeclaration(symbol: Symbol) {
+        return firstDefined(symbol.declarations, d => isExportAssignment(d) ? tryCast(skipOuterExpressions(d.expression), isIdentifier)?.text : undefined);
     }
 
     function getSymbolParentOrFail(symbol: Symbol) {

src/testRunner/unittests/tsserver/exportMapCache.ts

@@ -87,8 +87,8 @@ namespace ts.projectSystem {
             // transient symbols are recreated with every new checker.
             const programBefore = project.getCurrentProgram()!;
             let sigintPropBefore: readonly SymbolExportInfo[] | undefined;
-            exportMapCache.forEach(bTs.path as Path, (info, getSymbolName) => {
-                if (getSymbolName() === "SIGINT") sigintPropBefore = info;
+            exportMapCache.search(bTs.path as Path, /*preferCapitalized*/ false, returnTrue, (info, symbolName) => {
+                if (symbolName === "SIGINT") sigintPropBefore = info;
             });
             assert.ok(sigintPropBefore);
             assert.ok(sigintPropBefore![0].symbol.flags & SymbolFlags.Transient);
@@ -113,8 +113,8 @@ namespace ts.projectSystem {
 
             // Get same info from cache again
             let sigintPropAfter: readonly SymbolExportInfo[] | undefined;
-            exportMapCache.forEach(bTs.path as Path, (info, getSymbolName) => {
-                if (getSymbolName() === "SIGINT") sigintPropAfter = info;
+            exportMapCache.search(bTs.path as Path, /*preferCapitalized*/ false, returnTrue, (info, symbolName) => {
+                if (symbolName === "SIGINT") sigintPropAfter = info;
             });
             assert.ok(sigintPropAfter);
             assert.notEqual(symbolIdBefore, getSymbolId(sigintPropAfter![0].symbol));