Bug 1564887 - Simplify stack checks in Baseline. r=tcampbell

We no longer need the early-check mechanism if we always do environment chain
initialization + stack check before initializing locals. A comment mentions that
environment initialization can only happen _after_ pushing locals, but that
didn't match the code anyway: if the early stack check failed, we would end up
in initEnvironmentChain _without_ having initialized locals.

Ion prologue bailouts now resume in the prologue after environment chain
initialization, so the bailout code needed some small changes to always
initialize the environment chain instead of relying on the Baseline
prologue doing that.

Overall this is much simpler and deletes about 70 lines of code.

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

--HG--
extra : moz-landing-system : lando

js/src/jit/BaselineBailouts.cpp

@@ -2026,25 +2026,26 @@ bool jit::FinishBailoutToBaseline(BaselineBailoutInfo* bailoutInfo) {
   js_free(bailoutInfo);
   bailoutInfo = nullptr;
 
-  if (topFrame->environmentChain()) {
-    // Ensure the frame has a call object if it needs one. If the env chain
-    // is nullptr, we will enter baseline code at the prologue so no need to do
-    // anything in that case.
-    if (!EnsureHasEnvironmentObjects(cx, topFrame)) {
-      return false;
-    }
+  // If the frame has no environment chain, this must be a prologue bailout and
+  // we have to initialize with the function's initial environment to match
+  // emitInitFrameFields.
+  if (!topFrame->environmentChain()) {
+    topFrame->setEnvironmentChain(topFrame->callee()->environment());
+  }
 
-    // If we bailed out before Ion could do the global declaration
-    // conflicts check, because we resume in the body instead of the
-    // prologue for global frames.
-    if (checkGlobalDeclarationConflicts) {
-      Rooted<LexicalEnvironmentObject*> lexicalEnv(
-          cx, &cx->global()->lexicalEnvironment());
-      RootedScript script(cx, topFrame->script());
-      if (!CheckGlobalDeclarationConflicts(cx, script, lexicalEnv,
-                                           cx->global())) {
-        return false;
-      }
+  // Ensure the frame has a call object if it needs one.
+  if (!EnsureHasEnvironmentObjects(cx, topFrame)) {
+    return false;
+  }
+
+  // Check for global declaration conflicts if necessary.
+  if (checkGlobalDeclarationConflicts) {
+    Rooted<LexicalEnvironmentObject*> lexicalEnv(
+        cx, &cx->global()->lexicalEnvironment());
+    RootedScript script(cx, topFrame->script());
+    if (!CheckGlobalDeclarationConflicts(cx, script, lexicalEnv,
+                                         cx->global())) {
+      return false;
     }
   }
 
@@ -2111,14 +2112,6 @@ bool jit::FinishBailoutToBaseline(BaselineBailoutInfo* bailoutInfo) {
     ++iter;
   }
 
-  // If the frame has no environment chain, this must be a prologue bailout and
-  // we have to initialize with the function's initial environment to match
-  // emitInitFrameFields. This has to happen last because the earlier bailout
-  // code uses |envChain == nullptr| to indicate a prologue bailout.
-  if (!topFrame->environmentChain()) {
-    topFrame->setEnvironmentChain(topFrame->callee()->environment());
-  }
-
   MOZ_ASSERT(innerScript);
   MOZ_ASSERT(outerScript);
   MOZ_ASSERT(outerFp);
@@ -2162,6 +2155,14 @@ bool jit::FinishBailoutToBaseline(BaselineBailoutInfo* bailoutInfo) {
     UnwindEnvironment(cx, ei, tryPC);
   }
 
+  // Check for interrupts now because we might miss an interrupt check in JIT
+  // code when resuming in the prologue, after the stack/interrupt check.
+  if (!cx->isExceptionPending()) {
+    if (!CheckForInterrupt(cx)) {
+      return false;
+    }
+  }
+
   JitSpew(JitSpew_BaselineBailouts,
           "  Restored outerScript=(%s:%u:%u,%u) innerScript=(%s:%u:%u,%u) "
           "(bailoutKind=%u)",

js/src/jit/BaselineCompiler.cpp

@@ -713,30 +713,16 @@ bool BaselineCodeGen<Handler>::callVMInternal(VMFunctionId id,
                            BaselineFrame::reverseOffsetOfFrameSize());
   uint32_t frameBaseSize =
       BaselineFrame::FramePointerOffset + BaselineFrame::Size();
-  if (phase == POST_INITIALIZE) {
+  if (phase == CallVMPhase::AfterPushingLocals) {
     storeFrameSizeAndPushDescriptor(frameBaseSize, argSize, frameSizeAddress,
                                     R0.scratchReg(), R1.scratchReg());
   } else {
-    MOZ_ASSERT(phase == CHECK_OVER_RECURSED);
-    Label done, pushedFrameLocals;
-
-    // If OVER_RECURSED is set, then frame locals haven't been pushed yet.
-    masm.branchTest32(Assembler::Zero, frame.addressOfFlags(),
-                      Imm32(BaselineFrame::OVER_RECURSED), &pushedFrameLocals);
-    {
-      masm.store32(Imm32(frameBaseSize), frameSizeAddress);
-      uint32_t descriptor =
-          MakeFrameDescriptor(frameBaseSize + argSize, FrameType::BaselineJS,
-                              ExitFrameLayout::Size());
-      masm.push(Imm32(descriptor));
-      masm.jump(&done);
-    }
-    masm.bind(&pushedFrameLocals);
-    {
-      storeFrameSizeAndPushDescriptor(frameBaseSize, argSize, frameSizeAddress,
-                                      R0.scratchReg(), R1.scratchReg());
-    }
-    masm.bind(&done);
+    MOZ_ASSERT(phase == CallVMPhase::BeforePushingLocals);
+    masm.store32(Imm32(frameBaseSize), frameSizeAddress);
+    uint32_t descriptor =
+        MakeFrameDescriptor(frameBaseSize + argSize, FrameType::BaselineJS,
+                            ExitFrameLayout::Size());
+    masm.push(Imm32(descriptor));
   }
   MOZ_ASSERT(fun.expectTailCall == NonTailCall);
   // Perform the call.
@@ -773,36 +759,26 @@ bool BaselineCodeGen<Handler>::callVM(CallVMPhase phase) {
 
 template <typename Handler>
 bool BaselineCodeGen<Handler>::emitStackCheck() {
-  // If this is the late stack check for a frame which contains an early stack
-  // check, then the early stack check might have failed and skipped past the
-  // pushing of locals on the stack.
-  //
-  // If this is a possibility, then the OVER_RECURSED flag should be checked,
-  // and the VMCall to CheckOverRecursedBaseline done unconditionally if it's
-  // set.
-  Label forceCall;
-  if (handler.needsEarlyStackCheck()) {
-    masm.branchTest32(Assembler::NonZero, frame.addressOfFlags(),
-                      Imm32(BaselineFrame::OVER_RECURSED), &forceCall);
-  }
-
   Label skipCall;
-  masm.branchStackPtrRhs(Assembler::BelowOrEqual,
-                         AbsoluteAddress(cx->addressOfJitStackLimit()),
-                         &skipCall);
-
-  if (handler.needsEarlyStackCheck()) {
-    masm.bind(&forceCall);
+  if (handler.mustIncludeSlotsInStackCheck()) {
+    // Subtract the size of script->nslots() first.
+    Register scratch = R1.scratchReg();
+    masm.moveStackPtrTo(scratch);
+    subtractScriptSlotsSize(scratch, R2.scratchReg());
+    masm.branchPtr(Assembler::BelowOrEqual,
+                   AbsoluteAddress(cx->addressOfJitStackLimit()), scratch,
+                   &skipCall);
+  } else {
+    masm.branchStackPtrRhs(Assembler::BelowOrEqual,
+                           AbsoluteAddress(cx->addressOfJitStackLimit()),
+                           &skipCall);
   }
 
   prepareVMCall();
   masm.loadBaselineFramePtr(BaselineFrameReg, R1.scratchReg());
   pushArg(R1.scratchReg());
 
-  CallVMPhase phase = POST_INITIALIZE;
-  if (handler.needsEarlyStackCheck()) {
-    phase = CHECK_OVER_RECURSED;
-  }
+  const CallVMPhase phase = CallVMPhase::BeforePushingLocals;
 
   using Fn = bool (*)(JSContext*, BaselineFrame*);
   if (!callVMNonOp<Fn, CheckOverRecursedBaseline>(phase)) {
@@ -1287,19 +1263,16 @@ bool BaselineInterpreterCodeGen::initEnvironmentChainHelper(
 
 template <typename Handler>
 bool BaselineCodeGen<Handler>::initEnvironmentChain() {
-  CallVMPhase phase = POST_INITIALIZE;
-  if (handler.needsEarlyStackCheck()) {
-    phase = CHECK_OVER_RECURSED;
-  }
-
-  auto initFunctionEnv = [this, phase]() {
-    auto initEnv = [this, phase]() {
+  auto initFunctionEnv = [this]() {
+    auto initEnv = [this]() {
       // Call into the VM to create the proper environment objects.
       prepareVMCall();
 
       masm.loadBaselineFramePtr(BaselineFrameReg, R0.scratchReg());
       pushArg(R0.scratchReg());
 
+      const CallVMPhase phase = CallVMPhase::BeforePushingLocals;
+
       using Fn = bool (*)(JSContext*, BaselineFrame*);
       return callVMNonOp<Fn, jit::InitFunctionEnvironmentObjects>(phase);
     };
@@ -1308,7 +1281,7 @@ bool BaselineCodeGen<Handler>::initEnvironmentChain() {
         initEnv, R2.scratchReg());
   };
 
-  auto initGlobalOrEvalEnv = [this, phase]() {
+  auto initGlobalOrEvalEnv = [this]() {
     // EnvironmentChain pointer in BaselineFrame has already been initialized
     // in prologue, but we need to check for redeclaration errors in global and
     // eval scripts.
@@ -1319,6 +1292,8 @@ bool BaselineCodeGen<Handler>::initEnvironmentChain() {
     masm.loadPtr(frame.addressOfEnvironmentChain(), R0.scratchReg());
     pushArg(R0.scratchReg());
 
+    const CallVMPhase phase = CallVMPhase::BeforePushingLocals;
+
     using Fn = bool (*)(JSContext*, HandleObject, HandleScript);
     return callVMNonOp<Fn, js::CheckGlobalOrEvalDeclarationConflicts>(phase);
   };
@@ -6719,79 +6694,40 @@ bool BaselineCodeGen<Handler>::emitPrologue() {
     return false;
   }
 
-  // Functions with a large number of locals require two stack checks.
-  // The VMCall for a fallible stack check can only occur after the
-  // env chain has been initialized, as that is required for proper
-  // exception handling if the VMCall returns false.  The env chain
-  // initialization can only happen after the UndefinedValues for the
-  // local slots have been pushed. However by that time, the stack might
-  // have grown too much.
-  //
-  // In these cases, we emit an extra, early, infallible check before pushing
-  // the locals. The early check just sets a flag on the frame if the stack
-  // check fails. If the flag is set, then the jitcode skips past the pushing
-  // of the locals, and directly to env chain initialization followed by the
-  // actual stack check, which will throw the correct exception.
-  Label earlyStackCheckFailed;
-  if (handler.needsEarlyStackCheck()) {
-    // Subtract the size of script->nslots() from the stack pointer.
-    Register scratch = R1.scratchReg();
-    masm.moveStackPtrTo(scratch);
-    subtractScriptSlotsSize(scratch, R2.scratchReg());
+  // When compiling with Debugger instrumentation, set the debuggeeness of
+  // the frame before any operation that can call into the VM.
+  if (!emitIsDebuggeeCheck()) {
+    return false;
+  }
 
-    // Set the OVER_RECURSED flag on the frame if the computed stack pointer
-    // overflows the stack limit. We have to use the actual (*NoInterrupt)
-    // stack limit here because we don't want to set the flag and throw an
-    // overrecursion exception later in the interrupt case.
-    Label stackCheckOk;
-    masm.branchPtr(Assembler::BelowOrEqual,
-                   AbsoluteAddress(cx->addressOfJitStackLimitNoInterrupt()),
-                   scratch, &stackCheckOk);
-    {
-      masm.or32(Imm32(BaselineFrame::OVER_RECURSED), frame.addressOfFlags());
-      masm.jump(&earlyStackCheckFailed);
-    }
-    masm.bind(&stackCheckOk);
+  // Initialize the env chain before any operation that may call into the VM and
+  // trigger a GC.
+  if (!initEnvironmentChain()) {
+    return false;
+  }
+
+  // Check for overrecursion before initializing locals.
+  if (!emitStackCheck()) {
+    return false;
   }
 
   emitInitializeLocals();
 
-  if (handler.needsEarlyStackCheck()) {
-    masm.bind(&earlyStackCheckFailed);
-  }
-
 #ifdef JS_TRACE_LOGGING
   if (JS::TraceLoggerSupported() && !emitTraceLoggerEnter()) {
     return false;
   }
 #endif
 
-  // Record the offset of the prologue, because Ion can bailout before
-  // the env chain is initialized.
+  // Record prologue offset for Ion bailouts.
   bailoutPrologueOffset_ = CodeOffset(masm.currentOffset());
 
-  // When compiling with Debugger instrumentation, set the debuggeeness of
-  // the frame before any operation that can call into the VM.
-  if (!emitIsDebuggeeCheck()) {
-    return false;
-  }
-
-  // Initialize the env chain before any operation that may
-  // call into the VM and trigger a GC.
-  if (!initEnvironmentChain()) {
-    return false;
-  }
-
   frame.assertSyncedStack();
 
   if (JSScript* script = handler.maybeScript()) {
     masm.debugAssertContextRealm(script->realm(), R1.scratchReg());
   }
 
-  if (!emitStackCheck()) {
-    return false;
-  }
-
   if (!emitDebugPrologue()) {
     return false;
   }

js/src/jit/BaselineCompiler.h

@@ -363,14 +363,14 @@ class BaselineCodeGen {
                                        const Address& frameSizeAddr,
                                        Register scratch1, Register scratch2);
 
-  enum CallVMPhase { POST_INITIALIZE, CHECK_OVER_RECURSED };
+  enum class CallVMPhase { BeforePushingLocals, AfterPushingLocals };
   bool callVMInternal(VMFunctionId id, CallVMPhase phase);
 
   template <typename Fn, Fn fn>
-  bool callVM(CallVMPhase phase = POST_INITIALIZE);
+  bool callVM(CallVMPhase phase = CallVMPhase::AfterPushingLocals);
 
   template <typename Fn, Fn fn>
-  bool callVMNonOp(CallVMPhase phase = POST_INITIALIZE) {
+  bool callVMNonOp(CallVMPhase phase = CallVMPhase::AfterPushingLocals) {
     if (!callVM<Fn, fn>(phase)) {
       return false;
     }
@@ -597,11 +597,11 @@ class BaselineCompilerHandler {
     retAddrEntries_.back().setKind(kind);
   }
 
-  // If a script has more |nslots| than this, then emit code to do an
-  // early stack check.
-  bool needsEarlyStackCheck() const {
-    static const unsigned EARLY_STACK_CHECK_SLOT_COUNT = 128;
-    return script()->nslots() > EARLY_STACK_CHECK_SLOT_COUNT;
+  // If a script has more |nslots| than this the stack check must account
+  // for these slots explicitly.
+  bool mustIncludeSlotsInStackCheck() const {
+    static constexpr size_t NumSlotsLimit = 128;
+    return script()->nslots() > NumSlotsLimit;
   }
 
   JSObject* maybeNoCloneSingletonObject();
@@ -725,9 +725,9 @@ class BaselineInterpreterHandler {
 
   bool maybeIonCompileable() const { return true; }
 
-  // The interpreter always does the early stack check because we don't know the
-  // frame size statically.
-  bool needsEarlyStackCheck() const { return true; }
+  // The interpreter doesn't know the number of slots statically so we always
+  // include them.
+  bool mustIncludeSlotsInStackCheck() const { return true; }
 
   JSObject* maybeNoCloneSingletonObject() { return nullptr; }
 };

js/src/jit/BaselineFrame.h

@@ -51,14 +51,9 @@ class BaselineFrame {
     // of invariants of debuggee compartments, scripts, and frames.
     DEBUGGEE = 1 << 6,
 
-    // (1 << 7 and 1 << 8 are unused)
-
-    // Frame has over-recursed on an early check.
-    OVER_RECURSED = 1 << 9,
-
     // Frame has a BaselineRecompileInfo stashed in the scratch value
     // slot. See PatchBaselineFramesForDebugMode.
-    HAS_DEBUG_MODE_OSR_INFO = 1 << 10,
+    HAS_DEBUG_MODE_OSR_INFO = 1 << 7,
 
     // This flag is intended for use whenever the frame is settled on a
     // native code address without a corresponding RetAddrEntry. In this
@@ -72,12 +67,12 @@ class BaselineFrame {
     // This flag should never be set on the top frame while we're
     // executing JIT code. In debug builds, it is checked before and
     // after VM calls.
-    HAS_OVERRIDE_PC = 1 << 11,
+    HAS_OVERRIDE_PC = 1 << 8,
 
     // If set, we're handling an exception for this frame. This is set for
     // debug mode OSR sanity checking when it handles corner cases which
     // only arise during exception handling.
-    HANDLING_EXCEPTION = 1 << 12,
+    HANDLING_EXCEPTION = 1 << 9,
   };
 
  protected:  // Silence Clang warning about unused private fields.
@@ -315,10 +310,6 @@ class BaselineFrame {
   void setIsHandlingException() { flags_ |= HANDLING_EXCEPTION; }
   void unsetIsHandlingException() { flags_ &= ~HANDLING_EXCEPTION; }
 
-  bool overRecursed() const { return flags_ & OVER_RECURSED; }
-
-  void setOverRecursed() { flags_ |= OVER_RECURSED; }
-
   BaselineDebugModeOSRInfo* debugModeOSRInfo() {
     MOZ_ASSERT(flags_ & HAS_DEBUG_MODE_OSR_INFO);
     return debugModeOSRInfo_;

js/src/jit/VMFunctions.cpp

@@ -297,7 +297,7 @@ bool InvokeFromInterpreterStub(JSContext* cx,
   return true;
 }
 
-bool CheckOverRecursed(JSContext* cx) {
+static bool CheckOverRecursedImpl(JSContext* cx, size_t extra) {
   // We just failed the jitStackLimit check. There are two possible reasons:
   //  1) jitStackLimit was the real stack limit and we're over-recursed
   //  2) jitStackLimit was set to UINTPTR_MAX by JSContext::requestInterrupt
@@ -305,12 +305,12 @@ bool CheckOverRecursed(JSContext* cx) {
 
   // This handles 1).
 #ifdef JS_SIMULATOR
-  if (cx->simulator()->overRecursedWithExtra(0)) {
+  if (cx->simulator()->overRecursedWithExtra(extra)) {
     ReportOverRecursed(cx);
     return false;
   }
 #else
-  if (!CheckRecursionLimit(cx)) {
+  if (!CheckRecursionLimitWithExtra(cx, extra)) {
     return false;
   }
 #endif
@@ -320,19 +320,13 @@ bool CheckOverRecursed(JSContext* cx) {
   return cx->handleInterrupt();
 }
 
-// This function gets called when the overrecursion check fails for a Baseline
-// frame. This is just like CheckOverRecursed, with an extra check to handle
-// early stack check failures.
-bool CheckOverRecursedBaseline(JSContext* cx, BaselineFrame* frame) {
-  // The OVERRECURSED flag may have already been set on the frame by an
-  // early over-recursed check (before pushing the locals).  If so, throw
-  // immediately.
-  if (frame->overRecursed()) {
-    ReportOverRecursed(cx);
-    return false;
-  }
+bool CheckOverRecursed(JSContext* cx) { return CheckOverRecursedImpl(cx, 0); }
 
-  return CheckOverRecursed(cx);
+bool CheckOverRecursedBaseline(JSContext* cx, BaselineFrame* frame) {
+  // The stack check in Baseline happens before pushing locals so we have to
+  // account for that by including script->nslots() in the C++ recursion check.
+  size_t extra = frame->script()->nslots() * sizeof(Value);
+  return CheckOverRecursedImpl(cx, extra);
 }
 
 bool MutatePrototype(JSContext* cx, HandlePlainObject obj, HandleValue value) {

js/src/jsfriendapi.h

@@ -1058,6 +1058,18 @@ MOZ_ALWAYS_INLINE bool CheckRecursionLimitWithStackPointer(JSContext* cx,
   return true;
 }
 
+MOZ_ALWAYS_INLINE bool CheckRecursionLimitWithExtra(JSContext* cx,
+                                                    size_t extra) {
+  char stackDummy;
+  char* sp = &stackDummy;
+#if JS_STACK_GROWTH_DIRECTION > 0
+  sp += extra;
+#else
+  sp -= extra;
+#endif
+  return CheckRecursionLimitWithStackPointer(cx, sp);
+}
+
 MOZ_ALWAYS_INLINE bool CheckSystemRecursionLimit(JSContext* cx) {
   return CheckRecursionLimit(cx,
                              GetNativeStackLimit(cx, JS::StackForSystemCode));