Ensure forward progress on trace exit to BC_ITERN.

Also use a safer way to force a static dispatch for BC_RET*.
Reported by Bartel Eerdekens. Analyzed by Peter Cawley. #1000 #1045

8 files changed, 146 insertions, 35 deletions

diff --git a/src/lj_trace.c b/src/lj_trace.c
index 03c8d1d0..e019a79f 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -431,6 +431,12 @@ static void trace_start(jit_State *J)
     return;
   }
 
+  /* Ensuring forward progress for BC_ITERN can trigger hotcount again. */
+  if (!J->parent && bc_op(*J->pc) == BC_JLOOP) {  /* Already compiled. */
+    J->state = LJ_TRACE_IDLE;  /* Silently ignored. */
+    return;
+  }
+
   /* Get a new trace number. */
   traceno = trace_findfree(J);
   if (LJ_UNLIKELY(traceno == 0)) {  /* No free trace? */
@@ -867,7 +873,7 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
   ExitDataCP exd;
   int errcode, exitcode = J->exitcode;
   TValue exiterr;
-  const BCIns *pc;
+  const BCIns *pc, *retpc;
   void *cf;
   GCtrace *T;
 
@@ -919,22 +925,7 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
   } else {
     trace_hotside(J, pc);
   }
-  if (bc_op(*pc) == BC_JLOOP) {
-    BCIns *retpc = &traceref(J, bc_d(*pc))->startins;
-    int isret = bc_isret(bc_op(*retpc));
-    if (isret || bc_op(*retpc) == BC_ITERN) {
-      if (J->state == LJ_TRACE_RECORD) {
-        J->patchins = *pc;
-        J->patchpc = (BCIns *)pc;
-        *J->patchpc = *retpc;
-        J->bcskip = 1;
-      } else if (isret) {
-        pc = retpc;
-        setcframe_pc(cf, pc);
-      }
-    }
-  }
-  /* Return MULTRES or 0. */
+  /* Return MULTRES or 0 or -17. */
   ERRNO_RESTORE
   switch (bc_op(*pc)) {
   case BC_CALLM: case BC_CALLMT:
@@ -943,6 +934,18 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
     return (int)((BCReg)(L->top - L->base) + 1 - bc_a(*pc) - bc_d(*pc));
   case BC_TSETM:
     return (int)((BCReg)(L->top - L->base) + 1 - bc_a(*pc));
+  case BC_JLOOP:
+    retpc = &traceref(J, bc_d(*pc))->startins;
+    if (bc_isret(bc_op(*retpc)) || bc_op(*retpc) == BC_ITERN) {
+      /* Dispatch to original ins to ensure forward progress. */
+      if (J->state != LJ_TRACE_RECORD) return -17;
+      /* Unpatch bytecode when recording. */
+      J->patchins = *pc;
+      J->patchpc = (BCIns *)pc;
+      *J->patchpc = *retpc;
+      J->bcskip = 1;
+    }
+    return 0;
   default:
     if (bc_op(*pc) >= BC_FUNCF)
       return (int)((BCReg)(L->top - L->base) + 1);
diff --git a/src/vm_arm.dasc b/src/vm_arm.dasc
index 770c1602..4df4b488 100644
--- a/src/vm_arm.dasc
+++ b/src/vm_arm.dasc
@@ -2196,8 +2196,8 @@ static void build_subroutines(BuildCtx *ctx)
   |.if JIT
   |  ldr L, SAVE_L
   |1:
-  |  cmp CARG1, #0
-  |  blt >9                             // Check for error from exit.
+  |  cmn CARG1, #LUA_ERRERR
+  |  bhs >9                             // Check for error from exit.
   |   lsl RC, CARG1, #3
   |  ldr LFUNC:CARG2, [BASE, FRAME_FUNC]
   |   str RC, SAVE_MULTRES
@@ -2213,6 +2213,8 @@ static void build_subroutines(BuildCtx *ctx)
   |   ldr INS, [PC], #4
   |     lsl MASKR8, MASKR8, #3          // MASKR8 = 255*8.
   |    st_vmstate CARG4
+  |  cmn CARG1, #17                     // Static dispatch?
+  |  beq >5
   |  cmp OP, #BC_FUNCC+2                // Fast function?
   |  bhs >4
   |2:
@@ -2238,6 +2240,17 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldr KBASE, [CARG3, #PC2PROTO(k)]
   |  b <2
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  ldr CARG1, [DISPATCH, #DISPATCH_J(trace)]
+  |  decode_RD RC, INS
+  |  ldr TRACE:CARG1, [CARG1, RC, lsl #2]
+  |  ldr INS, TRACE:CARG1->startins
+  |  decode_OP OP, INS
+  |   decode_RA8 RA, INS
+  |  add OP, DISPATCH, OP, lsl #2
+  |   decode_RD RC, INS
+  |  ldr pc, [OP, #GG_DISP2STATIC]
+  |
   |9:  // Rethrow error from the right C frame.
   |  rsb CARG2, CARG1, #0
   |  mov CARG1, L
diff --git a/src/vm_arm64.dasc b/src/vm_arm64.dasc
index d45cc86b..effb8d91 100644
--- a/src/vm_arm64.dasc
+++ b/src/vm_arm64.dasc
@@ -2005,8 +2005,8 @@ static void build_subroutines(BuildCtx *ctx)
   |.if JIT
   |  ldr L, SAVE_L
   |1:
-  |  cmp CARG1w, #0
-  |  blt >9                             // Check for error from exit.
+  |  cmn CARG1w, #LUA_ERRERR
+  |  bhs >9                             // Check for error from exit.
   |   lsl RC, CARG1, #3
   |  ldr LFUNC:CARG2, [BASE, FRAME_FUNC]
   |    movz TISNUM, #(LJ_TISNUM>>1)&0xffff, lsl #48
@@ -2023,6 +2023,8 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldrb RBw, [PC, # OFS_OP]
   |   ldr INSw, [PC], #4
   |    st_vmstate CARG4w
+  |  cmn CARG1w, #17                    // Static dispatch?
+  |  beq >5
   |  cmp RBw, #BC_FUNCC+2               // Fast function?
   |   add TMP1, GL, INS, uxtb #3
   |  bhs >4
@@ -2033,12 +2035,12 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RA RA, INS
   |   lsr TMP0, INS, #16
   |   csel RC, TMP0, RC, lo
-  |   blo >5
+  |   blo >3
   |   ldr CARG3, [BASE, FRAME_FUNC]
   |   sub RC, RC, #8
   |   add RA, BASE, RA, lsl #3  // Yes: RA = BASE+framesize*8, RC = nargs*8
   |   and LFUNC:CARG3, CARG3, #LJ_GCVMASK
-  |5:
+  |3:
   |  br_auth RB
   |
   |4:  // Check frame below fast function.
@@ -2055,6 +2057,17 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldr KBASE, [CARG3, #PC2PROTO(k)]
   |  b <2
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  ldr RA, [GL, #GL_J(trace)]
+  |  decode_RD RC, INS
+  |  ldr TRACE:RA, [RA, RC, lsl #3]
+  |  ldr INSw, TRACE:RA->startins
+  |  add TMP0, GL, INS, uxtb #3
+  |   decode_RA RA, INS
+  |  ldr RB, [TMP0, #GG_G2DISP+GG_DISP2STATIC]
+  |   decode_RD RC, INS
+  |  br_auth RB
+  |
   |9:  // Rethrow error from the right C frame.
   |  neg CARG2w, CARG1w
   |  mov CARG1, L
diff --git a/src/vm_mips.dasc b/src/vm_mips.dasc
index 34645bf1..bfdcfc1e 100644
--- a/src/vm_mips.dasc
+++ b/src/vm_mips.dasc
@@ -2466,7 +2466,8 @@ static void build_subroutines(BuildCtx *ctx)
   |   addiu DISPATCH, JGL, -GG_DISP2G-32768
   |  sw BASE, L->base
   |1:
-  |  bltz CRET1, >9                     // Check for error from exit.
+  |  sltiu TMP0, CRET1, -LUA_ERRERR     // Check for error from exit.
+  |  beqz TMP0, >9
   |.  lw LFUNC:RB, FRAME_FUNC(BASE)
   |    .FPU lui TMP3, 0x59c0                    // TOBIT = 2^52 + 2^51 (float).
   |  sll MULTRES, CRET1, 3
@@ -2480,14 +2481,16 @@ static void build_subroutines(BuildCtx *ctx)
   |    .FPU cvt.d.s TOBIT, TOBIT
   |  // Modified copy of ins_next which handles function header dispatch, too.
   |  lw INS, 0(PC)
-  |   addiu PC, PC, 4
+  |  addiu CRET1, CRET1, 17             // Static dispatch?
   |    // Assumes TISNIL == ~LJ_VMST_INTERP == -1
   |    sw TISNIL, DISPATCH_GL(vmstate)(DISPATCH)
+  |   decode_RD8a RD, INS
+  |  beqz CRET1, >5
+  |.  addiu PC, PC, 4
   |  decode_OP4a TMP1, INS
   |  decode_OP4b TMP1
-  |    sltiu TMP2, TMP1, BC_FUNCF*4
   |  addu TMP0, DISPATCH, TMP1
-  |   decode_RD8a RD, INS
+  |    sltiu TMP2, TMP1, BC_FUNCF*4
   |  lw AT, 0(TMP0)
   |   decode_RA8a RA, INS
   |    beqz TMP2, >2
@@ -2515,6 +2518,22 @@ static void build_subroutines(BuildCtx *ctx)
   |  jr AT
   |.  addu RA, RA, BASE
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  lw TMP0, DISPATCH_J(trace)(DISPATCH)
+  |  decode_RD4b RD
+  |  addu TMP0, TMP0, RD
+  |  lw TRACE:TMP2, 0(TMP0)
+  |  lw INS, TRACE:TMP2->startins
+  |  decode_OP4a TMP1, INS
+  |  decode_OP4b TMP1
+  |  addu TMP0, DISPATCH, TMP1
+  |   decode_RD8a RD, INS
+  |  lw AT, GG_DISP2STATIC(TMP0)
+  |   decode_RA8a RA, INS
+  |   decode_RD8b RD
+  |  jr AT
+  |.  decode_RA8b RA
+  |
   |9:  // Rethrow error from the right C frame.
   |  load_got lj_err_trace
   |  sub CARG2, r0, CRET1
diff --git a/src/vm_mips64.dasc b/src/vm_mips64.dasc
index 651bc42e..801087b3 100644
--- a/src/vm_mips64.dasc
+++ b/src/vm_mips64.dasc
@@ -2571,7 +2571,8 @@ static void build_subroutines(BuildCtx *ctx)
   |   daddiu DISPATCH, JGL, -GG_DISP2G-32768
   |  sd BASE, L->base
   |1:
-  |  bltz CRET1, >9                     // Check for error from exit.
+  |  sltiu TMP0, CRET1, -LUA_ERRERR     // Check for error from exit.
+  |  beqz TMP0, >9
   |.  ld LFUNC:RB, FRAME_FUNC(BASE)
   |    .FPU lui TMP3, 0x59c0            // TOBIT = 2^52 + 2^51 (float).
   |  dsll MULTRES, CRET1, 3
@@ -2586,14 +2587,16 @@ static void build_subroutines(BuildCtx *ctx)
   |    .FPU cvt.d.s TOBIT, TOBIT
   |  // Modified copy of ins_next which handles function header dispatch, too.
   |  lw INS, 0(PC)
-  |   daddiu PC, PC, 4
+  |  addiu CRET1, CRET1, 17             // Static dispatch?
   |    // Assumes TISNIL == ~LJ_VMST_INTERP == -1
   |    sw TISNIL, DISPATCH_GL(vmstate)(DISPATCH)
+  |   decode_RD8a RD, INS
+  |  beqz CRET1, >5
+  |.  daddiu PC, PC, 4
   |  decode_OP8a TMP1, INS
   |  decode_OP8b TMP1
-  |    sltiu TMP2, TMP1, BC_FUNCF*8
   |  daddu TMP0, DISPATCH, TMP1
-  |   decode_RD8a RD, INS
+  |    sltiu TMP2, TMP1, BC_FUNCF*8
   |  ld AT, 0(TMP0)
   |   decode_RA8a RA, INS
   |    beqz TMP2, >2
@@ -2622,6 +2625,22 @@ static void build_subroutines(BuildCtx *ctx)
   |  jr AT
   |.  daddu RA, RA, BASE
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  ld TMP0, DISPATCH_J(trace)(DISPATCH)
+  |  decode_RD8b RD
+  |  daddu TMP0, TMP0, RD
+  |  ld TRACE:TMP2, 0(TMP0)
+  |  lw INS, TRACE:TMP2->startins
+  |  decode_OP8a TMP1, INS
+  |  decode_OP8b TMP1
+  |  daddu TMP0, DISPATCH, TMP1
+  |   decode_RD8a RD, INS
+  |  ld AT, GG_DISP2STATIC(TMP0)
+  |   decode_RA8a RA, INS
+  |   decode_RD8b RD
+  |  jr AT
+  |.  decode_RA8b RA
+  |
   |9:  // Rethrow error from the right C frame.
   |  load_got lj_err_trace
   |  sub CARG2, r0, CRET1
diff --git a/src/vm_ppc.dasc b/src/vm_ppc.dasc
index 3cad37d2..73d60ae4 100644
--- a/src/vm_ppc.dasc
+++ b/src/vm_ppc.dasc
@@ -3015,8 +3015,9 @@ static void build_subroutines(BuildCtx *ctx)
   |  addi DISPATCH, JGL, -GG_DISP2G-32768
   |  stp BASE, L->base
   |1:
-  |  cmpwi CARG1, 0
-  |  blt >9                             // Check for error from exit.
+  |  li TMP2, -LUA_ERRERR
+  |  cmplw CARG1, TMP2
+  |  bge >9                             // Check for error from exit.
   |  lwz LFUNC:RB, FRAME_FUNC(BASE)
   |   slwi MULTRES, CARG1, 3
   |    li TMP2, 0
@@ -3041,6 +3042,8 @@ static void build_subroutines(BuildCtx *ctx)
   |   addi PC, PC, 4
   |    // Assumes TISNIL == ~LJ_VMST_INTERP == -1.
   |    stw TISNIL, DISPATCH_GL(vmstate)(DISPATCH)
+  |  cmpwi CARG1, -17                   // Static dispatch?
+  |  beq >5
   |  decode_OPP TMP1, INS
   |   decode_RA8 RA, INS
   |  lpx TMP0, DISPATCH, TMP1
@@ -3070,6 +3073,21 @@ static void build_subroutines(BuildCtx *ctx)
   |   add RA, RA, BASE
   |  bctr
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  lwz TMP1, DISPATCH_J(trace)(DISPATCH)
+  |  decode_RD4 RD, INS
+  |  lwzx TRACE:TMP1, TMP1, RD
+  |  lwz INS, TRACE:TMP1->startins
+  |  decode_OPP TMP1, INS
+  |  addi TMP1, TMP1, GG_DISP2STATIC
+  |  lpx TMP0, DISPATCH, TMP1
+  |  mtctr TMP0
+  |   decode_RB8 RB, INS
+  |   decode_RD8 RD, INS
+  |   decode_RA8 RA, INS
+  |   decode_RC8 RC, INS
+  |  bctr
+  |
   |9:  // Rethrow error from the right C frame.
   |  neg CARG2, CARG1
   |  mr CARG1, L
diff --git a/src/vm_x64.dasc b/src/vm_x64.dasc
index 03d96557..5983eeed 100644
--- a/src/vm_x64.dasc
+++ b/src/vm_x64.dasc
@@ -2453,7 +2453,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov r12, [RA]
   |  mov rsp, RA                        // Reposition stack to C frame.
   |.endif
-  |  test RDd, RDd; js >9               // Check for error from exit.
+  |  cmp RDd, -LUA_ERRERR; jae >9       // Check for error from exit.
   |  mov L:RB, SAVE_L
   |  mov MULTRES, RDd
   |  mov LFUNC:KBASE, [BASE-16]
@@ -2469,6 +2469,8 @@ static void build_subroutines(BuildCtx *ctx)
   |  movzx OP, RCL
   |  add PC, 4
   |  shr RCd, 16
+  |  cmp MULTRES, -17                   // Static dispatch?
+  |  je >5
   |  cmp OP, BC_FUNCF                   // Function header?
   |  jb >3
   |  cmp OP, BC_FUNCC+2                 // Fast function?
@@ -2491,6 +2493,15 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov KBASE, [KBASE+PC2PROTO(k)]
   |  jmp <2
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  mov RA, [DISPATCH+DISPATCH_J(trace)]
+  |  mov TRACE:RA, [RA+RD*8]
+  |  mov RCd, TRACE:RA->startins
+  |  movzx RAd, RCH
+  |  movzx OP, RCL
+  |  shr RCd, 16
+  |  jmp aword [DISPATCH+OP*8+GG_DISP2STATIC]
+  |
   |9:  // Rethrow error from the right C frame.
   |  mov CARG2d, RDd
   |  mov CARG1, L:RB
diff --git a/src/vm_x86.dasc b/src/vm_x86.dasc
index 18ca87b5..f7847762 100644
--- a/src/vm_x86.dasc
+++ b/src/vm_x86.dasc
@@ -2902,7 +2902,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov r13, TMPa
   |  mov r12, TMPQ
   |.endif
-  |  test RD, RD; js >9                 // Check for error from exit.
+  |  cmp RD, -LUA_ERRERR; jae >9        // Check for error from exit.
   |  mov L:RB, SAVE_L
   |  mov MULTRES, RD
   |  mov LFUNC:KBASE, [BASE-8]
@@ -2917,6 +2917,8 @@ static void build_subroutines(BuildCtx *ctx)
   |  movzx OP, RCL
   |  add PC, 4
   |  shr RC, 16
+  |  cmp MULTRES, -17                   // Static dispatch?
+  |  je >5
   |  cmp OP, BC_FUNCF                   // Function header?
   |  jb >3
   |  cmp OP, BC_FUNCC+2                 // Fast function?
@@ -2942,6 +2944,19 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov KBASE, [KBASE+PC2PROTO(k)]
   |  jmp <2
   |
+  |5:  // Dispatch to static entry of original ins replaced by BC_JLOOP.
+  |  mov RA, [DISPATCH+DISPATCH_J(trace)]
+  |  mov TRACE:RA, [RA+RD*4]
+  |  mov RC, TRACE:RA->startins
+  |  movzx RA, RCH
+  |  movzx OP, RCL
+  |  shr RC, 16
+  |.if X64
+  |  jmp aword [DISPATCH+OP*8+GG_DISP2STATIC]
+  |.else
+  |  jmp aword [DISPATCH+OP*4+GG_DISP2STATIC]
+  |.endif
+  |
   |9:  // Rethrow error from the right C frame.
   |  mov FCARG2, RD
   |  mov FCARG1, L:RB