Cache the latest source file from document registry so we can keep it alive when script info is orphan

src/harness/unittests/tsserverProjectSystem.ts

@@ -8405,15 +8405,21 @@ new C();`
             service.openClientFile(file.path);
             const project = service.configuredProjects.get(configFile.path);
             checkProject(/*moduleIsOrphan*/ false);
+            const moduleInfo = service.getScriptInfo(moduleFile.path);
+            const sourceFile = moduleInfo.cacheSourceFile.sourceFile;
+            assert.equal(project.getSourceFile(moduleInfo.path), sourceFile);
 
             // edit file
             const info = service.getScriptInfo(file.path);
             service.applyChangesToFile(info, [{ span: { start: 0, length: importModuleContent.length }, newText: "" }]);
             checkProject(/*moduleIsOrphan*/ true);
+            assert.equal(moduleInfo.cacheSourceFile.sourceFile, sourceFile);
 
             // write content back
             service.applyChangesToFile(info, [{ span: { start: 0, length: 0 }, newText: importModuleContent }]);
             checkProject(/*moduleIsOrphan*/ false);
+            assert.equal(moduleInfo.cacheSourceFile.sourceFile, sourceFile);
+            assert.equal(project.getSourceFile(moduleInfo.path), sourceFile);
 
             function checkProject(moduleIsOrphan: boolean) {
                 // Update the project

src/server/editorServices.ts

@@ -475,7 +475,7 @@ namespace ts.server {
                 extraFileExtensions: []
             };
 
-            this.documentRegistry = createDocumentRegistry(this.host.useCaseSensitiveFileNames, this.currentDirectory);
+            this.documentRegistry = createDocumentRegistryInternal(this.host.useCaseSensitiveFileNames, this.currentDirectory, this);
             const watchLogLevel = this.logger.hasLevel(LogLevel.verbose) ? WatchLogLevel.Verbose :
                 this.logger.loggingEnabled() ? WatchLogLevel.TriggerOnly : WatchLogLevel.None;
             const log: (s: string) => void = watchLogLevel !== WatchLogLevel.None ? (s => this.logger.info(s)) : noop;
@@ -496,6 +496,19 @@ namespace ts.server {
             return getNormalizedAbsolutePath(fileName, this.host.getCurrentDirectory());
         }
 
+        /*@internal*/
+        setDocument(key: DocumentRegistryBucketKey, path: Path, sourceFile: SourceFile) {
+            const info = this.getScriptInfoForPath(path);
+            Debug.assert(!!info);
+            info.cacheSourceFile = { key, sourceFile };
+        }
+
+        /*@internal*/
+        getDocument(key: DocumentRegistryBucketKey, path: Path) {
+            const info = this.getScriptInfoForPath(path);
+            return info && info.cacheSourceFile && info.cacheSourceFile.key === key && info.cacheSourceFile.sourceFile;
+        }
+
         /* @internal */
         ensureInferredProjectsUpToDate_TestOnly() {
             this.ensureProjectStructuresUptoDate();

src/server/scriptInfo.ts

@@ -202,6 +202,12 @@ namespace ts.server {
         return fileName[0] === "^" || getBaseFileName(fileName)[0] === "^";
     }
 
+    /*@internal*/
+    export interface DocumentRegistrySourceFileCache {
+        key: DocumentRegistryBucketKey;
+        sourceFile: SourceFile;
+    }
+
     export class ScriptInfo {
         /**
          * All projects that include this file
@@ -221,6 +227,9 @@ namespace ts.server {
         /** Set to real path if path is different from info.path */
         private realpath: Path | undefined;
 
+        /*@internal*/
+        cacheSourceFile: DocumentRegistrySourceFileCache;
+
         constructor(
             private readonly host: ServerHost,
             readonly fileName: NormalizedPath,

src/services/documentRegistry.ts

@@ -93,6 +93,12 @@ namespace ts {
         reportStats(): string;
     }
 
+    /*@internal*/
+    export interface ExternalDocumentCache {
+        setDocument(key: DocumentRegistryBucketKey, path: Path, sourceFile: SourceFile): void;
+        getDocument(key: DocumentRegistryBucketKey, path: Path): SourceFile | undefined;
+    }
+
     export type DocumentRegistryBucketKey = string & { __bucketKey: any };
 
     interface DocumentRegistryEntry {
@@ -102,10 +108,14 @@ namespace ts {
         // language services are referencing the file, then the file can be removed from the
         // registry.
         languageServiceRefCount: number;
-        owners: string[];
     }
 
-    export function createDocumentRegistry(useCaseSensitiveFileNames?: boolean, currentDirectory = ""): DocumentRegistry {
+    export function createDocumentRegistry(useCaseSensitiveFileNames?: boolean, currentDirectory?: string): DocumentRegistry {
+        return createDocumentRegistryInternal(useCaseSensitiveFileNames, currentDirectory);
+    }
+
+    /*@internal*/
+    export function createDocumentRegistryInternal(useCaseSensitiveFileNames?: boolean, currentDirectory = "", externalCache?: ExternalDocumentCache): DocumentRegistry {
         // Maps from compiler setting target (ES3, ES5, etc.) to all the cached documents we have
         // for those settings.
         const buckets = createMap<Map<DocumentRegistryEntry>>();
@@ -126,12 +136,11 @@ namespace ts {
         function reportStats() {
             const bucketInfoArray = arrayFrom(buckets.keys()).filter(name => name && name.charAt(0) === "_").map(name => {
                 const entries = buckets.get(name);
-                const sourceFiles: { name: string; refCount: number; references: string[]; }[] = [];
+                const sourceFiles: { name: string; refCount: number; }[] = [];
                 entries.forEach((entry, name) => {
                     sourceFiles.push({
                         name,
-                        refCount: entry.languageServiceRefCount,
-                        references: entry.owners.slice(0)
+                        refCount: entry.languageServiceRefCount
                     });
                 });
                 sourceFiles.sort((x, y) => y.refCount - x.refCount);
@@ -176,14 +185,26 @@ namespace ts {
             const bucket = getBucketForCompilationSettings(key, /*createIfMissing*/ true);
             let entry = bucket.get(path);
             const scriptTarget = scriptKind === ScriptKind.JSON ? ScriptTarget.JSON : compilationSettings.target;
+            if (!entry && externalCache) {
+                const sourceFile = externalCache.getDocument(key, path);
+                if (sourceFile) {
+                    entry = {
+                        sourceFile,
+                        languageServiceRefCount: 1
+                    };
+                    bucket.set(path, entry);
+                }
+            }
+
             if (!entry) {
                 // Have never seen this file with these settings.  Create a new source file for it.
                 const sourceFile = createLanguageServiceSourceFile(fileName, scriptSnapshot, scriptTarget, version, /*setNodeParents*/ false, scriptKind);
-
+                if (externalCache) {
+                    externalCache.setDocument(key, path, sourceFile);
+                }
                 entry = {
                     sourceFile,
                     languageServiceRefCount: 1,
-                    owners: []
                 };
                 bucket.set(path, entry);
             }
@@ -194,6 +215,9 @@ namespace ts {
                 if (entry.sourceFile.version !== version) {
                     entry.sourceFile = updateLanguageServiceSourceFile(entry.sourceFile, scriptSnapshot, version,
                         scriptSnapshot.getChangeRange(entry.sourceFile.scriptSnapshot));
+                    if (externalCache) {
+                        externalCache.setDocument(key, path, entry.sourceFile);
+                    }
                 }
 
                 // If we're acquiring, then this is the first time this LS is asking for this document.