wasm: poll VMTraps at loop back-edges so pure-Wasm loops can be terminated

Source/JavaScriptCore/llint/InPlaceInterpreter.asm

@@ -1173,6 +1173,22 @@ macro handleDebuggerTrapIfNeeded()
     addp 4 * MachineRegisterSize, sp
 end
 
+# Slow path for ipintOp(_loop) when m_trapAwareSoftStackLimit has been poisoned by
+# VMTraps::requestStop(). Services the pending async trap; on Termination the call
+# throws (operationCallMayThrow unwinds), otherwise PC/MC (restored by
+# operationCallMayThrow, still pointing at the loop opcode) are advanced past it and
+# the next instruction is dispatched.
+op(ipint_loop_check_vm_traps, macro ()
+    operationCallMayThrow(macro()
+        move cfr, a1
+        cCall2(_ipint_extern_handle_vm_traps_at_loop)
+    end)
+    loadb IPInt::InstructionLengthMetadata::length[MC], t0
+    advancePCByReg(t0)
+    advanceMCByReg(constexpr (sizeof(IPInt::InstructionLengthMetadata)))
+    nextIPIntInstruction()
+end)
+
 op(wasm_ipint_check_debugger_hook_and_throw_trap, macro ()
     handleDebuggerTrapIfNeeded()
     # r0 == 0 i.e. DebuggerTrapStatus::ResolvedByDebugger i.e. this was purely a debugger trap / breakpoint,

Source/JavaScriptCore/llint/InPlaceInterpreter64.asm

@@ -272,6 +272,13 @@ ipintOp(_loop, macro()
     # loop
     # We already validateOpcodeConfig in ipintLoopOSR.
     ipintLoopOSR(1)
+    # VMTraps poll: requestStop() sets m_trapAwareSoftStackLimit = UINTPTR_MAX, so a pure
+    # Wasm loop (no calls) observes termination/watchdog requests at each back-edge instead of
+    # only at the next function prologue. sp as the first comparand (not second) because ARM64
+    # cannot encode sp as Rm in SUBS.
+    bpaeq sp, JSWebAssemblyInstance::m_stackMirror + StackManager::Mirror::m_trapAwareSoftStackLimit[wasmInstance], .ipint_loop_no_trap
+    jmp _ipint_loop_check_vm_traps
+.ipint_loop_no_trap:
     loadb IPInt::InstructionLengthMetadata::length[MC], t0
     advancePCByReg(t0)
     advanceMCByReg(constexpr (sizeof(IPInt::InstructionLengthMetadata)))

Source/JavaScriptCore/runtime/StackManager.h

@@ -49,6 +49,11 @@ class StackManager {
             return OBJECT_OFFSETOF(Mirror, m_softStackLimit);
         }
 
+        static constexpr ptrdiff_t offsetOfTrapAwareSoftStackLimit()
+        {
+            return OBJECT_OFFSETOF(Mirror, m_trapAwareSoftStackLimit);
+        }
+
     private:
         Atomic<void*> m_trapAwareSoftStackLimit { nullptr };
         void* m_softStackLimit { nullptr };

Source/JavaScriptCore/wasm/WasmBBQJIT.cpp

@@ -3648,6 +3648,24 @@ void BBQJIT::emitLoopTierUpCheckAndOSREntryData(const ControlData& data, std::sp
     ASSERT(enclosingStack.size() >= args.size());
     auto enclosingWithoutArgs = enclosingStack.first(enclosingStack.size() - args.size());
     emitLoopTierUpCheckAndOSREntryData(result, enclosingWithoutArgs, loopIndex);
+
+    // VMTraps poll: requestStop() (e.g. VM::notifyNeedTermination) poisons
+    // m_trapAwareSoftStackLimit so a pure-Wasm loop observes termination/watchdog requests at
+    // each back-edge instead of never. The slow path runs under jit.probe so all live state is
+    // preserved for the resume case (non-termination async traps like NeedStopTheWorld).
+    JIT_COMMENT(m_jit, "Loop VMTraps poll");
+    static_assert(GPRInfo::nonPreservedNonArgumentGPR0 == wasmScratchGPR);
+    // sp as the left comparand: matches the offsetOfSoftStackLimit() checks above and avoids
+    // ARM64's extra mov when sp is Rm.
+    Jump needTrapHandling = m_jit.branchPtr(CCallHelpers::Below, MacroAssembler::stackPointerRegister, CCallHelpers::Address(GPRInfo::wasmContextInstancePointer, JSWebAssemblyInstance::offsetOfTrapAwareSoftStackLimit()));
+    MacroAssembler::Label trapResume = m_jit.label();
+    addLatePath(origin(), [needTrapHandling, trapResume](BBQJIT& bbq, CCallHelpers& jit) {
+        needTrapHandling.link(&jit);
+        jit.probe(tagCFunction<JITProbePtrTag>(operationWasmHandleTrapsAtLoop), nullptr);
+        bbq.recordJumpToThrowException(ExceptionType::Termination, jit.branchTestPtr(CCallHelpers::NonZero, GPRInfo::nonPreservedNonArgumentGPR0));
+        jit.jump().linkTo(trapResume, &jit);
+    });
+
     return { };
 }
 

Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.cpp

@@ -1467,6 +1467,22 @@ WASM_IPINT_EXTERN_CPP_DECL(check_stack_and_vm_traps, void* candidateNewStackPoin
     IPINT_THROW(Wasm::ExceptionType::StackOverflow);
 }
 
+// Reached from the loop back-edge when m_trapAwareSoftStackLimit has been poisoned by
+// VMTraps::requestStop(). Services whatever async trap is pending (NeedTermination /
+// NeedWatchdogCheck / NeedStopTheWorld / NeedDebuggerBreak); resumes the loop when it
+// wasn't a termination. Without this a pure-Wasm loop never observes VMTraps at all.
+WASM_IPINT_EXTERN_CPP_DECL(handle_vm_traps_at_loop, CallFrame* callFrame)
+{
+    UNUSED_PARAM(callFrame);
+    VM& vm = instance->vm();
+    if (vm.traps().handleTrapsIfNeeded()) {
+        if (vm.hasPendingTerminationException())
+            IPINT_THROW(Wasm::ExceptionType::Termination);
+        ASSERT(!vm.exceptionForInspection());
+    }
+    IPINT_END();
+}
+
 #if ENABLE(WEBASSEMBLY_DEBUGGER)
 static UNUSED_FUNCTION void displayWasmDebugState(JSWebAssemblyInstance* instance, Wasm::IPIntCallee* callee, CallFrame* callFrame, IPIntStackEntry* sp)
 {

Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.h

@@ -139,6 +139,7 @@ WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(memory_atomic_wait64, IPIntStackEntry*);
 WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(memory_atomic_notify, IPIntStackEntry*);
 
 WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(check_stack_and_vm_traps, void* candidateNewStackPointer, Wasm::IPIntCallee*, CallFrame*);
+WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(handle_vm_traps_at_loop, CallFrame*);
 WASM_IPINT_EXTERN_CPP_DECL(handle_debugger_trap_if_needed, CallFrame*, Register*);
 
 

Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp

@@ -4619,6 +4619,49 @@
     }
 
     m_currentBlock = body;
+
+    // VMTraps poll: requestStop() (e.g. VM::notifyNeedTermination) poisons
+    // m_trapAwareSoftStackLimit to UINTPTR_MAX so a pure-Wasm loop observes termination /
+    // watchdog requests at each back-edge instead of never. Match FTL compileCheckTraps: lower
+    // the poll as plain B3 IR (Load + Above + Branch) so the hot path carries no clobber set,
+    // and confine the probe patchpoint to a Rare block. fp stands in for sp (B3 has no SP value);
+    // fp > sp so fp < limit only under the poisoned-limit, which is the branch we want.
+    {
+        BasicBlock* slowPath = m_proc.addBlock();
+        BasicBlock* continuation = m_proc.addBlock();
+
+        Value* limit = m_currentBlock->appendNew<MemoryValue>(m_proc, Load, pointerType(), origin(), instanceValue(), safeCast<int32_t>(JSWebAssemblyInstance::offsetOfTrapAwareSoftStackLimit()));
+        m_currentBlock->appendNewControlValue(m_proc, B3::Branch, origin(),
+            m_currentBlock->appendNew<Value>(m_proc, Above, origin(), limit, framePointer()),
+            FrequentedBlock(slowPath, FrequencyClass::Rare), FrequentedBlock(continuation));
+        slowPath->addPredecessor(m_currentBlock);
+        continuation->addPredecessor(m_currentBlock);
+
+        m_currentBlock = slowPath;
+        // The probe preserves all live state for the resume case (non-termination async traps
+        // like NeedStopTheWorld); clobbers apply only to this rare block.
+        B3::PatchpointValue* handle = m_currentBlock->appendNew<B3::PatchpointValue>(m_proc, B3::Void, origin());
+        Effects effects = Effects::none();
+        effects.reads = B3::HeapRange::top();
+        effects.exitsSideways = true;
+        handle->effects = effects;
+        RegisterSet clobbers = RegisterSet::macroClobberedGPRs();
+        clobbers.add(GPRInfo::nonPreservedNonArgumentGPR0, IgnoreVectors);
+        handle->clobber(clobbers);
+        handle->append(instanceValue(), ValueRep::reg(GPRInfo::wasmContextInstancePointer));
+        handle->setGenerator([this, origin = this->origin()](CCallHelpers& jit, const B3::StackmapGenerationParams&) {
+            AllowMacroScratchRegisterUsage allowScratch(jit);
+            jit.probe(tagCFunction<JITProbePtrTag>(operationWasmHandleTrapsAtLoop), nullptr);
+            CCallHelpers::Jump resume = jit.branchTestPtr(CCallHelpers::Zero, GPRInfo::nonPreservedNonArgumentGPR0);
+            this->emitExceptionCheck(jit, origin, ExceptionType::Termination);
+            resume.link(&jit);
+        });
+        m_currentBlock->appendNewControlValue(m_proc, Jump, origin(), continuation);
+        continuation->addPredecessor(m_currentBlock);
+
+        m_currentBlock = continuation;
+    }
+
     return { };
 }
 

Source/JavaScriptCore/wasm/WasmOperations.cpp

@@ -1171,6 +1171,20 @@ JSC_DEFINE_NOEXCEPT_JIT_OPERATION(operationWasmLoopOSREnterBBQJIT, void, (Probe:
     context.gpr(GPRInfo::nonPreservedNonArgumentGPR0) = std::bit_cast<UCPURegister>(callee.loopEntrypoints()[loopIndex].taggedPtr());
 }
 
+// Reached from BBQ/OMG loop heads when m_trapAwareSoftStackLimit has been poisoned by
+// VMTraps::requestStop() (e.g. VM::notifyNeedTermination()). Services whatever async trap
+// is pending and reports whether a TerminationException was raised so the caller can unwind;
+// otherwise all registers are restored by the probe and the loop resumes. Without this a
+// pure-Wasm loop that never calls into JS (e.g. `loop { br 0 }`) never observes VMTraps and
+// cannot be preempted by worker.terminate() / the watchdog.
+JSC_DEFINE_NOEXCEPT_JIT_OPERATION(operationWasmHandleTrapsAtLoop, void, (Probe::Context& context))
+{
+    JSWebAssemblyInstance* instance = context.gpr<JSWebAssemblyInstance*>(GPRInfo::wasmContextInstancePointer);
+    VM& vm = instance->vm();
+    vm.traps().handleTrapsIfNeeded();
+    context.gpr(GPRInfo::nonPreservedNonArgumentGPR0) = static_cast<UCPURegister>(vm.hasPendingTerminationException());
+}
+
 #endif
 
 #if ENABLE(WEBASSEMBLY_BBQJIT)

Source/JavaScriptCore/wasm/WasmOperations.h

@@ -69,6 +69,7 @@ JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmTriggerTierUpNow, void, (CallFra
 #if ENABLE(WEBASSEMBLY_OMGJIT) || ENABLE(WEBASSEMBLY_BBQJIT)
 JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmTriggerOSREntryNow, void, (Probe::Context&));
 JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmLoopOSREnterBBQJIT, void, (Probe::Context&));
+JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmHandleTrapsAtLoop, void, (Probe::Context&));
 #endif
 #if ENABLE(WEBASSEMBLY_BBQJIT)
 JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmMaterializeBaselineData, void, (CallFrame*, JSWebAssemblyInstance*));

Source/JavaScriptCore/wasm/js/JSWebAssemblyInstance.h

@@ -177,6 +177,7 @@ class JSWebAssemblyInstance final : public JSNonFinalObject {
     using FunctionWrapperMap = UncheckedKeyHashMap<uint32_t, WriteBarrier<Unknown>, IntHash<uint32_t>, WTF::UnsignedWithZeroKeyHashTraits<uint32_t>>;
 
     static constexpr ptrdiff_t offsetOfSoftStackLimit() { return OBJECT_OFFSETOF(JSWebAssemblyInstance, m_stackMirror) + StackManager::Mirror::offsetOfSoftStackLimit(); }
+    static constexpr ptrdiff_t offsetOfTrapAwareSoftStackLimit() { return OBJECT_OFFSETOF(JSWebAssemblyInstance, m_stackMirror) + StackManager::Mirror::offsetOfTrapAwareSoftStackLimit(); }
 
     Wasm::Module& module() const { return m_module.get(); }
     SourceTaintedOrigin taintedness() const { return m_sourceProvider->sourceTaintedOrigin(); }