Refactor export map cache to not store transient symbols (#44816)

* Add some failing tests around transient symbols

* Working, but slower

* A class is much faster, apparently

* This is probably best?

* Back to multimap

* Go back to single symbol cache

* Revert now-unnecessary generics

* Rename and reorganize

* Fix weird compound condition

* Clean up

Author: andrewbranch

src/server/moduleSpecifierCache.ts

@@ -0,0 +1,95 @@
+/*@internal*/
+namespace ts.server {
+    export interface ModuleSpecifierResolutionCacheHost {
+        watchNodeModulesForPackageJsonChanges(directoryPath: string): FileWatcher;
+    }
+
+    export function createModuleSpecifierCache(host: ModuleSpecifierResolutionCacheHost): ModuleSpecifierCache {
+        let containedNodeModulesWatchers: ESMap<string, FileWatcher> | undefined;
+        let cache: ESMap<Path, ResolvedModuleSpecifierInfo> | undefined;
+        let currentKey: string | undefined;
+        const result: ModuleSpecifierCache = {
+            get(fromFileName, toFileName, preferences) {
+                if (!cache || currentKey !== key(fromFileName, preferences)) return undefined;
+                return cache.get(toFileName);
+            },
+            set(fromFileName, toFileName, preferences, modulePaths, moduleSpecifiers) {
+                ensureCache(fromFileName, preferences).set(toFileName, createInfo(modulePaths, moduleSpecifiers, /*isAutoImportable*/ true));
+
+                // If any module specifiers were generated based off paths in node_modules,
+                // a package.json file in that package was read and is an input to the cached.
+                // Instead of watching each individual package.json file, set up a wildcard
+                // directory watcher for any node_modules referenced and clear the cache when
+                // it sees any changes.
+                if (moduleSpecifiers) {
+                    for (const p of modulePaths) {
+                        if (p.isInNodeModules) {
+                            // No trailing slash
+                            const nodeModulesPath = p.path.substring(0, p.path.indexOf(nodeModulesPathPart) + nodeModulesPathPart.length - 1);
+                            if (!containedNodeModulesWatchers?.has(nodeModulesPath)) {
+                                (containedNodeModulesWatchers ||= new Map()).set(
+                                    nodeModulesPath,
+                                    host.watchNodeModulesForPackageJsonChanges(nodeModulesPath),
+                                );
+                            }
+                        }
+                    }
+                }
+            },
+            setModulePaths(fromFileName, toFileName, preferences, modulePaths) {
+                const cache = ensureCache(fromFileName, preferences);
+                const info = cache.get(toFileName);
+                if (info) {
+                    info.modulePaths = modulePaths;
+                }
+                else {
+                    cache.set(toFileName, createInfo(modulePaths, /*moduleSpecifiers*/ undefined, /*isAutoImportable*/ undefined));
+                }
+            },
+            setIsAutoImportable(fromFileName, toFileName, preferences, isAutoImportable) {
+                const cache = ensureCache(fromFileName, preferences);
+                const info = cache.get(toFileName);
+                if (info) {
+                    info.isAutoImportable = isAutoImportable;
+                }
+                else {
+                    cache.set(toFileName, createInfo(/*modulePaths*/ undefined, /*moduleSpecifiers*/ undefined, isAutoImportable));
+                }
+            },
+            clear() {
+                containedNodeModulesWatchers?.forEach(watcher => watcher.close());
+                cache?.clear();
+                containedNodeModulesWatchers?.clear();
+                currentKey = undefined;
+            },
+            count() {
+                return cache ? cache.size : 0;
+            }
+        };
+        if (Debug.isDebugging) {
+            Object.defineProperty(result, "__cache", { get: () => cache });
+        }
+        return result;
+
+        function ensureCache(fromFileName: Path, preferences: UserPreferences) {
+            const newKey = key(fromFileName, preferences);
+            if (cache && (currentKey !== newKey)) {
+                result.clear();
+            }
+            currentKey = newKey;
+            return cache ||= new Map();
+        }
+
+        function key(fromFileName: Path, preferences: UserPreferences) {
+            return `${fromFileName},${preferences.importModuleSpecifierEnding},${preferences.importModuleSpecifierPreference}`;
+        }
+
+        function createInfo(
+            modulePaths: readonly ModulePath[] | undefined,
+            moduleSpecifiers: readonly string[] | undefined,
+            isAutoImportable: boolean | undefined,
+        ): ResolvedModuleSpecifierInfo {
+            return { modulePaths, moduleSpecifiers, isAutoImportable };
+        }
+    }
+}

src/server/project.ts

@@ -251,7 +251,7 @@ namespace ts.server {
         public readonly getCanonicalFileName: GetCanonicalFileName;
 
         /*@internal*/
-        private exportMapCache = createExportMapCache();
+        private exportMapCache: ExportInfoMap | undefined;
         /*@internal*/
         private changedFilesForExportMapCache: Set<Path> | undefined;
         /*@internal*/
@@ -793,6 +793,7 @@ namespace ts.server {
             this.cachedUnresolvedImportsPerFile = undefined!;
             this.moduleSpecifierCache = undefined!;
             this.directoryStructureHost = undefined!;
+            this.exportMapCache = undefined;
             this.projectErrors = undefined;
 
             // Clean up file watchers waiting for missing files
@@ -990,7 +991,7 @@ namespace ts.server {
         /*@internal*/
         markFileAsDirty(changedFile: Path) {
             this.markAsDirty();
-            if (!this.exportMapCache.isEmpty()) {
+            if (this.exportMapCache && !this.exportMapCache.isEmpty()) {
                 (this.changedFilesForExportMapCache ||= new Set()).add(changedFile);
             }
         }
@@ -1192,7 +1193,8 @@ namespace ts.server {
                 }
             }
 
-            if (!this.exportMapCache.isEmpty()) {
+            if (this.exportMapCache && !this.exportMapCache.isEmpty()) {
+                this.exportMapCache.releaseSymbols();
                 if (this.hasAddedorRemovedFiles || oldProgram && !this.program!.structureIsReused) {
                     this.exportMapCache.clear();
                 }
@@ -1201,10 +1203,10 @@ namespace ts.server {
                         const oldSourceFile = oldProgram.getSourceFileByPath(fileName);
                         const sourceFile = this.program!.getSourceFileByPath(fileName);
                         if (!oldSourceFile || !sourceFile) {
-                            this.exportMapCache.clear();
+                            this.exportMapCache!.clear();
                             return true;
                         }
-                        return this.exportMapCache.onFileChanged(oldSourceFile, sourceFile, !!this.getTypeAcquisition().enable);
+                        return this.exportMapCache!.onFileChanged(oldSourceFile, sourceFile, !!this.getTypeAcquisition().enable);
                     });
                 }
             }
@@ -1666,8 +1668,13 @@ namespace ts.server {
         }
 
         /*@internal*/
-        getExportMapCache() {
-            return this.exportMapCache;
+        getCachedExportInfoMap() {
+            return this.exportMapCache ||= createCacheableExportInfoMap(this);
+        }
+
+        /*@internal*/
+        clearCachedExportInfoMap() {
+            this.exportMapCache?.clear();
         }
 
         /*@internal*/
@@ -2017,7 +2024,7 @@ namespace ts.server {
             const oldProgram = this.getCurrentProgram();
             const hasSameSetOfFiles = super.updateGraph();
             if (oldProgram && oldProgram !== this.getCurrentProgram()) {
-                this.hostProject.getExportMapCache().clear();
+                this.hostProject.clearCachedExportInfoMap();
             }
             return hasSameSetOfFiles;
         }

src/server/tsconfig.json

@@ -23,6 +23,7 @@
         "typingsCache.ts",
         "project.ts",
         "editorServices.ts",
+        "moduleSpecifierCache.ts",
         "packageJsonCache.ts",
         "session.ts",
         "scriptVersionCache.ts"