[JSC] Module graph instances: instantiate a linked module graph more than once in one global object

Source/JavaScriptCore/CMakeLists.txt

@@ -1594,6 +1594,8 @@ set(JavaScriptCore_PRIVATE_FRAMEWORK_HEADERS
     runtime/ModuleMap.h
     runtime/ModuleProgramExecutable.h
     runtime/ModuleProgramExecutableInlines.h
+    runtime/ModuleGraphInstance.h
+    runtime/ModuleGraphInstanceInlines.h
     runtime/ModuleRegistryEntry.h
     runtime/NarrowingNumberPredictionFuzzerAgent.h
     runtime/NativeCallee.h

Source/JavaScriptCore/Sources.txt

@@ -1009,6 +1009,7 @@ runtime/ModuleGraphLoadingState.cpp
 runtime/ModuleLoadingContext.cpp
 runtime/ModuleLoaderPayload.cpp
 runtime/ModuleProgramExecutable.cpp
+runtime/ModuleGraphInstance.cpp
 runtime/ModuleRegistryEntry.cpp
 runtime/NarrowingNumberPredictionFuzzerAgent.cpp
 runtime/NativeCallee.cpp

Source/JavaScriptCore/builtins/BuiltinNames.h

@@ -194,6 +194,7 @@ namespace JSC {
     macro(copyDataProperties) \
     macro(cloneObject) \
     macro(meta) \
+    macro(moduleGraphInstance) \
     macro(instanceFieldInitializer) \
     macro(privateBrand) \
     macro(privateClassBrand) \

Source/JavaScriptCore/bytecode/BytecodeList.rb

@@ -522,6 +522,7 @@
     },
     metadata: {
         resolveType: ResolveType, # offset 4
+        moduleImportSlot: unsigned, # ModuleVar: 1 + import slot ScopeOffset in the importing environment, 0 = none
         _0: { # offset 5
             localScopeDepth: unsigned,
             globalLexicalBindingEpoch: unsigned,

Source/JavaScriptCore/bytecode/CodeBlock.cpp

@@ -607,6 +607,7 @@ bool CodeBlock::finishCreation(VM& vm, ScriptExecutable* ownerExecutable, Unlink
 
             metadata.m_resolveType = op.type;
             metadata.m_localScopeDepth = op.depth;
+            metadata.m_moduleImportSlot = op.moduleImportSlot;
             if (op.lexicalEnvironment) {
                 if (op.type == ModuleVar) {
                     // Keep the linked module environment strongly referenced.

Source/JavaScriptCore/dfg/DFGByteCodeParser.cpp

@@ -10395,13 +10395,15 @@ void ByteCodeParser::parseBlock(unsigned limit)
 
             ResolveType resolveType;
             unsigned depth;
+            unsigned moduleImportSlot = 0;
             JSScope* constantScope = nullptr;
             JSCell* lexicalEnvironment = nullptr;
             SymbolTable* symbolTable = nullptr;
             {
                 ConcurrentJSLocker locker(m_inlineStackTop->m_profiledBlock->m_lock);
                 resolveType = metadata.m_resolveType;
                 depth = metadata.m_localScopeDepth;
+                moduleImportSlot = metadata.m_moduleImportSlot;
                 switch (resolveType) {
                 case GlobalProperty:
                 case GlobalVar:
@@ -10455,11 +10457,46 @@ void ByteCodeParser::parseBlock(unsigned limit)
                 break;
             }
             case ModuleVar: {
-                // Module environment is already strongly referenced by the CodeBlock.
-                set(bytecode.m_dst, weakJSConstant(lexicalEnvironment));
-                // BytecodeUseDef reports m_scope as a use regardless of resolve type,
-                // so we need to keep it OSR-available even though LLInt won't read it.
-                addToGraph(Phantom, get(bytecode.m_scope));
+                Node* localBase = get(bytecode.m_scope);
+                addToGraph(Phantom, localBase);
+                if (!moduleImportSlot) {
+                    // No import slot (module graph instances off): the linked
+                    // exporter environment is the only one there is.
+                    set(bytecode.m_dst, weakJSConstant(lexicalEnvironment));
+                    break;
+                }
+                // Module graph instances: the exporter environment is the importing
+                // environment's import slot — a closure-variable-like load. It folds
+                // to the linked exporter environment while that is the exporting
+                // module's only environment (its symbol table's singleton; the
+                // watchpoint fires when an instance creates a second one), and to
+                // the slot's value when the importing environment is a constant here.
+                ScopeOffset slot(moduleImportSlot - 1);
+                if (lexicalEnvironment) {
+                    SymbolTable* exporterSymbolTable = uncheckedDowncast<JSSymbolTableObject>(lexicalEnvironment)->symbolTable();
+                    if (exporterSymbolTable->singleton().inferredValue() == lexicalEnvironment) {
+                        m_graph.watchpoints().addLazily(m_graph, exporterSymbolTable);
+                        set(bytecode.m_dst, weakJSConstant(lexicalEnvironment));
+                        break;
+                    }
+                }
+                if (JSScope* resolvedScope = localBase->dynamicCastConstant<JSScope*>()) {
+                    for (unsigned n = depth; n--;)
+                        resolvedScope = resolvedScope->next();
+                    if (auto* importer = dynamicDowncast<JSModuleEnvironment>(resolvedScope)) {
+                        // Import slots sit past the symbol table's scope size; index the storage directly.
+                        if (JSValue exporter = importer->variables()[slot.offset()].get()) {
+                            set(bytecode.m_dst, weakJSConstant(exporter.asCell()));
+                            break;
+                        }
+                    }
+                }
+                for (unsigned n = depth; n--;)
+                    localBase = addToGraph(SkipScope, localBase);
+                Node* exporter = addToGraph(GetClosureVar, OpInfo(slot.offset()), OpInfo(SpecObjectOther), localBase);
+                addToGraph(CheckNotEmpty, exporter);
+                addToGraph(Check, Edge(exporter, CellUse));
+                set(bytecode.m_dst, exporter);
                 break;
             }
             case ResolvedClosureVar:

Source/JavaScriptCore/heap/Heap.cpp

@@ -76,6 +76,7 @@
 #include "MarkedSpaceInlines.h"
 #include "MarkingConstraintSet.h"
 #include "MegamorphicCache.h"
+#include "ModuleGraphInstance.h"
 #include "ModuleLoadingContext.h"
 #include "ModuleProgramExecutable.h"
 #include "ModuleRegistryEntry.h"

Source/JavaScriptCore/heap/Heap.h

@@ -274,6 +274,8 @@ class Heap;
     v(jsModuleRecordSpace, jsModuleRecordHeapCellType, JSModuleRecord) \
     v(moduleRegistryEntrySpace, destructibleCellHeapCellType, ModuleRegistryEntry) \
     v(moduleLoadingContextSpace, destructibleCellHeapCellType, ModuleLoadingContext) \
+    v(moduleGraphInstanceSpace, destructibleObjectHeapCellType, ModuleGraphInstance) \
+    v(moduleRecordInstanceSpace, destructibleCellHeapCellType, ModuleRecordInstance) \
     v(sentinelSpace, cellHeapCellType, JSSentinel) \
     v(syntheticModuleRecordSpace, syntheticModuleRecordHeapCellType, SyntheticModuleRecord) \
     v(jsMicrotaskDispatcherSpace, destructibleCellHeapCellType, JSMicrotaskDispatcher) \

Source/JavaScriptCore/interpreter/CallFrame.cpp

@@ -31,6 +31,7 @@
 #include "ExecutableAllocator.h"
 #include "InlineCallFrame.h"
 #include "JSCInlines.h"
+#include "JSCallee.h"
 #include "JSWebAssemblyInstance.h"
 #include "JSWebAssemblyModule.h"
 #include "LLIntPCRanges.h"
@@ -194,6 +195,42 @@ SUPPRESS_ASAN CallFrame* CallFrame::unsafeCallerFrame(EntryFrame*& currEntryFram
     return static_cast<CallFrame*>(unsafeCallerFrameOrEntryFrame());
 }
 
+
+JSScope* CallFrame::callerScope(VM& vm)
+{
+    RELEASE_ASSERT(callee().isCell());
+    JSScope* found = nullptr;
+    bool haveSkippedFirstFrame = false;
+    StackVisitor::visit(this, vm, [&](StackVisitor& visitor) {
+        if (!std::exchange(haveSkippedFirstFrame, true))
+            return IterationStatus::Continue;
+        switch (visitor->codeType()) {
+        case StackVisitor::Frame::CodeType::Native:
+        case StackVisitor::Frame::CodeType::Wasm:
+        // Eval code runs in its caller's scope, but its callee is the global
+        // object's shared eval callee whose scope is only set while the eval is
+        // being entered: attribute to the frame that called eval instead.
+        case StackVisitor::Frame::CodeType::Eval:
+            return IterationStatus::Continue;
+        case StackVisitor::Frame::CodeType::Function:
+        case StackVisitor::Frame::CodeType::Module:
+        case StackVisitor::Frame::CodeType::Global:
+            break;
+        }
+        // The callee carries the scope its code was created with at every tier
+        // (the scope register is not materialized in optimized frames).
+        JSCell* calleeCell = visitor->callee().asCell();
+        if (auto* function = dynamicDowncast<JSFunction>(calleeCell)) {
+            if (!function->isHostFunction() && function->jsExecutable()->isPrivateBuiltinFunction())
+                return IterationStatus::Continue;
+            found = function->scope();
+        } else if (auto* callee = dynamicDowncast<JSCallee>(calleeCell))
+            found = callee->scope();
+        return IterationStatus::Done;
+    });
+    return found;
+}
+
 SourceOrigin CallFrame::callerSourceOrigin(VM& vm)
 {
     RELEASE_ASSERT(callee().isCell());

Source/JavaScriptCore/interpreter/CallFrame.h

@@ -228,6 +228,11 @@ using JSInstruction = BaseInstruction<JSOpcodeTraits>;
         JS_EXPORT_PRIVATE CallFrame* callerFrame(EntryFrame*&) const;
 
         JS_EXPORT_PRIVATE SourceOrigin callerSourceOrigin(VM&);
+        // Module graph instances: the scope the calling JS code closes over (the
+        // module environment for module code; the callee's scope chain for
+        // functions), or null for native callers. Lets import() load into the
+        // caller's graph instance.
+        JS_EXPORT_PRIVATE JSScope* callerScope(VM&);
 
         static constexpr ptrdiff_t callerFrameOffset() { return OBJECT_OFFSETOF(CallerFrameAndPC, callerFrame); }
 

Source/JavaScriptCore/interpreter/Interpreter.cpp

@@ -73,6 +73,7 @@
 #include "LLIntThunks.h"
 #include "LiteralParser.h"
 #include "MicrotaskCall.h"
+#include "ModuleGraphInstance.h"
 #include "ModuleProgramCodeBlock.h"
 #include "NativeCallee.h"
 #include "ProgramCodeBlock.h"
@@ -1708,6 +1709,11 @@ JSValue Interpreter::executeEval(EvalExecutable* eval, JSValue thisValue, JSScop
 }
 
 JSValue Interpreter::executeModuleProgram(JSModuleRecord* record, ModuleProgramExecutable* executable, JSGlobalObject* lexicalGlobalObject, JSModuleEnvironment* scope, JSValue sentValue, JSValue resumeMode)
+{
+    return executeModuleProgram(record, record, executable, lexicalGlobalObject, scope, sentValue, resumeMode);
+}
+
+JSValue Interpreter::executeModuleProgram(JSModuleRecord* record, JSObject* generatorState, ModuleProgramExecutable* executable, JSGlobalObject* lexicalGlobalObject, JSModuleEnvironment* scope, JSValue sentValue, JSValue resumeMode)
 {
     VM& vm = this->vm();
     auto throwScope = DECLARE_THROW_SCOPE(vm);
@@ -1737,9 +1743,15 @@ JSValue Interpreter::executeModuleProgram(JSModuleRecord* record, ModuleProgramE
     RefPtr<JSC::JITCode> jitCode;
 
     ProtoCallFrame protoCallFrame;
+    auto stateField = [&]() -> WriteBarrier<Unknown>& {
+        if (auto* recordInstance = dynamicDowncast<ModuleRecordInstance>(generatorState))
+            return recordInstance->internalField(ModuleRecordInstance::Field::State);
+        ASSERT(generatorState == record);
+        return record->internalField(JSModuleRecord::Field::State);
+    };
     EncodedJSValue args[numberOfArguments] = {
-        JSValue::encode(record),
-        JSValue::encode(record->internalField(JSModuleRecord::Field::State).get()),
+        JSValue::encode(generatorState),
+        JSValue::encode(stateField().get()),
         JSValue::encode(sentValue),
         JSValue::encode(resumeMode),
         JSValue::encode(scope),
@@ -1767,7 +1779,7 @@ JSValue Interpreter::executeModuleProgram(JSModuleRecord* record, ModuleProgramE
             protoCallFrame.init(codeBlock, globalObject, callee, jsUndefined(), nullptr, numberOfArguments + 1, args);
         }
 
-        record->internalField(JSModuleRecord::Field::State).set(vm, record, jsNumber(static_cast<int>(JSModuleRecord::State::Executing)));
+        stateField().set(vm, generatorState, jsNumber(static_cast<int>(JSModuleRecord::State::Executing)));
     }
 
     // Execute the code:

Source/JavaScriptCore/interpreter/Interpreter.h

@@ -161,6 +161,9 @@ using JSOrWasmInstruction = Variant<const JSInstruction*, uintptr_t /* IPIntOffs
         JSValue executeProgram(const SourceCode&, JSGlobalObject*, JSObject* thisObj);
 #endif
         JSValue executeModuleProgram(JSModuleRecord*, ModuleProgramExecutable*, JSGlobalObject*, JSModuleEnvironment*, JSValue sentValue, JSValue resumeMode);
+        // Module graph instances (prototype): as above with an explicit holder for the
+        // top-level-await generator state (State, Frame internal fields) instead of the record.
+        JSValue executeModuleProgram(JSModuleRecord*, JSObject* generatorState, ModuleProgramExecutable*, JSGlobalObject*, JSModuleEnvironment*, JSValue sentValue, JSValue resumeMode);
         JSValue executeCall(JSObject* function, const CallData&, JSValue thisValue, JSCell* context, const ArgList&);
         JSObject* executeConstruct(JSObject* function, const CallData&, const ArgList&, JSValue newTarget);
         JSValue executeEval(EvalExecutable*, JSValue thisValue, JSScope*);

Source/JavaScriptCore/jit/JITOperations.cpp

@@ -65,6 +65,8 @@ WTF_ALLOW_UNSAFE_BUFFER_USAGE_BEGIN
 #include "JSGlobalObjectFunctions.h"
 #include "JSLexicalEnvironmentInlines.h"
 #include "JSMapIterator.h"
+#include "JSModuleEnvironment.h"
+#include "AbstractModuleRecord.h"
 #include "JSMicrotask.h"
 #include "JSPromise.h"
 #include "JSRemoteFunction.h"
@@ -4593,16 +4595,20 @@ JSC_DEFINE_JIT_OPERATION(operationResolveScopeForBaseline, EncodedJSValue, (JSGl
     auto bytecode = pc->as<OpResolveScope>();
     const Identifier& ident = codeBlock->identifier(bytecode.m_var);
     JSScope* environment = callFrame->uncheckedR(bytecode.m_scope).Register::scope();
+    auto& metadata = bytecode.metadata(codeBlock);
+
+    if (metadata.m_resolveType == ModuleVar) {
+        JSObject* result = JSModuleEnvironment::resolveModuleVarScope(globalObject, environment, metadata.m_localScopeDepth, uncheckedDowncast<JSModuleEnvironment>(metadata.m_lexicalEnvironment.get()));
+        OPERATION_RETURN_IF_EXCEPTION(scope, encodedJSValue());
+        OPERATION_RETURN(scope, JSValue::encode(result));
+    }
+
     JSObject* resolvedScope = JSScope::resolve(globalObject, environment, ident);
     // Proxy can throw an error here, e.g. Proxy in with statement's @unscopables.
     OPERATION_RETURN_IF_EXCEPTION(scope, encodedJSValue());
 
-    auto& metadata = bytecode.metadata(codeBlock);
     ResolveType resolveType = metadata.m_resolveType;
 
-    // ModuleVar does not keep the scope register value alive in DFG.
-    ASSERT(resolveType != ModuleVar);
-
     switch (resolveType) {
     case GlobalProperty:
     case GlobalPropertyWithVarInjectionChecks:

Source/JavaScriptCore/jit/JITPropertyAccess.cpp

@@ -896,9 +896,34 @@ void JIT::emit_op_resolve_scope(const JSInstruction* currentInstruction)
     // If we profile certain resolve types, we're guaranteed all linked code will have the same
     // resolve type.
 
-    if (profiledResolveType == ModuleVar)
-        loadPtrFromMetadata(bytecode, Metadata::offsetOfLexicalEnvironment(), returnValueGPR);
-    else if (profiledResolveType == ClosureVar) {
+    if (profiledResolveType == ModuleVar) {
+        if (!Options::useModuleGraphInstances()) {
+            // The exporter environment is a link-time constant.
+            loadPtrFromMetadata(bytecode, Metadata::offsetOfLexicalEnvironment(), returnValueGPR);
+        } else {
+            // Module graph instances: the import slot index (like the exporter
+            // environment) is a per-CodeBlock link-time value, and baseline code
+            // is shared between CodeBlocks of one UnlinkedCodeBlock, so read it
+            // from metadata. Slot 0: constant environment. Otherwise walk to the
+            // importing module environment (a lexical depth) and load its import
+            // slot; an unfilled slot goes to the slow path.
+            load32FromMetadata(bytecode, Metadata::offsetOfModuleImportSlot(), scratch1GPR);
+            Jump hasImportSlot = branchTest32(NonZero, scratch1GPR);
+            loadPtrFromMetadata(bytecode, Metadata::offsetOfLexicalEnvironment(), returnValueGPR);
+            Jump done = jump();
+            hasImportSlot.link(this);
+            emitGetVirtualRegister(scope, scopeGPR);
+            static_assert(scopeGPR == returnValueGPR);
+            unsigned localScopeDepth = bytecode.metadata(m_profiledCodeBlock).m_localScopeDepth;
+            for (unsigned index = 0; index < localScopeDepth; ++index)
+                loadPtr(Address(returnValueGPR, JSScope::offsetOfNext()), returnValueGPR);
+            static_assert(sizeof(WriteBarrier<Unknown>) == 8);
+            sub32(TrustedImm32(1), scratch1GPR);
+            load64(BaseIndex(returnValueGPR, scratch1GPR, TimesEight, JSLexicalEnvironment::offsetOfVariables()), returnValueGPR);
+            addSlowCase(branchIfEmpty(returnValueGPR));
+            done.link(this);
+        }
+    } else if (profiledResolveType == ClosureVar) {
         emitGetVirtualRegister(scope, scopeGPR);
         static_assert(scopeGPR == returnValueGPR);
         unsigned localScopeDepth = bytecode.metadata(m_profiledCodeBlock).m_localScopeDepth;
@@ -1093,10 +1118,10 @@ MacroAssemblerCodeRef<JITThunkPtrTag> JIT::generateOpResolveScopeThunk(VM& vm)
             emitResolveClosure(needsVarInjectionChecks(resolveType));
             break;
         case Dynamic:
+        case ModuleVar:
             slowCase.append(jit.jump());
             break;
         case ResolvedClosureVar:
-        case ModuleVar:
         case UnresolvedProperty:
         case UnresolvedPropertyWithVarInjectionChecks:
             RELEASE_ASSERT_NOT_REACHED();

Source/JavaScriptCore/llint/LowLevelInterpreter64.asm

@@ -2849,6 +2849,25 @@ llintOpWithMetadata(op_resolve_scope, OpResolveScope, macro (size, get, dispatch
 
 .rModuleVar:
     bineq t0, ModuleVar, .rGlobalPropertyWithVarInjectionChecks
+    # With module graph instances the exporting environment is per instance: it
+    # is the importing environment's import slot (walk localScopeDepth to the
+    # importing module environment, load the slot; empty = not filled yet).
+    loadi OpResolveScope::Metadata::m_moduleImportSlot[t5], t1
+    btiz t1, .rModuleVarConstant
+    loadi OpResolveScope::Metadata::m_localScopeDepth[t5], t2
+    get(m_scope, t0)
+    loadq [cfr, t0, 8], t0
+    btiz t2, .rModuleVarLoad
+.rModuleVarWalk:
+    loadp JSScope::m_next[t0], t0
+    subi 1, t2
+    btinz t2, .rModuleVarWalk
+.rModuleVarLoad:
+    subi 1, t1
+    loadq JSLexicalEnvironment_variables[t0, t1, 8], t0
+    btqz t0, .rDynamic
+    return(t0)
+.rModuleVarConstant:
     returnConstantScope()
 
 .rGlobalPropertyWithVarInjectionChecks: