1 files changed, 37 insertions, 27 deletions

diff --git a/src/lj_asm_arm64.h b/src/lj_asm_arm64.h
index cbb186d3..baafa21a 100644
--- a/src/lj_asm_arm64.h
+++ b/src/lj_asm_arm64.h
@@ -56,11 +56,11 @@ static void asm_exitstub_setup(ASMState *as, ExitNo nexits)
     asm_mclimit(as);
   /* 1: str lr,[sp]; bl ->vm_exit_handler; movz w0,traceno; bl <1; bl <1; ... */
   for (i = nexits-1; (int32_t)i >= 0; i--)
-    *--mxp = A64I_LE(A64I_BL|((-3-i)&0x03ffffffu));
-  *--mxp = A64I_LE(A64I_MOVZw|A64F_U16(as->T->traceno));
+    *--mxp = A64I_LE(A64I_BL | A64F_S26(-3-i));
+  *--mxp = A64I_LE(A64I_MOVZw | A64F_U16(as->T->traceno));
   mxp--;
-  *mxp = A64I_LE(A64I_BL|(((MCode *)(void *)lj_vm_exit_handler-mxp)&0x03ffffffu));
-  *--mxp = A64I_LE(A64I_STRx|A64F_D(RID_LR)|A64F_N(RID_SP));
+  *mxp = A64I_LE(A64I_BL | A64F_S26(((MCode *)(void *)lj_vm_exit_handler-mxp)));
+  *--mxp = A64I_LE(A64I_STRx | A64F_D(RID_LR) | A64F_N(RID_SP));
   as->mctop = mxp;
 }
 
@@ -77,7 +77,7 @@ static void asm_guardcc(ASMState *as, A64CC cc)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_cond_branch(as, cc^1, p-1);
     return;
   }
@@ -91,7 +91,7 @@ static void asm_guardtnb(ASMState *as, A64Ins ai, Reg r, uint32_t bit)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_tnb(as, ai^0x01000000u, r, bit, p-1);
     return;
   }
@@ -105,7 +105,7 @@ static void asm_guardcnb(ASMState *as, A64Ins ai, Reg r)
   MCode *p = as->mcp;
   if (LJ_UNLIKELY(p == as->invmcp)) {
     as->loopinv = 1;
-    *p = A64I_B | ((target-p) & 0x03ffffffu);
+    *p = A64I_B | A64F_S26(target-p);
     emit_cnb(as, ai^0x01000000u, r, p-1);
     return;
   }
@@ -1850,7 +1850,7 @@ static void asm_loop_fixup(ASMState *as)
     p[-2] |= ((uint32_t)delta & mask) << 5;
   } else {
     ptrdiff_t delta = target - (p - 1);
-    p[-1] = A64I_B | ((uint32_t)(delta) & 0x03ffffffu);
+    p[-1] = A64I_B | A64F_S26(delta);
   }
 }
 
@@ -1919,7 +1919,7 @@ static void asm_tail_fixup(ASMState *as, TraceNo lnk)
   }
   /* Patch exit branch. */
   target = lnk ? traceref(as->J, lnk)->mcode : (MCode *)lj_vm_exit_interp;
-  p[-1] = A64I_B | (((target-p)+1)&0x03ffffffu);
+  p[-1] = A64I_B | A64F_S26((target-p)+1);
 }
 
 /* Prepare tail of code. */
@@ -1982,40 +1982,50 @@ void lj_asm_patchexit(jit_State *J, GCtrace *T, ExitNo exitno, MCode *target)
 {
   MCode *p = T->mcode;
   MCode *pe = (MCode *)((char *)p + T->szmcode);
-  MCode *cstart = NULL, *cend = p;
+  MCode *cstart = NULL;
   MCode *mcarea = lj_mcode_patch(J, p, 0);
   MCode *px = exitstub_trace_addr(T, exitno);
+  /* Note: this assumes a trace exit is only ever patched once. */
   for (; p < pe; p++) {
     /* Look for exitstub branch, replace with branch to target. */
+    ptrdiff_t delta = target - p;
     MCode ins = A64I_LE(*p);
     if ((ins & 0xff000000u) == 0x54000000u &&
         ((ins ^ ((px-p)<<5)) & 0x00ffffe0u) == 0) {
-      /* Patch bcc exitstub. */
-      *p = A64I_LE((ins & 0xff00001fu) | (((target-p)<<5) & 0x00ffffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch bcc, if within range. */
+      if (A64F_S_OK(delta, 19)) {
+        *p = A64I_LE((ins & 0xff00001fu) | A64F_S19(delta));
+        if (!cstart) cstart = p;
+      }
     } else if ((ins & 0xfc000000u) == 0x14000000u &&
                ((ins ^ (px-p)) & 0x03ffffffu) == 0) {
-      /* Patch b exitstub. */
-      *p = A64I_LE((ins & 0xfc000000u) | ((target-p) & 0x03ffffffu));
-      cend = p+1;
+      /* Patch b. */
+      lua_assert(A64F_S_OK(delta, 26));
+      *p = A64I_LE((ins & 0xfc000000u) | A64F_S26(delta));
       if (!cstart) cstart = p;
     } else if ((ins & 0x7e000000u) == 0x34000000u &&
                ((ins ^ ((px-p)<<5)) & 0x00ffffe0u) == 0) {
-      /* Patch cbz/cbnz exitstub. */
-      *p = A64I_LE((ins & 0xff00001f) | (((target-p)<<5) & 0x00ffffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch cbz/cbnz, if within range. */
+      if (A64F_S_OK(delta, 19)) {
+        *p = A64I_LE((ins & 0xff00001fu) | A64F_S19(delta));
+        if (!cstart) cstart = p;
+      }
     } else if ((ins & 0x7e000000u) == 0x36000000u &&
                ((ins ^ ((px-p)<<5)) & 0x0007ffe0u) == 0) {
-      /* Patch tbz/tbnz exitstub. */
-      *p = A64I_LE((ins & 0xfff8001fu) | (((target-p)<<5) & 0x0007ffe0u));
-      cend = p+1;
-      if (!cstart) cstart = p;
+      /* Patch tbz/tbnz, if within range. */
+      if (A64F_S_OK(delta, 14)) {
+        *p = A64I_LE((ins & 0xfff8001fu) | A64F_S14(delta));
+        if (!cstart) cstart = p;
+      }
     }
   }
-  lua_assert(cstart != NULL);
-  lj_mcode_sync(cstart, cend);
+  {  /* Always patch long-range branch in exit stub itself. */
+    ptrdiff_t delta = target - px;
+    lua_assert(A64F_S_OK(delta, 26));
+    *px = A64I_B | A64F_S26(delta);
+    if (!cstart) cstart = px;
+  }
+  lj_mcode_sync(cstart, px+1);
   lj_mcode_patch(J, mcarea, 1);
 }