vm.Script: parse once and reuse the unlinked code in every context it runs in

scripts/build/deps/webkit.ts

@@ -3,7 +3,7 @@
  * for local mode. Override via `--webkit-version=<hash>` to test a branch.
  * From https://github.com/oven-sh/WebKit releases.
  */
-export const WEBKIT_VERSION = "7b763944f0ecd0e18cb9ac89a04ef994e3ccb4ae";
+export const WEBKIT_VERSION = "autobuild-preview-pr-418-79de1898";
 
 /**
  * WebKit (JavaScriptCore) — the JS engine.

src/jsc/bindings/NodeVMScript.cpp

@@ -3,6 +3,7 @@
 
 #include "ErrorCode.h"
 
+#include "JavaScriptCore/CodeCache.h"
 #include "JavaScriptCore/Completion.h"
 #include "JavaScriptCore/JIT.h"
 #include "JavaScriptCore/JSWeakMap.h"
@@ -131,39 +132,21 @@ constructScript(JSGlobalObject* globalObject, CallFrame* callFrame, JSValue newT
 
     RefPtr fetcher(NodeVMScriptFetcher::create(vm, importer, jsUndefined()));
 
-    SourceCode source = makeSource(sourceString, JSC::SourceOrigin(WTF::URL::fileURLWithFileSystemPath(options.filename), *fetcher), JSC::SourceTaintedOrigin::Untainted, options.filename, TextPosition(options.lineOffset, options.columnOffset));
+    TextPosition startPosition(options.lineOffset, options.columnOffset);
+    Ref provider = NodeVMScriptSourceProvider::create(sourceString, JSC::SourceOrigin(WTF::URL::fileURLWithFileSystemPath(options.filename), *fetcher), String(options.filename), startPosition);
+    SourceCode source(provider.copyRef(), startPosition.m_line.oneBasedInt(), startPosition.m_column.oneBasedInt());
     RETURN_IF_EXCEPTION(scope, {});
 
-    // Node's vm.Script throws SyntaxError at construction; the REPL's
-    // recoverable-error flow (and user code) relies on that. This is a
-    // double-parse (checkSyntax discards its AST and runInThisContext reparses
-    // via JSC::evaluate); compile-once via m_cachedExecutable is the follow-up.
-    JSC::ParserError parseError;
-    if (!JSC::checkSyntax(vm, source, parseError)) {
-        auto exception = parseError.toErrorObject(globalObject, source, -1);
-        // Building the error materializes its stack, running a user
-        // Error.prepareStackTrace that may throw; Node throws the SyntaxError
-        // anyway. tryClearException leaves a termination for the check below.
-        if (exception)
-            (void)scope.tryClearException();
-        RETURN_IF_EXCEPTION(scope, {});
-        // Node always attaches the arrow header to compile-time SyntaxErrors
-        // (node_contextify.cc DecorateErrorStack), independent of displayErrors.
-        // An absent filename becomes evalmachine.<anonymous>; an explicitly
-        // provided one — including "" — is used verbatim.
-        String url = options.filenameProvided ? options.filename : "evalmachine.<anonymous>"_s;
-        decorateParseErrorStack(globalObject, vm, exception, sourceString, url, parseError, options.lineOffset);
-        throwException(globalObject, scope, exception);
-        return {};
-    }
-
     const bool produceCachedData = options.produceCachedData;
-    auto filename = options.filename;
+    const bool filenameProvided = options.filenameProvided;
+    const OrdinalNumber lineOffset = options.lineOffset;
+    String filename = options.filename;
 
     NodeVMScript* script = NodeVMScript::create(vm, globalObject, structure, WTF::move(source), WTF::move(options));
     RETURN_IF_EXCEPTION(scope, {});
 
     fetcher->owner(vm, script);
+    provider->setScript(script);
 
     WTF::Vector<uint8_t>& cachedData = script->cachedData();
 
@@ -194,11 +177,39 @@ constructScript(JSGlobalObject* globalObject, CallFrame* callFrame, JSValue newT
             if (compilationResult != JSC::CompilationResult::CompilationFailed) {
                 executable->installCode(codeBlock);
                 script->cachedDataRejected(TriState::False);
+                // Accepted cached data is what the runs link, so the source is never parsed.
+                provider->pinUnlinkedCode(key, unlinkedBlock);
             } else {
                 script->cachedDataRejected(TriState::True);
             }
         }
-    } else if (produceCachedData) {
+    }
+
+    if (!script->hasPinnedUnlinkedCode()) {
+        // Node's vm.Script throws SyntaxError at construction; the REPL's
+        // recoverable-error flow (and user code) relies on that. The same parse
+        // is pinned for every later run, see NodeVMScriptSourceProvider.
+        JSC::ParserError parseError;
+        if (!script->compile(globalObject, parseError)) {
+            auto exception = parseError.toErrorObject(globalObject, script->source(), -1);
+            // Building the error materializes its stack, running a user
+            // Error.prepareStackTrace that may throw; Node throws the SyntaxError
+            // anyway. tryClearException leaves a termination for the check below.
+            if (exception)
+                (void)scope.tryClearException();
+            RETURN_IF_EXCEPTION(scope, {});
+            // Node always attaches the arrow header to compile-time SyntaxErrors
+            // (node_contextify.cc DecorateErrorStack), independent of displayErrors.
+            // An absent filename becomes evalmachine.<anonymous>; an explicitly
+            // provided one — including "" — is used verbatim.
+            String url = filenameProvided ? filename : "evalmachine.<anonymous>"_s;
+            decorateParseErrorStack(globalObject, vm, exception, sourceString, url, parseError, lineOffset);
+            throwException(globalObject, scope, exception);
+            return {};
+        }
+    }
+
+    if (cachedData.isEmpty() && produceCachedData) {
         script->cacheBytecode();
         // TODO(@heimskr): is there ever a case where bytecode production fails?
         script->cachedDataProduced(true);
@@ -207,6 +218,37 @@ constructScript(JSGlobalObject* globalObject, CallFrame* callFrame, JSValue newT
     return JSValue::encode(script);
 }
 
+JSC::JSCell* NodeVMScriptSourceProvider::pinnedUnlinkedCode(const JSC::SourceCodeKey& key) const
+{
+    return m_script ? m_script->pinnedUnlinkedCode(key.hash()) : nullptr;
+}
+
+void NodeVMScriptSourceProvider::pinUnlinkedCode(const JSC::SourceCodeKey& key, JSC::JSCell* unlinkedCode) const
+{
+    if (m_script)
+        m_script->pinUnlinkedCode(key.hash(), uncheckedDowncast<JSC::UnlinkedProgramCodeBlock>(unlinkedCode));
+}
+
+NodeVMScript::~NodeVMScript()
+{
+    static_cast<NodeVMScriptSourceProvider*>(m_source.provider())->setScript(nullptr);
+}
+
+void NodeVMScript::pinUnlinkedCode(unsigned keyHash, JSC::UnlinkedProgramCodeBlock* unlinkedCode)
+{
+    m_pinnedKeyHash = keyHash;
+    m_pinnedUnlinkedCode.set(vm(), this, unlinkedCode);
+}
+
+bool NodeVMScript::compile(JSGlobalObject* globalObject, JSC::ParserError& error)
+{
+    VM& vm = JSC::getVM(globalObject);
+    JSC::ProgramExecutable* executable = m_cachedExecutable ? m_cachedExecutable.get() : createExecutable();
+    // A successful parse reaches NodeVMScriptSourceProvider::pinUnlinkedCode from inside the CodeCache.
+    vm.codeCache()->getUnlinkedProgramCodeBlock(vm, executable, m_source, globalObject->defaultCodeGenerationMode(), error);
+    return !error.isValid();
+}
+
 JSC_DEFINE_HOST_FUNCTION(scriptConstructorCall, (JSGlobalObject * globalObject, CallFrame* callFrame))
 {
     return constructScript(globalObject, callFrame);
@@ -268,6 +310,7 @@ void NodeVMScript::visitChildrenImpl(JSCell* cell, Visitor& visitor)
     Base::visitChildren(thisObject, visitor);
     visitor.append(thisObject->m_cachedExecutable);
     visitor.append(thisObject->m_cachedBytecodeBuffer);
+    visitor.append(thisObject->m_pinnedUnlinkedCode);
 }
 
 NodeVMScriptConstructor::NodeVMScriptConstructor(VM& vm, Structure* structure)

src/jsc/bindings/NodeVMScript.h

@@ -4,6 +4,11 @@
 
 #include "../vm/SigintReceiver.h"
 
+namespace JSC {
+class SourceCodeKey;
+class UnlinkedProgramCodeBlock;
+}
+
 namespace Bun {
 
 class ScriptOptions : public BaseVMOptions {
@@ -38,11 +43,43 @@ class NodeVMScriptConstructor final : public JSC::InternalFunction {
 
 STATIC_ASSERT_ISO_SUBSPACE_SHARABLE(NodeVMScriptConstructor, JSC::InternalFunction);
 
+class NodeVMScript;
+
+// The SourceProvider behind a vm.Script. JSC's CodeCache consults it before parsing the
+// script's source and hands it what it parses (SourceProvider::pinnedUnlinkedCode /
+// pinUnlinkedCode), so every runInContext / runInThisContext of one Script links the same
+// UnlinkedProgramCodeBlock, however many contexts it runs in and whatever else has gone
+// through the CodeCache since. The Script cell owns that code block (it is what keeps it
+// alive); the provider only knows which Script to ask. The provider can outlive the Script
+// through the executables and stack traces earlier runs created, so the Script unlinks
+// itself when it is destroyed.
+class NodeVMScriptSourceProvider final : public JSC::StringSourceProvider {
+public:
+    static Ref<NodeVMScriptSourceProvider> create(const WTF::String& source, const JSC::SourceOrigin& sourceOrigin, WTF::String&& sourceURL, const TextPosition& startPosition)
+    {
+        return adoptRef(*new NodeVMScriptSourceProvider(source, sourceOrigin, WTF::move(sourceURL), startPosition));
+    }
+
+    void setScript(NodeVMScript* script) { m_script = script; }
+
+    JSC::JSCell* pinnedUnlinkedCode(const JSC::SourceCodeKey&) const final;
+    void pinUnlinkedCode(const JSC::SourceCodeKey&, JSC::JSCell*) const final;
+
+private:
+    NodeVMScriptSourceProvider(const WTF::String& source, const JSC::SourceOrigin& sourceOrigin, WTF::String&& sourceURL, const TextPosition& startPosition)
+        : JSC::StringSourceProvider(source, sourceOrigin, JSC::SourceTaintedOrigin::Untainted, WTF::move(sourceURL), startPosition, JSC::SourceProviderSourceType::Program)
+    {
+    }
+
+    NodeVMScript* m_script = nullptr;
+};
+
 class NodeVMScript final : public JSC::JSDestructibleObject, public SigintReceiver {
 public:
     using Base = JSC::JSDestructibleObject;
 
     static NodeVMScript* create(JSC::VM& vm, JSC::JSGlobalObject* globalObject, JSC::Structure* structure, JSC::SourceCode source, ScriptOptions options);
+    ~NodeVMScript();
 
     DECLARE_EXPORT_INFO;
     template<typename, JSC::SubspaceAccess mode> static JSC::GCClient::IsoSubspace* subspaceFor(JSC::VM& vm)
@@ -66,9 +103,18 @@ class NodeVMScript final : public JSC::JSDestructibleObject, public SigintReceiv
     static JSObject* createPrototype(VM& vm, JSGlobalObject* globalObject);
 
     JSC::ProgramExecutable* createExecutable();
+    // Parses the source once and pins the result; false (with `error` filled in) on a syntax error.
+    bool compile(JSC::JSGlobalObject*, JSC::ParserError& error);
     void cacheBytecode();
     JSC::JSUint8Array* getBytecodeBuffer();
 
+    // The hash is SourceCodeKey::hash(): the source plus the parse flags (code type, strictness,
+    // code generation mode) the block was produced under. A lookup under different flags is a
+    // miss, and the block JSC then produces replaces the pinned one.
+    JSC::UnlinkedProgramCodeBlock* pinnedUnlinkedCode(unsigned keyHash) const { return m_pinnedKeyHash == keyHash ? m_pinnedUnlinkedCode.get() : nullptr; }
+    void pinUnlinkedCode(unsigned keyHash, JSC::UnlinkedProgramCodeBlock*);
+    bool hasPinnedUnlinkedCode() const { return !!m_pinnedUnlinkedCode; }
+
     const JSC::SourceCode& source() const { return m_source; }
     WTF::Vector<uint8_t>& cachedData() { return m_options.cachedData; }
     JSC::ProgramExecutable* cachedExecutable() const { return m_cachedExecutable.get(); }
@@ -86,6 +132,8 @@ class NodeVMScript final : public JSC::JSDestructibleObject, public SigintReceiv
     RefPtr<JSC::CachedBytecode> m_cachedBytecode;
     JSC::WriteBarrier<JSC::JSUint8Array> m_cachedBytecodeBuffer;
     JSC::WriteBarrier<JSC::ProgramExecutable> m_cachedExecutable;
+    JSC::WriteBarrier<JSC::UnlinkedProgramCodeBlock> m_pinnedUnlinkedCode;
+    unsigned m_pinnedKeyHash = 0;
     ScriptOptions m_options;
     bool m_cachedDataProduced = false;
     bool m_sourceMapURLParsed = false;

test/js/node/vm/vm.test.ts

@@ -892,6 +892,93 @@ test("can't use bytecode from a different script", () => {
   expect(secondScript.runInThisContext()).toBe(4);
 });
 
+describe.concurrent("Script compiles its source once for every context it runs in", () => {
+  // Runs one Script in several fresh contexts, keeping everything each run produced alive, and
+  // reports how many UnlinkedProgramCodeBlock cells (one per parse of a program) appeared between
+  // the first run and the last. A Script that pins its parse adds none; a Script that goes back
+  // to JSC's CodeCache adds one per context as soon as the cache does not hold the entry, which
+  // BUN_JSC_useCodeCache=0 forces for every run.
+  const fixture = String.raw`
+    const { Script, createContext } = require("node:vm");
+    const { heapStats } = require("bun:jsc");
+    const runs = 6;
+    let body = "";
+    for (let i = 0; i < 50; i++) body += "function f" + i + "(a) { return a + " + i + "; }\n";
+    const source = "(function (exports) {\n" + body + "exports.sum = f0(1) + f49(1);\n})";
+    const programBlocks = () => {
+      Bun.gc(true);
+      return heapStats().objectTypeCounts.UnlinkedProgramCodeBlock ?? 0;
+    };
+    const options = process.env.VM_FIXTURE_CACHED_DATA ? { cachedData: new Script(source).createCachedData() } : {};
+    const script = new Script(source, options);
+    const keep = [];
+    let afterFirstRun = 0;
+    for (let i = 0; i < runs; i++) {
+      const wrapper = script.runInContext(createContext({}));
+      const exports = {};
+      wrapper(exports);
+      if (exports.sum !== 51) throw new Error("run " + i + " computed " + exports.sum);
+      keep.push(wrapper, exports);
+      if (i === 0) afterFirstRun = programBlocks();
+    }
+    console.log(JSON.stringify({
+      programBlocksAddedByLaterRuns: programBlocks() - afterFirstRun,
+      cachedDataRejected: script.cachedDataRejected,
+      cachedDataStillProducible: script.createCachedData().length > 0,
+    }));
+  `;
+
+  async function runFixture(extraEnv: Record<string, string>) {
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "-e", fixture],
+      env: { ...bunEnv, ...extraEnv },
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect(stderr).toBe("");
+    expect(exitCode).toBe(0);
+    return JSON.parse(stdout);
+  }
+
+  test("with the CodeCache", async () => {
+    expect(await runFixture({})).toEqual({
+      programBlocksAddedByLaterRuns: 0,
+      cachedDataStillProducible: true,
+    });
+  });
+
+  test("without the CodeCache", async () => {
+    expect(await runFixture({ BUN_JSC_useCodeCache: "0" })).toEqual({
+      programBlocksAddedByLaterRuns: 0,
+      cachedDataStillProducible: true,
+    });
+  });
+
+  test("from accepted cachedData, without the CodeCache", async () => {
+    expect(await runFixture({ BUN_JSC_useCodeCache: "0", VM_FIXTURE_CACHED_DATA: "1" })).toEqual({
+      programBlocksAddedByLaterRuns: 0,
+      cachedDataRejected: false,
+      cachedDataStillProducible: true,
+    });
+  });
+
+  test("each context still gets its own global declarations", () => {
+    const script = new Script(
+      "var counter = (typeof counter === 'number' ? counter : 0) + 1; function id() { return tag; } counter;",
+    );
+    const first = createContext({ tag: "first" });
+    const second = createContext({ tag: "second" });
+    expect(script.runInContext(first)).toBe(1);
+    expect(script.runInContext(second)).toBe(1);
+    expect(script.runInContext(first)).toBe(2);
+    expect(runInContext("id()", first)).toBe("first");
+    expect(runInContext("id()", second)).toBe("second");
+    expect(first.counter).toBe(2);
+    expect(second.counter).toBe(1);
+  });
+});
+
 describe("codeGeneration options", () => {
   test("disabling codeGeneration.strings should block eval and Function constructor", () => {
     const context = createContext(