perf(lsp): release unused documents (#23398)

1 files changed, 44 insertions, 23 deletions

diff --git a/cli/tsc/99_main_compiler.js b/cli/tsc/99_main_compiler.js
index b76c95aa5..0677b4c27 100644
--- a/cli/tsc/99_main_compiler.js
+++ b/cli/tsc/99_main_compiler.js
@@ -155,6 +155,12 @@ delete Object.prototype.__proto__;
   /** @type {Map<string, ts.SourceFile>} */
   const sourceFileCache = new Map();
 
+  /** @type {Map<string, string>} */
+  const sourceTextCache = new Map();
+
+  /** @type {Map<string, number>} */
+  const sourceRefCounts = new Map();
+
   /** @type {string[]=} */
   let scriptFileNamesCache;
 
@@ -172,6 +178,8 @@ delete Object.prototype.__proto__;
   /** @type {number | null} */
   let projectVersionCache = null;
 
+  let lastRequestMethod = null;
+
   const ChangeKind = {
     Opened: 0,
     Modified: 1,
@@ -250,6 +258,8 @@ delete Object.prototype.__proto__;
         );
         documentRegistrySourceFileCache.set(mapKey, sourceFile);
       }
+      const sourceRefCount = sourceRefCounts.get(fileName) ?? 0;
+      sourceRefCounts.set(fileName, sourceRefCount + 1);
       return sourceFile;
     },
 
@@ -333,8 +343,20 @@ delete Object.prototype.__proto__;
     },
 
     releaseDocumentWithKey(path, key, _scriptKind, _impliedNodeFormat) {
-      const mapKey = path + key;
-      documentRegistrySourceFileCache.delete(mapKey);
+      const sourceRefCount = sourceRefCounts.get(path) ?? 1;
+      if (sourceRefCount <= 1) {
+        sourceRefCounts.delete(path);
+        // We call `cleanupSemanticCache` for other purposes, don't bust the
+        // source cache in this case.
+        if (lastRequestMethod != "cleanupSemanticCache") {
+          const mapKey = path + key;
+          documentRegistrySourceFileCache.delete(mapKey);
+          sourceTextCache.delete(path);
+          ops.op_release(path);
+        }
+      } else {
+        sourceRefCounts.set(path, sourceRefCount - 1);
+      }
     },
 
     reportStats() {
@@ -807,19 +829,26 @@ delete Object.prototype.__proto__;
       if (logDebug) {
         debug(`host.getScriptSnapshot("${specifier}")`);
       }
-      let sourceFile = sourceFileCache.get(specifier);
-      if (!sourceFile) {
-        sourceFile = this.getSourceFile(
-          specifier,
-          specifier.endsWith(".json")
-            ? ts.ScriptTarget.JSON
-            : ts.ScriptTarget.ESNext,
-        );
-      }
+      const sourceFile = sourceFileCache.get(specifier);
       if (sourceFile) {
+        // This case only occurs for assets.
         return ts.ScriptSnapshot.fromString(sourceFile.text);
       }
-      return undefined;
+      let sourceText = sourceTextCache.get(specifier);
+      if (sourceText == undefined) {
+        /** @type {{ data: string, version: string, isCjs: boolean }} */
+        const fileInfo = ops.op_load(specifier);
+        if (!fileInfo) {
+          return undefined;
+        }
+        if (fileInfo.isCjs) {
+          isCjsCache.add(specifier);
+        }
+        sourceTextCache.set(specifier, fileInfo.data);
+        scriptVersionCache.set(specifier, fileInfo.version);
+        sourceText = fileInfo.data;
+      }
+      return ts.ScriptSnapshot.fromString(sourceText);
     },
   };
 
@@ -1047,6 +1076,7 @@ delete Object.prototype.__proto__;
     if (logDebug) {
       debug(`serverRequest()`, id, method, args);
     }
+    lastRequestMethod = method;
     switch (method) {
       case "$projectChanged": {
         /** @type {[string, number][]} */
@@ -1058,6 +1088,7 @@ delete Object.prototype.__proto__;
 
         if (configChanged) {
           tsConfigCache = null;
+          isNodeSourceFileCache.clear();
         }
 
         projectVersionCache = newProjectVersion;
@@ -1068,7 +1099,7 @@ delete Object.prototype.__proto__;
             opened = true;
           }
           scriptVersionCache.delete(script);
-          sourceFileCache.delete(script);
+          sourceTextCache.delete(script);
         }
 
         if (configChanged || opened) {
@@ -1077,10 +1108,6 @@ delete Object.prototype.__proto__;
 
         return respond(id);
       }
-      case "$restart": {
-        serverRestart();
-        return respond(id, true);
-      }
       case "$getSupportedCodeFixes": {
         return respond(
           id,
@@ -1152,12 +1179,6 @@ delete Object.prototype.__proto__;
     debug("serverInit()");
   }
 
-  function serverRestart() {
-    languageService = ts.createLanguageService(host, documentRegistry);
-    isNodeSourceFileCache.clear();
-    debug("serverRestart()");
-  }
-
   // A build time only op that provides some setup information that is used to
   // ensure the snapshot is setup properly.
   /** @type {{ buildSpecifier: string; libs: string[]; nodeBuiltInModuleNames: string[] }} */