Fix trace stitching by storing trace object itself in the continuation.

src/lj_ffrecord.c

@@ -96,18 +96,10 @@ static ptrdiff_t results_wanted(jit_State *J)
     return -1;
 }
 
-#ifdef LUAJIT_TRACE_STITCHING
-/* This feature is disabled for now due to a design mistake. Sorry.
-**
-** It causes unpredictable behavior and crashes when a full trace flush
-** happens with a stitching continuation still in the stack somewhere.
-*/
-
 /* Trace stitching: add continuation below frame to start a new trace. */
 static void recff_stitch(jit_State *J)
 {
   ASMFunction cont = lj_cont_stitch;
-  TraceNo traceno = J->cur.traceno;
   lua_State *L = J->L;
   TValue *base = L->base;
   const BCIns *pc = frame_pc(base-1);
@@ -120,7 +112,7 @@ static void recff_stitch(jit_State *J)
   setframe_ftsz(base+1, ((char *)(base+1) - (char *)pframe) + FRAME_CONT);
   setcont(base, cont);
   setframe_pc(base, pc);
-  if (LJ_DUALNUM) setintV(base-1, traceno); else base[-1].u64 = traceno;
+  setnilV(base-1);
   L->base += 2;
   L->top += 2;
 
@@ -131,8 +123,9 @@ static void recff_stitch(jit_State *J)
 #else
   trcont = lj_ir_kptr(J, (void *)cont);
 #endif
+  J->selfref = lj_ir_k64_reserve(J);
   J->base[0] = trcont | TREF_CONT;
-  J->base[-1] = LJ_DUALNUM ? lj_ir_kint(J,traceno) : lj_ir_knum_u64(J,traceno);
+  J->base[-1] = emitir(IRT(IR_XLOAD, IRT_P64), lj_ir_kptr(J, J->selfref), 0);
   J->base += 2;
   J->baseslot += 2;
   J->framedepth++;
@@ -181,31 +174,6 @@ static void LJ_FASTCALL recff_nyi(jit_State *J, RecordFFData *rd)
 
 /* Must stop the trace for classic C functions with arbitrary side-effects. */
 #define recff_c		recff_nyi
-#else
-/* Fallback handler for fast functions that are not recorded (yet). */
-static void LJ_FASTCALL recff_nyi(jit_State *J, RecordFFData *rd)
-{
-  setfuncV(J->L, &J->errinfo, J->fn);
-  lj_trace_err_info(J, LJ_TRERR_NYIFF);
-  UNUSED(rd);
-}
-
-/* Throw error for unsupported variant of fast function. */
-LJ_NORET static void recff_nyiu(jit_State *J, RecordFFData *rd)
-{
-  setfuncV(J->L, &J->errinfo, J->fn);
-  lj_trace_err_info(J, LJ_TRERR_NYIFFU);
-  UNUSED(rd);
-}
-
-/* Must abort the trace for classic C functions with arbitrary side-effects. */
-static void LJ_FASTCALL recff_c(jit_State *J, RecordFFData *rd)
-{
-  setfuncV(J->L, &J->errinfo, J->fn);
-  lj_trace_err_info(J, LJ_TRERR_NYICF);
-  UNUSED(rd);
-}
-#endif
 
 /* Emit BUFHDR for the global temporary buffer. */
 static TRef recff_bufhdr(jit_State *J)

src/lj_gc.c

@@ -69,6 +69,10 @@ static void gc_mark(global_State *g, GCobj *o)
       gray2black(o);  /* Closed upvalues are never gray. */
   } else if (gct != ~LJ_TSTR && gct != ~LJ_TCDATA) {
     lua_assert(gct == ~LJ_TFUNC || gct == ~LJ_TTAB ||
+#if LJ_HASJIT
+               /* Can occur as a part of trace stitching continuation. */
+               gct == ~LJ_TTRACE ||
+#endif
 	       gct == ~LJ_TTHREAD || gct == ~LJ_TPROTO);
     setgcrefr(o->gch.gclist, g->gc.gray);
     setgcref(g->gc.gray, o);

src/lj_ir.c

@@ -209,13 +209,46 @@ void lj_ir_k64_freeall(jit_State *J)
     lj_mem_free(J2G(J), k, sizeof(K64Array));
     k = next;
   }
+  setmref(J->k64, NULL);
+}
+
+/* Add 64 bit constant to the last K64Array or add a new array. */
+static TValue *ir_k64_add(jit_State *J, K64Array *kl, uint64_t u64)
+{
+  TValue *ntv;
+  /* Constant was not found, need to add it. */
+  if (!(kl && kl->numk < LJ_MIN_K64SZ)) {  /* Allocate a new array. */
+    K64Array *kn = lj_mem_newt(J->L, sizeof(K64Array), K64Array);
+    setmref(kn->next, NULL);
+    kn->numk = 0;
+    if (kl)
+      setmref(kl->next, kn);  /* Chain to the end of the list. */
+    else
+      setmref(J->k64, kn);  /* Link first array. */
+    kl = kn;
+  }
+  ntv = &kl->k[kl->numk++];  /* Add to current array. */
+  ntv->u64 = u64;
+  return ntv;
+}
+
+/* Reserve space for a 64 bit constant in chained array. */
+TValue *lj_ir_k64_reserve(jit_State *J)
+{
+  K64Array *k, *kl = NULL;
+  /* Intern 0 to ensure that no code can by mistake reuse mutable slot. */
+  lj_ir_k64_find(J, 0);
+  /* Search for the last element in the chain of arrays. */
+  for (k = mref(J->k64, K64Array); k; k = mref(k->next, K64Array)) {
+    kl = k;
+  }
+  return ir_k64_add(J, kl, 0);
 }
 
 /* Find 64 bit constant in chained array or add it. */
 cTValue *lj_ir_k64_find(jit_State *J, uint64_t u64)
 {
   K64Array *k, *kp = NULL;
-  TValue *ntv;
   MSize idx;
   /* Search for the constant in the whole chain of arrays. */
   for (k = mref(J->k64, K64Array); k; k = mref(k->next, K64Array)) {
@@ -226,20 +259,7 @@ cTValue *lj_ir_k64_find(jit_State *J, uint64_t u64)
 	return tv;
     }
   }
-  /* Constant was not found, need to add it. */
+  return ir_k64_add(J, kp, u64);
-  if (!(kp && kp->numk < LJ_MIN_K64SZ)) {  /* Allocate a new array. */
-    K64Array *kn = lj_mem_newt(J->L, sizeof(K64Array), K64Array);
-    setmref(kn->next, NULL);
-    kn->numk = 0;
-    if (kp)
-      setmref(kp->next, kn);  /* Chain to the end of the list. */
-    else
-      setmref(J->k64, kn);  /* Link first array. */
-    kp = kn;
-  }
-  ntv = &kp->k[kp->numk++];  /* Add to current array. */
-  ntv->u64 = u64;
-  return ntv;
 }
 
 /* Intern 64 bit constant, given by its address. */

src/lj_iropt.h

@@ -40,6 +40,7 @@ static LJ_AINLINE IRRef lj_ir_nextins(jit_State *J)
 LJ_FUNC TRef LJ_FASTCALL lj_ir_kint(jit_State *J, int32_t k);
 LJ_FUNC void lj_ir_k64_freeall(jit_State *J);
 LJ_FUNC TRef lj_ir_k64(jit_State *J, IROp op, cTValue *tv);
+LJ_FUNC TValue *lj_ir_k64_reserve(jit_State *J);
 LJ_FUNC cTValue *lj_ir_k64_find(jit_State *J, uint64_t u64);
 LJ_FUNC TRef lj_ir_knum_u64(jit_State *J, uint64_t u64);
 LJ_FUNC TRef lj_ir_knumint(jit_State *J, lua_Number n);

src/lj_jit.h

@@ -418,6 +418,8 @@ typedef struct jit_State {
 
   TValue errinfo;	/* Additional info element for trace errors. */
 
+  TValue* selfref;  /* Slot for the trace self reference. */
+
 #if LJ_HASPROFILE
   GCproto *prev_pt;	/* Previous prototype. */
   BCLine prev_line;	/* Previous line. */

src/lj_trace.c

@@ -274,7 +274,7 @@ int lj_trace_flushall(lua_State *L)
       if (T->root == 0)
 	trace_flushroot(J, T);
       lj_gdbjit_deltrace(J, T);
-      T->traceno = 0;
+      T->traceno = T->link = 0;  /* Blacklist the link for cont_stitch. */
       setgcrefnull(J->trace[i]);
     }
   }
@@ -284,6 +284,7 @@ int lj_trace_flushall(lua_State *L)
   memset(J->penalty, 0, sizeof(J->penalty));
   /* Free the whole machine code and invalidate all exit stub groups. */
   lj_mcode_free(J);
+  lj_ir_k64_freeall(J);
   memset(J->exitstubgroup, 0, sizeof(J->exitstubgroup));
   lj_vmevent_send(L, TRACE,
     setstrV(L, L->top++, lj_str_newlit(L, "flush"));
@@ -403,6 +404,7 @@ static void trace_start(jit_State *J)
   lj_resetsplit(J);
   J->retryrec = 0;
   setgcref(J->cur.startpt, obj2gco(J->pt));
+  J->selfref = NULL;
 
   L = J->L;
   lj_vmevent_send(L, TRACE,
@@ -426,7 +428,7 @@ static void trace_stop(jit_State *J)
   GCproto *pt = &gcref(J->cur.startpt)->pt;
   TraceNo traceno = J->cur.traceno;
   GCtrace *T = trace_save_alloc(J);  /* Do this first. May throw OOM. */
-  lua_State *L;
+  lua_State *L = J->L;
 
   switch (op) {
   case BC_FORL:
@@ -477,8 +479,10 @@ static void trace_stop(jit_State *J)
   lj_mcode_commit(J, J->cur.mcode);
   J->postproc = LJ_POST_NONE;
   trace_save(J, T);
+  if (J->cur.linktype == LJ_TRLINK_STITCH) {
+    setgcV(L, J->selfref, obj2gco(T), LJ_TTRACE);
+  }
 
-  L = J->L;
   lj_vmevent_send(L, TRACE,
     setstrV(L, L->top++, lj_str_newlit(L, "stop"));
     setintV(L->top++, traceno);

src/vm_arm.dasc

@@ -2086,7 +2086,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  // RA = resultptr, CARG4 = meta base
   |   ldr RB, SAVE_MULTRES
   |  ldr INS, [PC, #-4]
-  |    ldr CARG3, [CARG4, #-24]		// Save previous trace number.
+  |    ldr TRACE:CARG3, [CARG4, #-24]		// Save previous trace.
   |   subs RB, RB, #8
   |  decode_RA8 RC, INS			// Call base.
   |   beq >2
@@ -2101,23 +2101,20 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RA8 RA, INS
   |   decode_RB8 RB, INS
   |   add RA, RA, RB
-  |  ldr CARG1, [DISPATCH, #DISPATCH_J(trace)]
   |3:
   |   cmp RA, RC
   |  mvn CARG2, #~LJ_TNIL
   |   bhi >9				// More results wanted?
   |
-  |  ldr TRACE:RA, [CARG1, CARG3, lsl #2]
+  |  ldrh RA, TRACE:CARG3->traceno
-  |  cmp TRACE:RA, #0
+  |  ldrh RC, TRACE:CARG3->link
-  |  beq ->cont_nop
+  |  cmp RC, RA
-  |  ldrh RC, TRACE:RA->link
-  |  cmp RC, CARG3
   |  beq ->cont_nop			// Blacklisted.
   |  cmp RC, #0
   |  bne =>BC_JLOOP			// Jump to stitched trace.
   |
   |  // Stitch a new trace to the previous trace.
-  |  str CARG3, [DISPATCH, #DISPATCH_J(exitno)]
+  |  str RA, [DISPATCH, #DISPATCH_J(exitno)]
   |  str L, [DISPATCH, #DISPATCH_J(L)]
   |  str BASE, L->base
   |  sub CARG1, DISPATCH, #-GG_DISP2J

src/vm_mips.dasc

@@ -2015,7 +2015,7 @@ static void build_subroutines(BuildCtx *ctx)
   |.if JIT
   |  // RA = resultptr, RB = meta base
   |  lw INS, -4(PC)
-  |    lw TMP3, -24+LO(RB)		// Save previous trace number.
+  |    lw TMP2, -24+LO(RB)		// Save previous trace.
   |  decode_RA8a RC, INS
   |   addiu AT, MULTRES, -8
   |  decode_RA8b RC
@@ -2034,17 +2034,13 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RA8b RA
   |    decode_RB8b RB
   |   addu RA, RA, RB
-  |  lw TMP1, DISPATCH_J(trace)(DISPATCH)
   |   addu RA, BASE, RA
   |3:
   |   sltu AT, RC, RA
   |   bnez AT, >9			// More results wanted?
-  |. sll TMP2, TMP3, 2
+  |.   nop
   |
-  |  addu TMP2, TMP1, TMP2
+  |  lhu TMP3, TRACE:TMP2->traceno
-  |  lw TRACE:TMP2, 0(TMP2)
-  |  beqz TRACE:TMP2, ->cont_nop
-  |.  nop
   |  lhu RD, TRACE:TMP2->link
   |  beq RD, TMP3, ->cont_nop		// Blacklisted.
   |.  load_got lj_dispatch_stitch

src/vm_ppc.dasc

@@ -2525,7 +2525,7 @@ static void build_subroutines(BuildCtx *ctx)
   |.if JIT
   |  // RA = resultptr, RB = meta base
   |  lwz INS, -4(PC)
-  |    lwz TMP3, -20(RB)		// Save previous trace number.
+  |    lwz TRACE:TMP2, -20(RB)		// Save previous trace.
   |   addic. TMP1, MULTRES, -8
   |  decode_RA8 RC, INS			// Call base.
   |   beq >2
@@ -2540,15 +2540,11 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RA8 RA, INS
   |   decode_RB8 RB, INS
   |   add RA, RA, RB
-  |  lwz TMP1, DISPATCH_J(trace)(DISPATCH)
   |3:
   |   cmplw RA, RC
   |   bgt >9				// More results wanted?
   |
-  |  slwi TMP2, TMP3, 2
+  |  lhz TMP3, TRACE:TMP2->traceno
-  |  lwzx TRACE:TMP2, TMP1, TMP2
-  |  cmpwi TRACE:TMP2, 0
-  |  beq ->cont_nop
   |  lhz RD, TRACE:TMP2->link
   |  cmpw RD, TMP3
   |   cmpwi cr1, RD, 0

src/vm_x86.dasc

@@ -2667,7 +2667,7 @@ static void build_subroutines(BuildCtx *ctx)
   |->cont_stitch:			// Trace stitching.
   |.if JIT
   |  // BASE = base, RC = result, RB = mbase
-  |  mov RA, [RB-24]			// Save previous trace number.
+  |  mov RA, [RB-24]			// Save previous trace.
   |  mov TMP1, RA
   |  mov TMP3, DISPATCH			// Need one more register.
   |  mov DISPATCH, MULTRES
@@ -2699,11 +2699,8 @@ static void build_subroutines(BuildCtx *ctx)
   |  ja >9				// More results wanted?
   |
   |  mov DISPATCH, TMP3
-  |  mov RB, TMP1			// Get previous trace number.
+  |  mov TRACE:RD, TMP1			// Get previous trace.
-  |  mov RA, [DISPATCH+DISPATCH_J(trace)]
+  |  movzx RB, word TRACE:RD->traceno
-  |  mov TRACE:RD, [RA+RB*4]
-  |  test TRACE:RD, TRACE:RD
-  |  jz ->cont_nop
   |  movzx RD, word TRACE:RD->link
   |  cmp RD, RB
   |  je ->cont_nop			// Blacklisted.