4 files changed, 31 insertions, 7 deletions

diff --git a/src/lj_asm_arm.h b/src/lj_asm_arm.h
index 087530b2..e884df5e 100644
--- a/src/lj_asm_arm.h
+++ b/src/lj_asm_arm.h
@@ -2040,6 +2040,9 @@ static void asm_stack_restore(ASMState *as, SnapShot *snap)
 
 /* -- GC handling --------------------------------------------------------- */
 
+/* Marker to prevent patching the GC check exit. */
+#define ARM_NOPATCH_GC_CHECK    (ARMI_BIC|ARMI_K12)
+
 /* Check GC threshold and do one or more GC steps. */
 static void asm_gc_check(ASMState *as)
 {
@@ -2051,6 +2054,7 @@ static void asm_gc_check(ASMState *as)
   l_end = emit_label(as);
   /* Exit trace if in GCSatomic or GCSfinalize. Avoids syncing GC objects. */
   asm_guardcc(as, CC_NE);  /* Assumes asm_snap_prep() already done. */
+  *--as->mcp = ARM_NOPATCH_GC_CHECK;
   emit_n(as, ARMI_CMP|ARMI_K12|0, RID_RET);
   args[0] = ASMREF_TMP1;  /* global_State *g */
   args[1] = ASMREF_TMP2;  /* MSize steps     */
@@ -2349,7 +2353,8 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
     /* Look for bl_cc exitstub, replace with b_cc target. */
     uint32_t ins = *p;
     if ((ins & 0x0f000000u) == 0x0b000000u && ins < 0xf0000000u &&
-        ((ins ^ (px-p)) & 0x00ffffffu) == 0) {
+        ((ins ^ (px-p)) & 0x00ffffffu) == 0 &&
+        p[-1] != ARM_NOPATCH_GC_CHECK) {
       *p = (ins & 0xfe000000u) | (((target-p)-2) & 0x00ffffffu);
       cend = p+1;
       if (!cstart) cstart = p;
diff --git a/src/lj_asm_mips.h b/src/lj_asm_mips.h
index 190a55eb..22031d11 100644
--- a/src/lj_asm_mips.h
+++ b/src/lj_asm_mips.h
@@ -1665,6 +1665,9 @@ static void asm_stack_restore(ASMState *as, SnapShot *snap)
 
 /* -- GC handling --------------------------------------------------------- */
 
+/* Marker to prevent patching the GC check exit. */
+#define MIPS_NOPATCH_GC_CHECK   MIPSI_OR
+
 /* Check GC threshold and do one or more GC steps. */
 static void asm_gc_check(ASMState *as)
 {
@@ -1680,6 +1683,7 @@ static void asm_gc_check(ASMState *as)
   args[0] = ASMREF_TMP1;  /* global_State *g */
   args[1] = ASMREF_TMP2;  /* MSize steps     */
   asm_gencall(as, ci, args);
+  l_end[-3] = MIPS_NOPATCH_GC_CHECK;  /* Replace the nop after the call. */
   emit_tsi(as, MIPSI_ADDIU, ra_releasetmp(as, ASMREF_TMP1), RID_JGL, -32768);
   tmp = ra_releasetmp(as, ASMREF_TMP2);
   emit_loadi(as, tmp, as->gcsteps);
@@ -1936,7 +1940,8 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
       if (((p[-1] ^ (px-p)) & 0xffffu) == 0 &&
           ((p[-1] & 0xf0000000u) == MIPSI_BEQ ||
            (p[-1] & 0xfc1e0000u) == MIPSI_BLTZ ||
-           (p[-1] & 0xffe00000u) == MIPSI_BC1F)) {
+           (p[-1] & 0xffe00000u) == MIPSI_BC1F) &&
+          p[-2] != MIPS_NOPATCH_GC_CHECK) {
         ptrdiff_t delta = target - p;
         if (((delta + 0x8000) >> 16) == 0) {  /* Patch in-range branch. */
         patchbranch:
diff --git a/src/lj_asm_ppc.h b/src/lj_asm_ppc.h
index dc092db2..32546c78 100644
--- a/src/lj_asm_ppc.h
+++ b/src/lj_asm_ppc.h
@@ -1852,6 +1852,9 @@ static void asm_stack_restore(ASMState *as, SnapShot *snap)
 
 /* -- GC handling --------------------------------------------------------- */
 
+/* Marker to prevent patching the GC check exit. */
+#define PPC_NOPATCH_GC_CHECK    PPCI_ORIS
+
 /* Check GC threshold and do one or more GC steps. */
 static void asm_gc_check(ASMState *as)
 {
@@ -1863,6 +1866,7 @@ static void asm_gc_check(ASMState *as)
   l_end = emit_label(as);
   /* Exit trace if in GCSatomic or GCSfinalize. Avoids syncing GC objects. */
   asm_guardcc(as, CC_NE);  /* Assumes asm_snap_prep() already done. */
+  *--as->mcp = PPC_NOPATCH_GC_CHECK;
   emit_ai(as, PPCI_CMPWI, RID_RET, 0);
   args[0] = ASMREF_TMP1;  /* global_State *g */
   args[1] = ASMREF_TMP2;  /* MSize steps     */
@@ -2124,7 +2128,7 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
   MCode *px = exitstub_trace_addr(T, exitno);
   MCode *cstart = NULL;
   MCode *mcarea = lj_mcode_patch(J, p, 0);
-  int clearso = 0;
+  int clearso = 0, patchlong = 1;
   for (; p < pe; p++) {
     /* Look for exitstub branch, try to replace with branch to target. */
     uint32_t ins = *p;
@@ -2136,7 +2140,9 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
         delta -= sizeof(MCode);
       }
       /* Many, but not all short-range branches can be patched directly. */
-      if (((delta + 0x8000) >> 16) == 0) {
+      if (p[-1] == PPC_NOPATCH_GC_CHECK) {
+        patchlong = 0;
+      } else if (((delta + 0x8000) >> 16) == 0) {
         *p = (ins & 0xffdf0000u) | ((uint32_t)delta & 0xffffu) |
              ((delta & 0x8000) * (PPCF_Y/0x8000));
         if (!cstart) cstart = p;
@@ -2149,7 +2155,8 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
       if (!cstart) cstart = p;
     }
   }
-  {  /* Always patch long-range branch in exit stub itself. */
+  /* Always patch long-range branch in exit stub itself. Except, if we can't. */
+  if (patchlong) {
     ptrdiff_t delta = (char *)target - (char *)px - clearso;
     lua_assert(((delta + 0x02000000) >> 26) == 0);
     *px = PPCI_B | ((uint32_t)delta & 0x03ffffffu);
diff --git a/src/lj_asm_x86.h b/src/lj_asm_x86.h
index 68b40b31..9b61b397 100644
--- a/src/lj_asm_x86.h
+++ b/src/lj_asm_x86.h
@@ -2884,6 +2884,7 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
   MSize len = T->szmcode;
   MCode *px = exitstub_addr(J, exitno) - 6;
   MCode *pe = p+len-6;
+  MCode *pgc = NULL;
   uint32_t stateaddr = u32ptr(&J2G(J)->vmstate);
   if (len > 5 && p[len-5] == XI_JMP && p+len-6 + *(int32_t *)(p+len-4) == px)
     *(int32_t *)(p+len-4) = jmprel(p+len, target);
@@ -2892,9 +2893,15 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
     if (*(uint32_t *)(p+(LJ_64 ? 3 : 2)) == stateaddr && p[0] == XI_MOVmi)
       break;
   lua_assert(p < pe);
-  for (; p < pe; p += asm_x86_inslen(p))
-    if ((*(uint16_t *)p & 0xf0ff) == 0x800f && p + *(int32_t *)(p+2) == px)
+  for (; p < pe; p += asm_x86_inslen(p)) {
+    if ((*(uint16_t *)p & 0xf0ff) == 0x800f && p + *(int32_t *)(p+2) == px &&
+        p != pgc) {
       *(int32_t *)(p+2) = jmprel(p+6, target);
+    } else if (*p == XI_CALL &&
+              (void *)(p+5+*(int32_t *)(p+1)) == (void *)lj_gc_step_jit) {
+      pgc = p+7;  /* Do not patch GC check exit. */
+    }
+  }
   lj_mcode_sync(T->mcode, T->mcode + T->szmcode);
   lj_mcode_patch(J, mcarea, 1);
 }