NETSCRIPT: Compiled modules will be even more shared (#468)

src/NetscriptJSEvaluator.ts

@@ -5,9 +5,9 @@
 import * as walk from "acorn-walk";
 import { parse } from "acorn";
 
-import { Script, ScriptURL } from "./Script/Script";
-import { areImportsEquals } from "./Terminal/DirectoryHelpers";
-import { ScriptModule } from "./Script/ScriptModule";
+import { LoadedModule, ScriptURL, ScriptModule } from "./Script/LoadedModule";
+import { Script } from "./Script/Script";
+import { areImportsEquals, removeLeadingSlash } from "./Terminal/DirectoryHelpers";
 
 // Acorn type def is straight up incomplete so we have to fill with our own.
 export type Node = any;
@@ -17,24 +17,6 @@ function makeScriptBlob(code: string): Blob {
   return new Blob([code], { type: "text/javascript" });
 }
 
-const urlsToRevoke: ScriptURL[] = [];
-let activeCompilations = 0;
-/** Function to queue up revoking of script URLs. If there's no active compilation, just revoke it now. */
-export const queueUrlRevoke = (url: ScriptURL) => {
-  if (!activeCompilations) return URL.revokeObjectURL(url);
-  urlsToRevoke.push(url);
-};
-
-/** Function to revoke any expired urls */
-function triggerURLRevokes() {
-  if (activeCompilations === 0) {
-    // Revoke all pending revoke URLS
-    urlsToRevoke.forEach((url) => URL.revokeObjectURL(url));
-    // Remove all url strings from array
-    urlsToRevoke.length = 0;
-  }
-}
-
 // Webpack likes to turn the import into a require, which sort of
 // but not really behaves like import. So we use a "magic comment"
 // to disable that and leave it as a dynamic import.
@@ -51,50 +33,53 @@ export const config = {
   },
 };
 
+// Maps code to LoadedModules, so we can reuse compiled code across servers,
+// or possibly across files (if someone makes two copies of the same script,
+// or changes a script and then changes it back).
+// Modules can never be garbage collected by Javascript, so it's good to try
+// to keep from making more than we need.
+const moduleCache: Map<string, WeakRef<LoadedModule>> = new Map();
+const cleanup = new FinalizationRegistry((mapKey: string) => {
+  // A new entry can be created with the same key, before this callback is called.
+  if (moduleCache.get(mapKey)?.deref() === undefined) {
+    moduleCache.delete(mapKey);
+  }
+});
+
 export function compile(script: Script, scripts: Script[]): Promise<ScriptModule> {
   // Return the module if it already exists
-  if (script.module) return script.module;
-  // While importing, use an existing url or generate a new one.
-  if (!script.url) script.url = generateScriptUrl(script, scripts, []);
-  activeCompilations++;
-  script.module = config
-    .doImport(script.url)
-    .catch((e) => {
-      script.invalidateModule();
-      console.error(`Error occurred while attempting to compile ${script.filename} on ${script.server}:`);
-      console.error(e);
-      throw e;
-    })
-    .finally(() => {
-      activeCompilations--;
-      triggerURLRevokes();
-    });
-  return script.module;
+  if (script.mod) return script.mod.module;
+
+  script.mod = generateLoadedModule(script, scripts, []);
+  return script.mod.module;
 }
 
 /** Add the necessary dependency relationships for a script.
  * Dependents are used only for passing invalidation up an import tree, so only direct dependents need to be stored.
  * Direct and indirect dependents need to have the current url/script added to their dependency map for error text.
  *
- * This should only be called once the script has an assigned URL. */
-function addDependencyInfo(script: Script, dependents: Script[]) {
-  if (!script.url) throw new Error(`addDependencyInfo called without an assigned script URL (${script.filename})`);
-  if (dependents.length) {
-    script.dependents.add(dependents[dependents.length - 1]);
-    for (const dependent of dependents) dependent.dependencies.set(script.url, script);
+ * This should only be called once the script has a LoadedModule. */
+function addDependencyInfo(script: Script, seenStack: Script[]) {
+  if (!script.mod) throw new Error(`addDependencyInfo called without a LoadedModule (${script.filename})`);
+  if (seenStack.length) {
+    script.dependents.add(seenStack[seenStack.length - 1]);
+    for (const dependent of seenStack) dependent.dependencies.set(script.mod.url, script);
   }
+  // Add self to dependencies (it's not part of the stack, since we don't want
+  // it in dependents.)
+  script.dependencies.set(script.mod.url, script);
 }
 
 /**
  * @param script the script that needs a URL assigned
  * @param scripts array of other scripts on the server
- * @param dependents All scripts that were higher up in the import tree in a recursive call.
+ * @param seenStack A stack of scripts that were higher up in the import tree in a recursive call.
  */
-function generateScriptUrl(script: Script, scripts: Script[], dependents: Script[]): ScriptURL {
+function generateLoadedModule(script: Script, scripts: Script[], seenStack: Script[]): LoadedModule {
   // Early return for recursive calls where the script already has a URL
-  if (script.url) {
-    addDependencyInfo(script, dependents);
-    return script.url;
+  if (script.mod) {
+    addDependencyInfo(script, seenStack);
+    return script.mod;
   }
 
   // Inspired by: https://stackoverflow.com/a/43834063/91401
@@ -145,14 +130,40 @@ function generateScriptUrl(script: Script, scripts: Script[], dependents: Script
     const importedScript = scripts.find((s) => areImportsEquals(s.filename, filename));
     if (!importedScript) continue;
 
-    importedScript.url = generateScriptUrl(importedScript, scripts, [...dependents, script]);
-    newCode = newCode.substring(0, node.start) + importedScript.url + newCode.substring(node.end);
+    seenStack.push(script);
+    importedScript.mod = generateLoadedModule(importedScript, scripts, seenStack);
+    seenStack.pop();
+    newCode = newCode.substring(0, node.start) + importedScript.mod.url + newCode.substring(node.end);
   }
 
-  newCode += `\n//# sourceURL=${script.server}/${script.filename}`;
+  const cachedMod = moduleCache.get(newCode)?.deref();
+  if (cachedMod) {
+    script.mod = cachedMod;
+  } else {
+    // Add an inline source-map to make debugging nicer. This won't be right
+    // in all cases, since we can share the same script across multiple
+    // servers; it will be listed under the first server it was compiled for.
+    // We don't include this in the cache key, so that other instances of the
+    // script dedupe properly.
+    const adjustedCode = newCode + `\n//# sourceURL=${script.server}/${removeLeadingSlash(script.filename)}`;
+    // At this point we have the full code and can construct a new blob / assign the URL.
+    const url = URL.createObjectURL(makeScriptBlob(adjustedCode)) as ScriptURL;
+    const module = config.doImport(url).catch((e) => {
+      script.invalidateModule();
+      console.error(`Error occurred while attempting to compile ${script.filename} on ${script.server}:`);
+      console.error(e);
+      throw e;
+    }) as Promise<ScriptModule>;
+    // We can *immediately* invalidate the Blob, because we've already started the fetch
+    // by starting the import. From now on, any imports using the blob's URL *must*
+    // directly return the module, without even attempting to fetch, due to the way
+    // modules work.
+    URL.revokeObjectURL(url);
+    script.mod = new LoadedModule(url, module);
+    moduleCache.set(newCode, new WeakRef(script.mod));
+    cleanup.register(script.mod, newCode);
+  }
 
-  // At this point we have the full code and can construct a new blob / assign the URL.
-  script.url = URL.createObjectURL(makeScriptBlob(newCode)) as ScriptURL;
-  addDependencyInfo(script, dependents);
-  return script.url;
+  addDependencyInfo(script, seenStack);
+  return script.mod;
 }