5 files changed, 33 insertions, 37 deletions

diff --git a/src/lj_dispatch.c b/src/lj_dispatch.c
index b83912c5..f1579c8d 100644
--- a/src/lj_dispatch.c
+++ b/src/lj_dispatch.c
@@ -359,10 +359,7 @@ void LJ_FASTCALL lj_dispatch_ins(lua_State *L, const BCIns *pc)
     jit_State *J = G2J(g);
     if (J->state != LJ_TRACE_IDLE) {
       J->L = L;
-      J->pc = pc-1;
-      J->fn = fn;
-      J->pt = pt;
-      lj_trace_ins(J);
+      lj_trace_ins(J, pc-1);  /* The interpreter bytecode PC is offset by 1. */
     }
   }
 #endif
diff --git a/src/lj_snap.c b/src/lj_snap.c
index 547b2089..04b9a7f6 100644
--- a/src/lj_snap.c
+++ b/src/lj_snap.c
@@ -179,7 +179,7 @@ void lj_snap_regspmap(uint16_t *rsmap, Trace *T, SnapNo snapno)
 }
 
 /* Restore interpreter state from exit state with the help of a snapshot. */
-void lj_snap_restore(jit_State *J, void *exptr)
+const BCIns *lj_snap_restore(jit_State *J, void *exptr)
 {
   ExitState *ex = (ExitState *)exptr;
   SnapNo snapno = J->exitno;  /* For now, snapno == exitno. */
@@ -187,7 +187,7 @@ void lj_snap_restore(jit_State *J, void *exptr)
   SnapShot *snap = &T->snap[snapno];
   MSize n, nent = snap->nent;
   SnapEntry *map = &T->snapmap[snap->mapofs];
-  SnapEntry *flinks = map + nent;
+  SnapEntry *flinks = map + nent + 1;
   int32_t ftsz0;
   BCReg nslots = snap->nslots;
   TValue *frame;
@@ -201,7 +201,6 @@ void lj_snap_restore(jit_State *J, void *exptr)
   }
 
   /* Fill stack slots with data from the registers and spill slots. */
-  J->pc = snap_pc(*flinks++);
   frame = L->base-1;
   ftsz0 = frame_ftsz(frame);  /* Preserve link to previous frame in slot #0. */
   for (n = 0; n < nent; n++) {
@@ -266,6 +265,7 @@ void lj_snap_restore(jit_State *J, void *exptr)
   }
   L->top = curr_topL(L);
   lua_assert(map + nent + 1 + snap->depth == flinks);
+  return snap_pc(map[nent]);
 }
 
 #undef IR
diff --git a/src/lj_snap.h b/src/lj_snap.h
index ed7d98a1..80be33b7 100644
--- a/src/lj_snap.h
+++ b/src/lj_snap.h
@@ -13,7 +13,7 @@
 LJ_FUNC void lj_snap_add(jit_State *J);
 LJ_FUNC void lj_snap_shrink(jit_State *J);
 LJ_FUNC void lj_snap_regspmap(uint16_t *rsmap, Trace *T, SnapNo snapno);
-LJ_FUNC void lj_snap_restore(jit_State *J, void *exptr);
+LJ_FUNC const BCIns *lj_snap_restore(jit_State *J, void *exptr);
 LJ_FUNC void lj_snap_grow_buf_(jit_State *J, MSize need);
 LJ_FUNC void lj_snap_grow_map_(jit_State *J, MSize need);
 
diff --git a/src/lj_trace.c b/src/lj_trace.c
index 0a973d51..b5946346 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -532,46 +532,52 @@ static TValue *trace_state(lua_State *L, lua_CFunction dummy, void *ud)
 /* -- Event handling ------------------------------------------------------ */
 
 /* A bytecode instruction is about to be executed. Record it. */
-void lj_trace_ins(jit_State *J)
+void lj_trace_ins(jit_State *J, const BCIns *pc)
 {
-  while (lj_vm_cpcall(J->L, NULL, (void *)J, trace_state) != 0)
-    J->state = LJ_TRACE_ERR;
-}
-
-/* Start recording a new trace. */
-static void trace_new(jit_State *J)
-{
-  lua_assert(J->state == LJ_TRACE_IDLE);
-  J->state = LJ_TRACE_START;
+  /* Note: J->L must already be set. pc is the true bytecode PC here. */
+  J->pc = pc;
   J->fn = curr_func(J->L);
   J->pt = funcproto(J->fn);
-  lj_trace_ins(J);
+  while (lj_vm_cpcall(J->L, NULL, (void *)J, trace_state) != 0)
+    J->state = LJ_TRACE_ERR;
 }
 
 /* A hotcount triggered. Start recording a root trace. */
 void LJ_FASTCALL lj_trace_hot(jit_State *J, const BCIns *pc)
 {
+  /* Note: pc is the interpreter bytecode PC here. It's offset by 1. */
   hotcount_set(J2GG(J), pc, J->param[JIT_P_hotloop]+1);  /* Reset hotcount. */
   /* Only start a new trace if not recording or inside __gc call or vmevent. */
   if (J->state == LJ_TRACE_IDLE &&
       !(J2G(J)->hookmask & (HOOK_GC|HOOK_VMEVENT))) {
     J->parent = 0;  /* Root trace. */
     J->exitno = 0;
-    J->pc = pc-1;  /* The interpreter bytecode PC is offset by 1. */
-    trace_new(J);
+    J->state = LJ_TRACE_START;
+    lj_trace_ins(J, pc-1);
   }
 }
 
-/* A trace exited. Restore interpreter state and check for hot exits. */
+/* Check for a hot side exit. If yes, start recording a side trace. */
+static void trace_hotside(jit_State *J, const BCIns *pc)
+{
+  SnapShot *snap = &J->trace[J->parent]->snap[J->exitno];
+  if (!(J2G(J)->hookmask & (HOOK_GC|HOOK_VMEVENT)) &&
+      snap->count != SNAPCOUNT_DONE &&
+      ++snap->count >= J->param[JIT_P_hotexit]) {
+    lua_assert(J->state == LJ_TRACE_IDLE);
+    /* J->parent is non-zero for a side trace. */
+    J->state = LJ_TRACE_START;
+    lj_trace_ins(J, pc);
+  }
+}
+
+/* A trace exited. Restore interpreter state. */
 void * LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
 {
+  const BCIns *pc = lj_snap_restore(J, exptr);
   lua_State *L = J->L;
-  void *cf;
-
-  /* Restore interpreter state. */
-  lj_snap_restore(J, exptr);
-  cf = cframe_raw(L->cframe);
-  setcframe_pc(cf, J->pc);
+  void *cf = cframe_raw(L->cframe);
+  setcframe_pc(cf, pc);  /* Restart interpreter at this PC. */
 
   lj_vmevent_send(L, TEXIT,
     ExitState *ex = (ExitState *)exptr;
@@ -591,14 +597,7 @@ void * LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
     }
   );
 
-  {  /* Check for a hot exit. */
-    SnapShot *snap = &J->trace[J->parent]->snap[J->exitno];
-    if (!(J2G(J)->hookmask & (HOOK_GC|HOOK_VMEVENT)) &&
-        snap->count != SNAPCOUNT_DONE &&
-        ++snap->count >= J->param[JIT_P_hotexit])
-      trace_new(J);  /* Start recording a side trace. */
-  }
-
+  trace_hotside(J, pc);
   return cf;  /* Return the interpreter C frame. */
 }
 
diff --git a/src/lj_trace.h b/src/lj_trace.h
index 0d4e30e0..a8e45974 100644
--- a/src/lj_trace.h
+++ b/src/lj_trace.h
@@ -31,7 +31,7 @@ LJ_FUNC int lj_trace_flushall(lua_State *L);
 LJ_FUNC void lj_trace_freestate(global_State *g);
 
 /* Event handling. */
-LJ_FUNC void lj_trace_ins(jit_State *J);
+LJ_FUNC void lj_trace_ins(jit_State *J, const BCIns *pc);
 LJ_FUNCA void LJ_FASTCALL lj_trace_hot(jit_State *J, const BCIns *pc);
 LJ_FUNCA void * LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr);