Bug 1360254 - Baldr: remove JSContext::wasmActivationStack (r=bbouvier)

MozReview-Commit-ID: Ftzs7mTUzWN

--HG--
extra : rebase_source : b4670defcde377465b9d51681bb51a60a942b214

js/src/jit/MacroAssembler.cpp

@@ -3092,13 +3092,6 @@ MacroAssembler::wasmAssertNonExitInvariants(Register activation)
 
 //}}} check_macroassembler_style
 
-void
-MacroAssembler::loadWasmActivationFromTls(Register dest)
-{
-    loadPtr(Address(WasmTlsReg, offsetof(wasm::TlsData, cx)), dest);
-    loadPtr(Address(dest, JSContext::offsetOfWasmActivation()), dest);
-}
-
 void
 MacroAssembler::loadWasmTlsRegFromFrame(Register dest)
 {

js/src/jit/MacroAssembler.h

@@ -1539,7 +1539,6 @@ class MacroAssembler : public MacroAssemblerSpecific
 
     void guardGroupHasUnanalyzedNewScript(Register group, Register scratch, Label* fail);
 
-    void loadWasmActivationFromTls(Register dest);
     void loadWasmTlsRegFromFrame(Register dest = WasmTlsReg);
 
     template<typename T>

js/src/jit/arm/Simulator-arm.cpp

@@ -1558,7 +1558,7 @@ Simulator::handleWasmInterrupt()
     void* pc = (void*)get_pc();
     uint8_t* fp = (uint8_t*)get_register(r11);
 
-    WasmActivation* activation = JSContext::innermostWasmActivation();
+    WasmActivation* activation = wasm::MaybeActiveActivation(cx_);
     const wasm::Code* code = activation->compartment()->wasm.lookupCode(pc);
     if (!code || !code->segment().containsFunctionPC(pc))
         return;
@@ -1580,7 +1580,7 @@ Simulator::handleWasmInterrupt()
 bool
 Simulator::handleWasmFault(int32_t addr, unsigned numBytes)
 {
-    WasmActivation* act = cx_->wasmActivationStack();
+    WasmActivation* act = wasm::MaybeActiveActivation(cx_);
     if (!act)
         return false;
 
@@ -1614,7 +1614,7 @@ Simulator::readQ(int32_t addr, SimInstruction* instr, UnalignedPolicy f)
     }
 
     // See the comments below in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         uint64_t value;
         memcpy(&value, ptr, sizeof(value));
@@ -1638,7 +1638,7 @@ Simulator::writeQ(int32_t addr, uint64_t value, SimInstruction* instr, Unaligned
     }
 
     // See the comments below in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         memcpy(ptr, &value, sizeof(value));
         return;
@@ -1663,7 +1663,7 @@ Simulator::readW(int32_t addr, SimInstruction* instr, UnalignedPolicy f)
     // do the right thing. Making this simulator properly emulate the behavior
     // of raising a signal is complex, so as a special-case, when in wasm code,
     // we just do the right thing.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         int value;
         memcpy(&value, ptr, sizeof(value));
@@ -1687,7 +1687,7 @@ Simulator::writeW(int32_t addr, int value, SimInstruction* instr, UnalignedPolic
     }
 
     // See the comments above in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         memcpy(ptr, &value, sizeof(value));
         return;
@@ -1763,7 +1763,7 @@ Simulator::readHU(int32_t addr, SimInstruction* instr)
     }
 
     // See comments above in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         uint16_t value;
         memcpy(&value, ptr, sizeof(value));
@@ -1787,7 +1787,7 @@ Simulator::readH(int32_t addr, SimInstruction* instr)
     }
 
     // See comments above in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         int16_t value;
         memcpy(&value, ptr, sizeof(value));
@@ -1812,7 +1812,7 @@ Simulator::writeH(int32_t addr, uint16_t value, SimInstruction* instr)
     }
 
     // See the comments above in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         memcpy(ptr, &value, sizeof(value));
         return;
@@ -1835,7 +1835,7 @@ Simulator::writeH(int32_t addr, int16_t value, SimInstruction* instr)
     }
 
     // See the comments above in readW.
-    if (FixupFault() && wasm::IsPCInWasmCode(reinterpret_cast<void *>(get_pc()))) {
+    if (FixupFault() && wasm::InCompiledCode(reinterpret_cast<void*>(get_pc()))) {
         char* ptr = reinterpret_cast<char*>(addr);
         memcpy(ptr, &value, sizeof(value));
         return;

js/src/jit/arm64/vixl/Debugger-vixl.cpp

@@ -545,8 +545,8 @@ const char* RegisterToken::kWAliases[kNumberOfRegisters][kMaxAliasNumber] = {
 };
 
 
-Debugger::Debugger(Decoder* decoder, FILE* stream)
-    : Simulator(decoder, stream),
+Debugger::Debugger(JSContext* cx, Decoder* decoder, FILE* stream)
+    : Simulator(cx, decoder, stream),
       debug_parameters_(DBG_INACTIVE),
       pending_request_(false),
       steps_(0),

js/src/jit/arm64/vixl/Debugger-vixl.h

@@ -54,7 +54,7 @@ class FormatToken;
 
 class Debugger : public Simulator {
  public:
-  explicit Debugger(Decoder* decoder, FILE* stream = stdout);
+  explicit Debugger(JSContext* cx, Decoder* decoder, FILE* stream = stdout);
   ~Debugger();
 
   virtual void Run();

js/src/jit/arm64/vixl/MozSimulator-vixl.cpp

@@ -43,8 +43,9 @@ using mozilla::DebugOnly;
 using js::jit::ABIFunctionType;
 using js::jit::SimulatorProcess;
 
-Simulator::Simulator(Decoder* decoder, FILE* stream)
-  : stream_(nullptr)
+Simulator::Simulator(JSContext* cx, Decoder* decoder, FILE* stream)
+  : cx_(cx)
+  , stream_(nullptr)
   , print_disasm_(nullptr)
   , instrumentation_(nullptr)
   , stack_(nullptr)
@@ -168,9 +169,9 @@ Simulator* Simulator::Create(JSContext* cx) {
   // FIXME: Note that it can't be stored in the SimulatorRuntime due to lifetime conflicts.
   Simulator *sim;
   if (getenv("USE_DEBUGGER") != nullptr)
-    sim = js_new<Debugger>(decoder, stdout);
+    sim = js_new<Debugger>(cx, decoder, stdout);
   else
-    sim = js_new<Simulator>(decoder, stdout);
+    sim = js_new<Simulator>(cx, decoder, stdout);
 
   // Check if Simulator:init ran out of memory.
   if (sim && sim->oom()) {
@@ -242,7 +243,7 @@ void Simulator::handle_wasm_interrupt() {
   void* pc = (void*)get_pc();
   uint8_t* fp = (uint8_t*)xreg(30);
 
-  js::WasmActivation* activation = JSContext::innermostWasmActivation();
+  js::WasmActivation* activation = js::wasm::MaybeActiveActivation(cx_);
   const js::wasm::Code* code = activation->compartment()->wasm.lookupCode(pc);
   if (!code || !code->segment().containsFunctionPC(pc))
     return;

js/src/jit/arm64/vixl/Simulator-vixl.h

@@ -697,7 +697,7 @@ class Redirection;
 
 class Simulator : public DecoderVisitor {
  public:
-  explicit Simulator(Decoder* decoder, FILE* stream = stdout);
+  explicit Simulator(JSContext* cx, Decoder* decoder, FILE* stream = stdout);
   ~Simulator();
 
   // Moz changes.
@@ -2510,6 +2510,8 @@ class Simulator : public DecoderVisitor {
 
   // Processor state ---------------------------------------
 
+  JSContext* const cx_;
+
   // Simulated monitors for exclusive access instructions.
   SimExclusiveLocalMonitor local_monitor_;
   SimExclusiveGlobalMonitor global_monitor_;

js/src/jit/none/MacroAssembler-none.h

@@ -413,7 +413,6 @@ class MacroAssemblerNone : public Assembler
     void buildFakeExitFrame(Register, uint32_t*) { MOZ_CRASH(); }
     bool buildOOLFakeExitFrame(void*) { MOZ_CRASH(); }
     void loadWasmGlobalPtr(uint32_t, Register) { MOZ_CRASH(); }
-    void loadWasmActivationFromTls(Register) { MOZ_CRASH(); }
     void loadWasmPinnedRegsFromTls() { MOZ_CRASH(); }
 
     void setPrinter(Sprinter*) { MOZ_CRASH(); }

js/src/jscntxt.cpp

@@ -1133,7 +1133,6 @@ JSContext::JSContext(JSRuntime* runtime, const JS::ContextOptions& options)
     jitActivation(nullptr),
     activation_(nullptr),
     profilingActivation_(nullptr),
-    wasmActivationStack_(nullptr),
     nativeStackBase(GetNativeStackBase()),
     entryMonitor(nullptr),
     noExecuteDebuggerTop(nullptr),

js/src/jscntxt.h

@@ -293,9 +293,6 @@ struct JSContext : public JS::RootingContext,
     static size_t offsetOfActivation() {
         return offsetof(JSContext, activation_);
     }
-    static size_t offsetOfWasmActivation() {
-        return offsetof(JSContext, wasmActivationStack_);
-    }
     static size_t offsetOfProfilingActivation() {
         return offsetof(JSContext, profilingActivation_);
      }
@@ -369,16 +366,6 @@ struct JSContext : public JS::RootingContext,
     js::Activation* volatile profilingActivation_;
 
   public:
-    /* See WasmActivation comment. */
-    js::WasmActivation* volatile wasmActivationStack_;
-
-    js::WasmActivation* wasmActivationStack() const {
-        return wasmActivationStack_;
-    }
-    static js::WasmActivation* innermostWasmActivation() {
-        return js::TlsContext.get()->wasmActivationStack_;
-    }
-
     js::Activation* activation() const {
         return activation_;
     }

js/src/vm/Stack.cpp

@@ -1647,9 +1647,6 @@ WasmActivation::WasmActivation(JSContext* cx)
     exitFP_(nullptr),
     exitReason_(wasm::ExitReason::Fixed::None)
 {
-    prevWasm_ = cx->wasmActivationStack_;
-    cx->wasmActivationStack_ = this;
-
     // Now that the WasmActivation is fully initialized, make it visible to
     // asynchronous profiling.
     registerProfiling();
@@ -1663,9 +1660,6 @@ WasmActivation::~WasmActivation()
     MOZ_ASSERT(!interrupted());
     MOZ_ASSERT(exitFP_ == nullptr);
     MOZ_ASSERT(exitReason_.isNone());
-
-    MOZ_ASSERT(cx_->wasmActivationStack_ == this);
-    cx_->wasmActivationStack_ = prevWasm_;
 }
 
 void

js/src/vm/Stack.h

@@ -1360,6 +1360,9 @@ class Activation
         return hideScriptedCallerCount_ > 0;
     }
 
+    static size_t offsetOfPrev() {
+        return offsetof(Activation, prev_);
+    }
     static size_t offsetOfPrevProfiling() {
         return offsetof(Activation, prevProfiling_);
     }
@@ -1721,18 +1724,12 @@ class InterpreterFrameIterator
     }
 };
 
-// A WasmActivation is part of two activation linked lists:
-//  - the normal Activation list used by FrameIter
-//  - a list of only WasmActivations that is signal-safe since it is accessed
-//    from the profiler at arbitrary points
-//
 // An eventual goal is to remove WasmActivation and to run asm code in a
 // JitActivation interleaved with Ion/Baseline jit code. This would allow
 // efficient calls back and forth but requires that we can walk the stack for
 // all kinds of jit code.
 class WasmActivation : public Activation
 {
-    WasmActivation* prevWasm_;
     wasm::Frame* exitFP_;
     wasm::ExitReason exitReason_;
 
@@ -1740,8 +1737,6 @@ class WasmActivation : public Activation
     explicit WasmActivation(JSContext* cx);
     ~WasmActivation();
 
-    WasmActivation* prevWasm() const { return prevWasm_; }
-
     bool isProfiling() const {
         return true;
     }

js/src/wasm/WasmBuiltins.cpp

@@ -61,10 +61,22 @@ __aeabi_uidivmod(int, int);
 }
 #endif
 
+// This utility function can only be called for builtins that are called
+// directly from wasm code. Note that WasmCall pushes both an outer
+// WasmActivation and an inner JitActivation that becomes active when calling
+// JIT code.
+static WasmActivation*
+CallingActivation()
+{
+    Activation* act = TlsContext.get()->activation();
+    MOZ_ASSERT(!act->asJit()->isActive(), "WasmCall pushes an inactive JitActivation");
+    return act->prev()->asWasm();
+}
+
 static void*
 WasmHandleExecutionInterrupt()
 {
-    WasmActivation* activation = JSContext::innermostWasmActivation();
+    WasmActivation* activation = CallingActivation();
     MOZ_ASSERT(activation->interrupted());
 
     if (!CheckForInterrupt(activation->cx())) {
@@ -86,7 +98,7 @@ WasmHandleExecutionInterrupt()
 static bool
 WasmHandleDebugTrap()
 {
-    WasmActivation* activation = JSContext::innermostWasmActivation();
+    WasmActivation* activation = CallingActivation();
     MOZ_ASSERT(activation);
     JSContext* cx = activation->cx();
 
@@ -154,10 +166,8 @@ WasmHandleDebugTrap()
 static void*
 WasmHandleThrow()
 {
-    JSContext* cx = TlsContext.get();
-
-    WasmActivation* activation = cx->wasmActivationStack();
-    MOZ_ASSERT(activation);
+    WasmActivation* activation = CallingActivation();
+    JSContext* cx = activation->cx();
 
     // FrameIterator iterates down wasm frames in the activation starting at
     // WasmActivation::exitFP. Pass Unwind::True to pop WasmActivation::exitFP
@@ -215,7 +225,7 @@ WasmHandleThrow()
 static void
 WasmReportTrap(int32_t trapIndex)
 {
-    JSContext* cx = JSContext::innermostWasmActivation()->cx();
+    JSContext* cx = TlsContext.get();
 
     MOZ_ASSERT(trapIndex < int32_t(Trap::Limit) && trapIndex >= 0);
     Trap trap = Trap(trapIndex);
@@ -259,21 +269,21 @@ WasmReportTrap(int32_t trapIndex)
 static void
 WasmReportOutOfBounds()
 {
-    JSContext* cx = JSContext::innermostWasmActivation()->cx();
+    JSContext* cx = TlsContext.get();
     JS_ReportErrorNumberASCII(cx, GetErrorMessage, nullptr, JSMSG_WASM_OUT_OF_BOUNDS);
 }
 
 static void
 WasmReportUnalignedAccess()
 {
-    JSContext* cx = JSContext::innermostWasmActivation()->cx();
+    JSContext* cx = TlsContext.get();
     JS_ReportErrorNumberASCII(cx, GetErrorMessage, nullptr, JSMSG_WASM_UNALIGNED_ACCESS);
 }
 
 static int32_t
 CoerceInPlace_ToInt32(MutableHandleValue val)
 {
-    JSContext* cx = JSContext::innermostWasmActivation()->cx();
+    JSContext* cx = TlsContext.get();
 
     int32_t i32;
     if (!ToInt32(cx, val, &i32))
@@ -286,7 +296,7 @@ CoerceInPlace_ToInt32(MutableHandleValue val)
 static int32_t
 CoerceInPlace_ToNumber(MutableHandleValue val)
 {
-    JSContext* cx = JSContext::innermostWasmActivation()->cx();
+    JSContext* cx = TlsContext.get();
 
     double dbl;
     if (!ToNumber(cx, val, &dbl))

js/src/wasm/WasmFrameIterator.cpp

@@ -300,6 +300,16 @@ PushRetAddr(MacroAssembler& masm, unsigned entry)
 #endif
 }
 
+static void
+LoadActivation(MacroAssembler& masm, Register dest)
+{
+    // WasmCall pushes a WasmActivation and an inactive JitActivation. The
+    // JitActivation only becomes active when calling into JS from wasm.
+    masm.loadPtr(Address(WasmTlsReg, offsetof(wasm::TlsData, cx)), dest);
+    masm.loadPtr(Address(dest, JSContext::offsetOfActivation()), dest);
+    masm.loadPtr(Address(dest, Activation::offsetOfPrev()), dest);
+}
+
 static void
 GenerateCallablePrologue(MacroAssembler& masm, unsigned framePushed, ExitReason reason,
                          uint32_t* entry)
@@ -326,15 +336,13 @@ GenerateCallablePrologue(MacroAssembler& masm, unsigned framePushed, ExitReason
     }
 
     if (!reason.isNone()) {
+        // Native callers expect the native ABI, which assume that nonvolatile
+        // registers are preserved.
         Register scratch = ABINonArgReg0;
-
-        // Native callers expect the native ABI, which assume that non-saved
-        // registers are preserved. Explicitly preserve the scratch register
-        // in that case.
         if (reason.isNative() && !scratch.volatile_())
             masm.Push(scratch);
 
-        masm.loadWasmActivationFromTls(scratch);
+        LoadActivation(masm, scratch);
         masm.wasmAssertNonExitInvariants(scratch);
         Address exitReason(scratch, WasmActivation::offsetOfExitReason());
         masm.store32(Imm32(reason.raw()), exitReason);
@@ -357,13 +365,12 @@ GenerateCallableEpilogue(MacroAssembler& masm, unsigned framePushed, ExitReason
         masm.addToStackPtr(Imm32(framePushed));
 
     if (!reason.isNone()) {
-        Register scratch = ABINonArgReturnReg0;
-
         // See comment in GenerateCallablePrologue.
+        Register scratch = ABINonArgReturnReg0;
         if (reason.isNative() && !scratch.volatile_())
             masm.Push(scratch);
 
-        masm.loadWasmActivationFromTls(scratch);
+        LoadActivation(masm, scratch);
         Address exitFP(scratch, WasmActivation::offsetOfExitFP());
         masm.storePtr(ImmWord(0), exitFP);
         Address exitReason(scratch, WasmActivation::offsetOfExitReason());
@@ -963,3 +970,33 @@ wasm::LookupFaultingInstance(WasmActivation* activation, void* pc, void* fp)
     MOZ_RELEASE_ASSERT(&instance->code() == code);
     return instance;
 }
+
+WasmActivation*
+wasm::MaybeActiveActivation(JSContext* cx)
+{
+    // WasmCall pushes both an outer WasmActivation and an inner JitActivation
+    // that only becomes active when calling JIT code.
+    Activation* act = cx->activation();
+    while (act && act->isJit() && !act->asJit()->isActive())
+        act = act->prev();
+    if (!act || !act->isWasm())
+        return nullptr;
+    return act->asWasm();
+}
+
+bool
+wasm::InCompiledCode(void* pc)
+{
+    JSContext* cx = TlsContext.get();
+    if (!cx)
+        return false;
+
+    MOZ_RELEASE_ASSERT(!cx->handlingSegFault);
+
+    if (cx->compartment()->wasm.lookupCode(pc))
+        return true;
+
+    const CodeRange* codeRange;
+    uint8_t* codeBase;
+    return LookupBuiltinThunk(pc, &codeRange, &codeBase);
+}

js/src/wasm/WasmFrameIterator.h

@@ -187,6 +187,16 @@ TraceActivations(JSContext* cx, const CooperatingContext& target, JSTracer* trc)
 Instance*
 LookupFaultingInstance(WasmActivation* activation, void* pc, void* fp);
 
+// If the innermost (active) Activation is a WasmActivation, return it.
+
+WasmActivation*
+MaybeActiveActivation(JSContext* cx);
+
+// Return whether the given PC is in wasm code.
+
+bool
+InCompiledCode(void* pc);
+
 } // namespace wasm
 } // namespace js
 

js/src/wasm/WasmSignalHandlers.cpp

@@ -858,14 +858,14 @@ HandleFault(PEXCEPTION_POINTERS exception)
         return false;
     AutoSetHandlingSegFault handling(cx);
 
-    WasmActivation* activation = cx->wasmActivationStack();
-    if (!activation)
-        return false;
-
     const Code* code = activation->compartment()->wasm.lookupCode(pc);
     if (!code)
         return false;
 
+    WasmActivation* activation = MaybeActiveActivation(cx);
+    if (!activation)
+        return false;
+
     if (!code->segment().containsFunctionPC(pc)) {
         // On Windows, it is possible for InterruptRunningJitCode to execute
         // between a faulting heap access and the handling of the fault due
@@ -1012,7 +1012,7 @@ HandleMachException(JSContext* cx, const ExceptionRequest& request)
     if (request.body.exception != EXC_BAD_ACCESS || request.body.codeCnt != 2)
         return false;
 
-    WasmActivation* activation = cx->wasmActivationStack();
+    WasmActivation* activation = MaybeActiveActivation(cx);
     if (!activation)
         return false;
 
@@ -1224,7 +1224,7 @@ HandleFault(int signum, siginfo_t* info, void* ctx)
         return false;
     AutoSetHandlingSegFault handling(cx);
 
-    WasmActivation* activation = cx->wasmActivationStack();
+    WasmActivation* activation = MaybeActiveActivation(cx);
     if (!activation)
         return false;
 
@@ -1332,38 +1332,54 @@ RedirectJitCodeToInterruptCheck(JSContext* cx, CONTEXT* context)
     if (cx != cx->runtime()->activeContext())
         return false;
 
+    // The faulting thread is suspended so we can access cx fields that can
+    // normally only be accessed by the cx's active thread.
+    AutoNoteSingleThreadedRegion anstr;
+
     RedirectIonBackedgesToInterruptCheck(cx);
 
-    if (WasmActivation* activation = cx->wasmActivationStack()) {
 #ifdef JS_SIMULATOR
-        (void)ContextToPC(context);  // silence static 'unused' errors
-
-        void* pc = cx->simulator()->get_pc_as<void*>();
-
-        const Code* code = activation->compartment()->wasm.lookupCode(pc);
-        if (code && code->segment().containsFunctionPC(pc))
-            cx->simulator()->trigger_wasm_interrupt();
+    uint8_t* pc = cx->simulator()->get_pc_as<uint8_t*>();
 #else
-        uint8_t** ppc = ContextToPC(context);
-        uint8_t* pc = *ppc;
-        uint8_t* fp = ContextToFP(context);
+    uint8_t* pc = *ContextToPC(context);
+#endif
 
     // Only interrupt in function code so that the frame iterators have the
     // invariant that resumePC always has a function CodeRange and we can't
-        // get into any weird interrupt-during-interrupt-stub cases. Note that
-        // the out-of-bounds/unaligned trap paths which call startInterrupt() go
-        // through function code, so test if already interrupted. All these
-        // paths are temporary though, so this case can be removed later.
-        const Code* code = activation->compartment()->wasm.lookupCode(pc);
-        if (code && code->segment().containsFunctionPC(pc) && fp && !activation->interrupted()) {
-            activation->startInterrupt(pc, fp);
-            *ppc = code->segment().interruptCode();
-            return true;
-        }
-#endif
-    }
-
+    // get into any weird interrupt-during-interrupt-stub cases.
+    if (!cx->compartment())
         return false;
+    const Code* code = cx->compartment()->wasm.lookupCode(pc);
+    if (!code || !code->segment().containsFunctionPC(pc))
+        return false;
+
+    // Only probe cx->activation() via MaybeActiveActivation after we know the
+    // pc is in wasm code. This way we don't depend on signal-safe update of
+    // cx->activation().
+    WasmActivation* activation = MaybeActiveActivation(cx);
+    MOZ_ASSERT(activation);
+
+#ifdef JS_SIMULATOR
+    // The checks performed by the !JS_SIMULATOR path happen in
+    // Simulator::handleWasmInterrupt.
+    cx->simulator()->trigger_wasm_interrupt();
+#else
+    // fp may be null when first entering wasm code from an entry stub.
+    uint8_t* fp = ContextToFP(context);
+    if (!fp)
+        return false;
+
+    // The out-of-bounds/unaligned trap paths which call startInterrupt() go
+    // through function code, so test if already interrupted. These paths are
+    // temporary though, so this case can be removed later.
+    if (activation->interrupted())
+        return false;
+
+    activation->startInterrupt(pc, fp);
+    *ContextToPC(context) = code->segment().interruptCode();
+#endif
+
+    return true;
 }
 
 #if !defined(XP_WIN)
@@ -1577,24 +1593,3 @@ js::InterruptRunningJitCode(JSContext* cx)
     pthread_kill(thread, sInterruptSignal);
 #endif
 }
-
-MOZ_COLD bool
-js::wasm::IsPCInWasmCode(void* pc)
-{
-    JSContext* cx = TlsContext.get();
-    if (!cx)
-        return false;
-
-    MOZ_RELEASE_ASSERT(!cx->handlingSegFault);
-
-    WasmActivation* activation = cx->wasmActivationStack();
-    if (!activation)
-        return false;
-
-    if (activation->compartment()->wasm.lookupCode(pc))
-        return true;
-
-    const CodeRange* codeRange;
-    uint8_t* codeBase;
-    return LookupBuiltinThunk(pc, &codeRange, &codeBase);
-}

js/src/wasm/WasmSignalHandlers.h

@@ -72,10 +72,6 @@ class MachExceptionHandler
 };
 #endif
 
-// Test whether the given PC is within the innermost wasm activation. Return
-// false if it is not, or it cannot be determined.
-bool IsPCInWasmCode(void *pc);
-
 } // namespace wasm
 } // namespace js