Bug 1910880 - Use a real JIT exit frame for Wasm JitEntry stub. r=rhunt

JS JIT callers can call Wasm code either directly from Ion or through the JitEntry stub.
In the JitEntry case, we didn't have a real JS JIT exit frame. To make JS frame iteration
still work in this case, a special `JSJitToWasm` frame type was used to transition from Wasm
frame iteration to JIT frame iteration.

Now that we can unwind Wasm and JS frames directly, this caused some issues because there
were places where we were expecting an actual JIT exit frame.

This patch changes the JitEntry stub to initialize a JIT exit frame. To do this, we just need
to store an `ExitFooterFrame` below the frame pointer. This is similar to what we were already
doing in `GenerateJitEntryThrow`.

This lets us remove the `JSJitToWasm` frame type and simplify frame iteration a bit.

Differential Revision: https://phabricator.services.mozilla.com/D218346

Author: jandem

js/src/gdb/mozilla/unwind.py

@@ -49,7 +49,6 @@ def debug(something):
     "FrameType::BaselineStub": "BaselineStubFrameLayout",
     "FrameType::CppToJSJit": "JitFrameLayout",
     "FrameType::WasmToJSJit": "JitFrameLayout",
-    "FrameType::JSJitToWasm": "JitFrameLayout",
     "FrameType::Rectifier": "RectifierFrameLayout",
     "FrameType::IonAccessorIC": "IonAccessorICFrameLayout",
     "FrameType::IonICCall": "IonICCallFrameLayout",

js/src/jit-test/tests/ion/bug1910880.js

@@ -0,0 +1,38 @@
+// |jit-test| --fast-warmup; skip-if: !wasmIsSupported()
+
+gczeal(2);
+
+function wasmEvalText(str) {
+    var bin = wasmTextToBinary(str);
+    var m = new WebAssembly.Module(bin);
+    return new WebAssembly.Instance(m);
+}
+function test() {
+    var instance = wasmEvalText(`
+        (module (type $a (array (mut i32)))
+        (func (export "createDefault") (param i32) (result eqref)
+            local.get 0
+            array.new_default $a
+        )
+        )
+    `);
+    var createDefault = instance.exports.createDefault;
+
+    var g = newGlobal({newCompartment: true});
+    g.debuggeeGlobal = this;
+    g.eval("(" + function () {
+        var dbg = new Debugger(debuggeeGlobal);
+        dbg.onExceptionUnwind = function () {
+            throw new Error("x");
+        };
+    } + ")();");
+
+    for (var i = 0; i < 8; i++) {
+        try {
+            createDefault(-1);
+        } catch (e) {
+        }
+    }
+}
+test();
+quit(0); // Ensure exit code is 0, not 3.

js/src/jit/Ion.cpp

@@ -2325,20 +2325,15 @@ static void InvalidateActivation(JS::GCContext* gcx,
 
   for (OnlyJSJitFrameIter iter(activations); !iter.done(); ++iter, ++frameno) {
     const JSJitFrameIter& frame = iter.frame();
-    MOZ_ASSERT_IF(frameno == 1, frame.isExitFrame() ||
-                                    frame.type() == FrameType::Bailout ||
-                                    frame.type() == FrameType::JSJitToWasm);
+    MOZ_ASSERT_IF(frameno == 1,
+                  frame.isExitFrame() || frame.type() == FrameType::Bailout);
 
 #ifdef JS_JITSPEW
     switch (frame.type()) {
       case FrameType::Exit:
         JitSpew(JitSpew_IonInvalidate, "#%zu exit frame @ %p", frameno,
                 frame.fp());
         break;
-      case FrameType::JSJitToWasm:
-        JitSpew(JitSpew_IonInvalidate, "#%zu wasm exit frame @ %p", frameno,
-                frame.fp());
-        break;
       case FrameType::BaselineJS:
       case FrameType::IonJS:
       case FrameType::Bailout: {

js/src/jit/JSJitFrameIter.cpp

@@ -38,20 +38,16 @@ JSJitFrameIter::JSJitFrameIter(const JitActivation* activation)
   MOZ_ASSERT(!TlsContext.get()->inUnsafeCallWithABI);
 }
 
-JSJitFrameIter::JSJitFrameIter(const JitActivation* activation,
-                               FrameType frameType, uint8_t* fp, bool unwinding)
-    : current_(fp), type_(frameType), activation_(activation) {
+JSJitFrameIter::JSJitFrameIter(const JitActivation* activation, uint8_t* fp,
+                               bool unwinding)
+    : current_(fp), type_(FrameType::Exit), activation_(activation) {
   // This constructor is only used when resuming iteration after iterating Wasm
   // frames in the same JitActivation so ignore activation_->bailoutData().
-  //
-  // Note: FrameType::JSJitToWasm is used for JIT => Wasm calls through the Wasm
-  // JIT entry trampoline. FrameType::Exit is used for direct Ion => Wasm calls.
   if (unwinding) {
     MOZ_ASSERT(fp == activation->jsExitFP());
   } else {
     MOZ_ASSERT(fp > activation->jsOrWasmExitFP());
   }
-  MOZ_ASSERT(type_ == FrameType::JSJitToWasm || type_ == FrameType::Exit);
   MOZ_ASSERT(!TlsContext.get()->inUnsafeCallWithABI);
 }
 
@@ -409,9 +405,6 @@ void JSJitFrameIter::dump() const {
     case FrameType::Exit:
       fprintf(stderr, " Exit frame\n");
       break;
-    case FrameType::JSJitToWasm:
-      fprintf(stderr, " Wasm exit frame\n");
-      break;
   };
   fputc('\n', stderr);
 }
@@ -815,7 +808,6 @@ void JSJitProfilingFrameIterator::moveToNextFrame(CommonFrameLayout* frame) {
     case FrameType::TrampolineNative:
     case FrameType::Exit:
     case FrameType::Bailout:
-    case FrameType::JSJitToWasm:
       // Rectifier and Baseline Interpreter entry frames are handled before
       // this switch. The other frame types can't call JS functions directly.
       break;

js/src/jit/JSJitFrameIter.h

@@ -69,11 +69,6 @@ enum class FrameType {
   // point of view of JS JITs, this is just another kind of entry frame.
   WasmToJSJit,
 
-  // A JS to wasm frame is constructed during fast calls from any JS jits to
-  // wasm, and is a special kind of exit frame that doesn't have the exit
-  // footer. From the point of view of the jit, it can be skipped as an exit.
-  JSJitToWasm,
-
   // Frame for a TrampolineNative, a JS builtin implemented with a JIT
   // trampoline. See jit/TrampolineNatives.h.
   TrampolineNative,
@@ -133,8 +128,7 @@ class JSJitFrameIter {
 
   // A constructor specialized for jit->wasm frames, which starts at a
   // specific FP.
-  JSJitFrameIter(const JitActivation* activation, FrameType frameType,
-                 uint8_t* fp, bool unwinding);
+  JSJitFrameIter(const JitActivation* activation, uint8_t* fp, bool unwinding);
 
   void setResumePCInCurrentFrame(uint8_t* newAddr) {
     resumePCinCurrentFrame_ = newAddr;

js/src/jit/JitFrames.cpp

@@ -1429,14 +1429,6 @@ static void TraceRectifierFrame(JSTracer* trc, const JSJitFrameIter& frame) {
   TraceRoot(trc, &layout->thisv(), "rectifier-thisv");
 }
 
-static void TraceJSJitToWasmFrame(JSTracer* trc, const JSJitFrameIter& frame) {
-  // This is doing a subset of TraceIonJSFrame, since the callee doesn't
-  // have a script.
-  JitFrameLayout* layout = (JitFrameLayout*)frame.fp();
-  layout->replaceCalleeToken(TraceCalleeToken(trc, layout->calleeToken()));
-  TraceThisAndArguments(trc, frame, layout);
-}
-
 static void TraceTrampolineNativeFrame(JSTracer* trc,
                                        const JSJitFrameIter& frame) {
   auto* layout = (TrampolineNativeFrameLayout*)frame.fp();
@@ -1507,9 +1499,6 @@ static void TraceJitActivation(JSTracer* trc, JitActivation* activation) {
           // JitFrameIter know the frame above is a wasm frame, handled
           // in the next iteration.
           break;
-        case FrameType::JSJitToWasm:
-          TraceJSJitToWasmFrame(trc, jitFrame);
-          break;
         default:
           MOZ_CRASH("unexpected frame type");
       }

js/src/vm/FrameIter.cpp

@@ -202,15 +202,13 @@ void JitFrameIter::settle() {
 
     MOZ_ASSERT(wasmFrame.done());
     uint8_t* prevFP = wasmFrame.unwoundCallerFP();
-    jit::FrameType prevFrameType = wasmFrame.unwoundJitFrameType();
 
     if (mustUnwindActivation_) {
       act_->setJSExitFP(prevFP);
     }
 
     iter_.destroy();
-    iter_.construct<jit::JSJitFrameIter>(act_, prevFrameType, prevFP,
-                                         mustUnwindActivation_);
+    iter_.construct<jit::JSJitFrameIter>(act_, prevFP, mustUnwindActivation_);
     MOZ_ASSERT(!asJSJit().done());
     return;
   }

js/src/wasm/WasmFrameIter.cpp

@@ -180,17 +180,21 @@ void WasmFrameIter::operator++() {
   popFrame();
 }
 
-static inline void AssertDirectJitCall(const void* fp) {
-  // Called via an inlined fast JIT to wasm call: in this case, FP is
-  // pointing in the middle of the exit frame, right before the exit
-  // footer; ensure the exit frame type is the expected one.
+static inline void AssertJitExitFrame(const void* fp,
+                                      jit::ExitFrameType expected) {
+  // Called via a JIT to wasm call: in this case, FP is pointing in the middle
+  // of the exit frame, right before the exit footer; ensure the exit frame type
+  // is the expected one.
 #ifdef DEBUG
   auto* jitCaller = (ExitFrameLayout*)fp;
-  MOZ_ASSERT(jitCaller->footer()->type() ==
-             jit::ExitFrameType::DirectWasmJitCall);
+  MOZ_ASSERT(jitCaller->footer()->type() == expected);
 #endif
 }
 
+static inline void AssertDirectJitCall(const void* fp) {
+  AssertJitExitFrame(fp, jit::ExitFrameType::DirectWasmJitCall);
+}
+
 void WasmFrameIter::popFrame() {
   uint8_t* returnAddress = fp_->returnAddress();
   code_ = LookupCode(returnAddress, &codeRange_);
@@ -215,7 +219,7 @@ void WasmFrameIter::popFrame() {
     AssertDirectJitCall(fp_->jitEntryCaller());
 
     unwoundCallerFP_ = fp_->jitEntryCaller();
-    unwoundJitFrameType_.emplace(FrameType::Exit);
+    hasUnwoundJitFrame_ = true;
 
     if (unwind_ == Unwind::True) {
       activation_->setJSExitFP(unwoundCallerFP());
@@ -239,6 +243,7 @@ void WasmFrameIter::popFrame() {
   if (codeRange_->isInterpEntry()) {
     // Interpreter entry has a simple frame, record FP from it.
     unwoundCallerFP_ = reinterpret_cast<uint8_t*>(fp_);
+    MOZ_ASSERT(!hasUnwoundJitFrame_);
 
     fp_ = nullptr;
     code_ = nullptr;
@@ -266,10 +271,12 @@ void WasmFrameIter::popFrame() {
     // |      WASM FRAME     | (already unwound)
     // |---------------------|
     //
-    // The next value of FP is just a regular jit frame used as a marker to
-    // know that we should transition to a JSJit frame iterator.
+    // The next value of FP is a jit exit frame with type WasmGenericJitEntry.
+    // This lets us transition to a JSJit frame iterator.
     unwoundCallerFP_ = reinterpret_cast<uint8_t*>(fp_);
-    unwoundJitFrameType_.emplace(FrameType::JSJitToWasm);
+    hasUnwoundJitFrame_ = true;
+    AssertJitExitFrame(unwoundCallerFP_,
+                       jit::ExitFrameType::WasmGenericJitEntry);
 
     fp_ = nullptr;
     code_ = nullptr;
@@ -411,13 +418,8 @@ DebugFrame* WasmFrameIter::debugFrame() const {
 }
 
 bool WasmFrameIter::hasUnwoundJitFrame() const {
-  return unwoundCallerFP_ && unwoundJitFrameType_.isSome();
-}
-
-jit::FrameType WasmFrameIter::unwoundJitFrameType() const {
-  MOZ_ASSERT(unwoundCallerFP_);
-  MOZ_ASSERT(unwoundJitFrameType_.isSome());
-  return *unwoundJitFrameType_;
+  MOZ_ASSERT_IF(hasUnwoundJitFrame_, unwoundCallerFP_);
+  return hasUnwoundJitFrame_;
 }
 
 uint8_t* WasmFrameIter::resumePCinCurrentFrame() const {

js/src/wasm/WasmFrameIter.h

@@ -69,7 +69,8 @@ class WasmFrameIter {
   Frame* fp_;
   Instance* instance_;
   uint8_t* unwoundCallerFP_;
-  mozilla::Maybe<jit::FrameType> unwoundJitFrameType_;
+  // Whether unwoundCallerFP_ is a JS JIT exit frame.
+  bool hasUnwoundJitFrame_ = false;
   Unwind unwind_;
   void** unwoundAddressOfReturnAddress_;
   uint8_t* resumePCinCurrentFrame_;
@@ -98,7 +99,6 @@ class WasmFrameIter {
   void** unwoundAddressOfReturnAddress() const;
   bool debugEnabled() const;
   DebugFrame* debugFrame() const;
-  jit::FrameType unwoundJitFrameType() const;
   bool hasUnwoundJitFrame() const;
   uint8_t* unwoundCallerFP() const { return unwoundCallerFP_; }
   Frame* frame() const { return fp_; }

js/src/wasm/WasmStubs.cpp

@@ -932,29 +932,40 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
   //
   // GenerateJitEntryPrologue has additionally pushed the caller's frame
   // pointer. The stack pointer is now JitStackAlignment-aligned.
+  //
+  // We initialize an ExitFooterFrame (with ExitFrameType::WasmGenericJitEntry)
+  // immediately below the frame pointer to ensure FP is a valid JS JIT exit
+  // frame.
 
   MOZ_ASSERT(masm.framePushed() == 0);
 
-  unsigned normalBytesNeeded = StackArgBytesForWasmABI(funcType);
+  unsigned normalBytesNeeded =
+      ExitFooterFrame::Size() + StackArgBytesForWasmABI(funcType);
 
   MIRTypeVector coerceArgTypes;
   MOZ_ALWAYS_TRUE(coerceArgTypes.append(MIRType::Int32));
   MOZ_ALWAYS_TRUE(coerceArgTypes.append(MIRType::Pointer));
   MOZ_ALWAYS_TRUE(coerceArgTypes.append(MIRType::Pointer));
-  unsigned oolBytesNeeded = StackArgBytesForWasmABI(coerceArgTypes);
+  unsigned oolBytesNeeded =
+      ExitFooterFrame::Size() + StackArgBytesForWasmABI(coerceArgTypes);
 
   unsigned bytesNeeded = std::max(normalBytesNeeded, oolBytesNeeded);
 
   // Note the jit caller ensures the stack is aligned *after* the call
   // instruction.
-  unsigned frameSizeExclFP = StackDecrementForCall(
-      WasmStackAlignment, masm.framePushed(), bytesNeeded);
+  unsigned frameSize = StackDecrementForCall(WasmStackAlignment,
+                                             masm.framePushed(), bytesNeeded);
 
   // Reserve stack space for wasm ABI arguments, set up like this:
   // <-- ABI args | padding
-  masm.reserveStack(frameSizeExclFP);
+  masm.reserveStack(frameSize);
+
+  MOZ_ASSERT(masm.framePushed() == frameSize);
 
-  uint32_t frameSize = masm.framePushed();
+  // Initialize the ExitFooterFrame.
+  static_assert(ExitFooterFrame::Size() == sizeof(uintptr_t));
+  masm.storePtr(ImmWord(uint32_t(ExitFrameType::WasmGenericJitEntry)),
+                Address(FramePointer, -int32_t(ExitFooterFrame::Size())));
 
   GenerateJitEntryLoadInstance(masm);
 
@@ -1198,9 +1209,6 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
   CallFuncExport(masm, fe, funcPtr);
   masm.assertStackAlignment(WasmStackAlignment);
 
-  // Pop arguments.
-  masm.freeStackTo(frameSize - frameSizeExclFP);
-
   GenPrintf(DebugChannel::Function, masm, "wasm-function[%d]; returns ",
             fe.funcIndex());
 
@@ -1234,18 +1242,11 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
         break;
       }
       case ValType::I64: {
-        Label fail, done;
         GenPrintI64(DebugChannel::Function, masm, ReturnReg64);
-        GenerateBigIntInitialization(masm, 0, ReturnReg64, scratchG, fe, &fail);
+        MOZ_ASSERT(masm.framePushed() == frameSize);
+        GenerateBigIntInitialization(masm, 0, ReturnReg64, scratchG, fe,
+                                     &exception);
         masm.boxNonDouble(JSVAL_TYPE_BIGINT, scratchG, JSReturnOperand);
-        masm.jump(&done);
-        masm.bind(&fail);
-        // Fixup the stack for the exception tail so that we can share it.
-        masm.reserveStack(frameSizeExclFP);
-        masm.jump(&exception);
-        masm.bind(&done);
-        // Un-fixup the stack for the benefit of the assertion below.
-        masm.setFramePushed(0);
         break;
       }
       case ValType::V128: {
@@ -1263,9 +1264,11 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
     }
   }
 
-  GenPrintf(DebugChannel::Function, masm, "\n");
+  // Pop frame.
+  masm.moveToStackPtr(FramePointer);
+  masm.setFramePushed(0);
 
-  MOZ_ASSERT(masm.framePushed() == 0);
+  GenPrintf(DebugChannel::Function, masm, "\n");
 
   AssertExpectedSP(masm);
   GenerateJitEntryEpilogue(masm, offsets);
@@ -1318,14 +1321,14 @@ static bool GenerateJitEntry(MacroAssembler& masm, size_t funcExportIndex,
     // No widening is required, as the return value is used as a bool.
     masm.branchTest32(Assembler::NonZero, ReturnReg, ReturnReg,
                       &rejoinBeforeCall);
+
+    MOZ_ASSERT(masm.framePushed() == frameSize);
     hasFallThroughForException = true;
   }
 
-  // Prepare to throw: reload InstanceReg from the frame.
-  masm.bind(&exception);
-  masm.setFramePushed(frameSize);
   if (exception.used() || hasFallThroughForException) {
-    masm.freeStackTo(frameSize);
+    masm.bind(&exception);
+    masm.setFramePushed(frameSize);
     GenerateJitEntryThrow(masm, frameSize);
   }