Generalize handling of stack checks indicated by highest exit + 1.

4 files changed, 22 insertions, 12 deletions

diff --git a/src/lj_asm.c b/src/lj_asm.c
index c795b99e..bb134ac2 100644
--- a/src/lj_asm.c
+++ b/src/lj_asm.c
@@ -1357,9 +1357,15 @@ static void asm_head_side(ASMState *as)
   /* Inherit top stack slot already checked by parent trace. */
   as->T->topslot = as->parent->topslot;
   if (as->topslot > as->T->topslot) {  /* Need to check for higher slot? */
-    as->T->topslot = (uint8_t)as->topslot;  /* Remember for child traces. */
+#ifdef EXITSTATE_CHECKEXIT
+    /* Highest exit + 1 indicates stack check. */
+    ExitNo exitno = as->T->nsnap;
+#else
     /* Reuse the parent exit in the context of the parent trace. */
-    asm_stack_check(as, as->topslot, irp, allow & RSET_GPR, as->J->exitno);
+    ExitNo exitno = as->J->exitno;
+#endif
+    as->T->topslot = (uint8_t)as->topslot;  /* Remember for child traces. */
+    asm_stack_check(as, as->topslot, irp, allow & RSET_GPR, exitno);
   }
 }
 
diff --git a/src/lj_asm_arm.h b/src/lj_asm_arm.h
index 27f1d6c1..786dd831 100644
--- a/src/lj_asm_arm.h
+++ b/src/lj_asm_arm.h
@@ -1418,7 +1418,6 @@ static void asm_stack_check(ASMState *as, BCReg topslot,
   Reg pbase;
   uint32_t k;
   if (irp) {
-    exitno = as->T->nsnap;  /* Highest exit + 1 indicates stack check. */
     if (ra_hasreg(irp->r)) {
       pbase = irp->r;
     } else if (allow) {
diff --git a/src/lj_target_arm.h b/src/lj_target_arm.h
index 210fe497..0f50cf9c 100644
--- a/src/lj_target_arm.h
+++ b/src/lj_target_arm.h
@@ -117,6 +117,8 @@ typedef struct {
 
 /* PC after instruction that caused an exit. Used to find the trace number. */
 #define EXITSTATE_PCREG         RID_PC
+/* Highest exit + 1 indicates stack check. */
+#define EXITSTATE_CHECKEXIT     1
 
 #define EXITSTUB_SPACING        4
 #define EXITSTUBS_PER_GROUP     32
diff --git a/src/lj_trace.c b/src/lj_trace.c
index f907a77f..9645ecba 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -727,14 +727,8 @@ static TraceNo trace_exit_find(jit_State *J, MCode *pc)
   TraceNo traceno;
   for (traceno = 1; traceno < J->sizetrace; traceno++) {
     GCtrace *T = traceref(J, traceno);
-    if (T && pc >= T->mcode && pc < (MCode *)((char *)T->mcode + T->szmcode)) {
-      if (J->exitno == T->nsnap) {  /* Treat stack check like a parent exit. */
-        lua_assert(T->root != 0);
-        traceno = T->ir[REF_BASE].op1;
-        J->exitno = T->ir[REF_BASE].op2;
-      }
+    if (T && pc >= T->mcode && pc < (MCode *)((char *)T->mcode + T->szmcode))
       return traceno;
-    }
   }
   lua_assert(0);
   return 0;
@@ -751,11 +745,20 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
   int errcode;
   const BCIns *pc;
   void *cf;
+  GCtrace *T;
 #ifdef EXITSTATE_PCREG
   J->parent = trace_exit_find(J, (MCode *)(intptr_t)ex->gpr[EXITSTATE_PCREG]);
 #endif
-  lua_assert(traceref(J, J->parent) != NULL &&
-             J->exitno < traceref(J, J->parent)->nsnap);
+  T = traceref(J, J->parent); UNUSED(T);
+#ifdef EXITSTATE_CHECKEXIT
+  if (J->exitno == T->nsnap) {  /* Treat stack check like a parent exit. */
+    lua_assert(T->root != 0);
+    J->exitno = T->ir[REF_BASE].op2;
+    J->parent = T->ir[REF_BASE].op1;
+    T = traceref(J, J->parent);
+  }
+#endif
+  lua_assert(T != NULL && J->exitno < T->nsnap);
   exd.J = J;
   exd.exptr = exptr;
   errcode = lj_vm_cpcall(L, NULL, &exd, trace_exit_cp);