From 47bc358f44e8f1b3f4e77bf5de7722b89d80fe3a Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <jridgewell@google.com>
Date: Wed, 6 May 2020 03:03:37 -0400
Subject: [PATCH] Cache module code when running tests

---
 .../package.json                              |  1 +
 .../src/index.js                              | 78 +++++++++++++------
 2 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/packages/babel-helper-transform-fixture-test-runner/package.json b/packages/babel-helper-transform-fixture-test-runner/package.json
index 3138c4c0dfbb..a4cd5bcd7b95 100644
--- a/packages/babel-helper-transform-fixture-test-runner/package.json
+++ b/packages/babel-helper-transform-fixture-test-runner/package.json
@@ -19,6 +19,7 @@
     "jest": "^24.8.0",
     "jest-diff": "^24.8.0",
     "lodash": "^4.17.13",
+    "quick-lru": "5.1.0",
     "resolve": "^1.3.2",
     "source-map": "^0.5.0"
   }
diff --git a/packages/babel-helper-transform-fixture-test-runner/src/index.js b/packages/babel-helper-transform-fixture-test-runner/src/index.js
index d5474296262e..6d5647d8d022 100644
--- a/packages/babel-helper-transform-fixture-test-runner/src/index.js
+++ b/packages/babel-helper-transform-fixture-test-runner/src/index.js
@@ -16,9 +16,11 @@ import fs from "fs";
 import path from "path";
 import vm from "vm";
 import checkDuplicatedNodes from "babel-check-duplicated-nodes";
+import QuickLRU from "quick-lru";
 
 import diff from "jest-diff";
 
+const cachedScripts = new QuickLRU({ maxSize: 512 });
 const contextModuleCache = new WeakMap();
 const sharedTestContext = createContext();
 
@@ -40,26 +42,61 @@ function createContext() {
   runModuleInTestContext("@babel/polyfill", __filename, context);
 
   // Populate the "babelHelpers" global with Babel's helper utilities.
-  runCodeInTestContext(
-    buildExternalHelpers(),
-    {
-      filename: path.join(__dirname, "babel-helpers-in-memory.js"),
-    },
+  runCacheableScriptInTestContext(
+    path.join(__dirname, "babel-helpers-in-memory.js"),
+    buildExternalHelpers,
     context,
   );
 
   return context;
 }
 
+function runCacheableScriptInTestContext(
+  filename: string,
+  srcFn: () => string,
+  context: Context,
+) {
+  let cached = cachedScripts.get(filename);
+  if (!cached) {
+    const code = `(function (exports, require, module, __filename, __dirname) {\n${srcFn()}\n});`;
+    cached = {
+      code,
+      cachedData: undefined,
+    };
+    cachedScripts.set(filename, cached);
+  }
+
+  const script = new vm.Script(cached.code, {
+    filename,
+    displayErrors: true,
+    lineOffset: -1,
+    cachedData: cached.cachedData,
+    produceCachedData: true,
+  });
+
+  if (script.cachedDataProduced) {
+    cached.cachedData = script.cachedData;
+  }
+
+  const module = {
+    id: filename,
+    exports: {},
+  };
+  const req = id => runModuleInTestContext(id, filename, context);
+  const dirname = path.dirname(filename);
+
+  script
+    .runInContext(context)
+    .call(module.exports, module.exports, req, module, filename, dirname);
+
+  return module;
+}
+
 /**
  * A basic implementation of CommonJS so we can execute `@babel/polyfill` inside our test context.
  * This allows us to run our unittests
  */
-function runModuleInTestContext(
-  id: string,
-  relativeFilename: string,
-  context = sharedTestContext,
-) {
+function runModuleInTestContext(id: string, relativeFilename: string, context) {
   const filename = resolve.sync(id, {
     basedir: path.dirname(relativeFilename),
   });
@@ -68,24 +105,17 @@ function runModuleInTestContext(
   // the context's global scope.
   if (filename === id) return require(id);
 
+  // Modules can only evaluate once per context, so the moduleCache is a
+  // stronger cache guarantee than the LRU's Script cache.
   const moduleCache = contextModuleCache.get(context);
   if (moduleCache[filename]) return moduleCache[filename].exports;
 
-  const module = (moduleCache[filename] = {
-    id: filename,
-    exports: {},
-  });
-  const dirname = path.dirname(filename);
-  const req = id => runModuleInTestContext(id, filename, context);
-
-  const src = fs.readFileSync(filename, "utf8");
-  const code = `(function (exports, require, module, __filename, __dirname) {\n${src}\n});`;
-
-  vm.runInContext(code, context, {
+  const module = runCacheableScriptInTestContext(
     filename,
-    displayErrors: true,
-    lineOffset: -1,
-  }).call(module.exports, module.exports, req, module, filename, dirname);
+    () => fs.readFileSync(filename, "utf8"),
+    context,
+  );
+  moduleCache[filename] = module;
 
   return module.exports;
 }