7 files changed, 25 insertions, 17 deletions

diff --git a/src/lj_asm_arm.h b/src/lj_asm_arm.h
index d4f1113a..0ab4917e 100644
--- a/src/lj_asm_arm.h
+++ b/src/lj_asm_arm.h
@@ -1460,8 +1460,8 @@ static void asm_stack_check(ASMState *as, BCReg topslot,
 static void asm_stack_restore(ASMState *as, SnapShot *snap)
 {
   SnapEntry *map = &as->T->snapmap[snap->mapofs];
+  SnapEntry *flinks = &as->T->snapmap[snap_nextofs(as->T, snap)-1];
   MSize n, nent = snap->nent;
-  SnapEntry *flinks = map + nent + snap->depth;
   /* Store the value of all modified slots to the Lua stack. */
   for (n = 0; n < nent; n++) {
     SnapEntry sn = map[n];
diff --git a/src/lj_asm_ppc.h b/src/lj_asm_ppc.h
index 196ca2ed..c03d0102 100644
--- a/src/lj_asm_ppc.h
+++ b/src/lj_asm_ppc.h
@@ -1793,8 +1793,8 @@ static void asm_stack_check(ASMState *as, BCReg topslot,
 static void asm_stack_restore(ASMState *as, SnapShot *snap)
 {
   SnapEntry *map = &as->T->snapmap[snap->mapofs];
+  SnapEntry *flinks = &as->T->snapmap[snap_nextofs(as->T, snap)-1];
   MSize n, nent = snap->nent;
-  SnapEntry *flinks = map + nent + snap->depth;
   /* Store the value of all modified slots to the Lua stack. */
   for (n = 0; n < nent; n++) {
     SnapEntry sn = map[n];
diff --git a/src/lj_asm_x86.h b/src/lj_asm_x86.h
index 391e2de9..4d47e389 100644
--- a/src/lj_asm_x86.h
+++ b/src/lj_asm_x86.h
@@ -2316,8 +2316,8 @@ static void asm_stack_check(ASMState *as, BCReg topslot,
 static void asm_stack_restore(ASMState *as, SnapShot *snap)
 {
   SnapEntry *map = &as->T->snapmap[snap->mapofs];
+  SnapEntry *flinks = &as->T->snapmap[snap_nextofs(as->T, snap)-1];
   MSize n, nent = snap->nent;
-  SnapEntry *flinks = map + nent + snap->depth;
   /* Store the value of all modified slots to the Lua stack. */
   for (n = 0; n < nent; n++) {
     SnapEntry sn = map[n];
diff --git a/src/lj_jit.h b/src/lj_jit.h
index 8f87899c..e80547ab 100644
--- a/src/lj_jit.h
+++ b/src/lj_jit.h
@@ -139,8 +139,8 @@ typedef struct SnapShot {
   IRRef1 ref;           /* First IR ref for this snapshot. */
   uint8_t nslots;       /* Number of valid slots. */
   uint8_t nent;         /* Number of compressed entries. */
-  uint8_t depth;        /* Number of frame links. */
   uint8_t count;        /* Count of taken exits for this snapshot. */
+  uint8_t unused;
 } SnapShot;
 
 #define SNAPCOUNT_DONE  255     /* Already compiled and linked a side trace. */
@@ -224,6 +224,14 @@ typedef struct GCtrace {
 
 LJ_STATIC_ASSERT(offsetof(GChead, gclist) == offsetof(GCtrace, gclist));
 
+static LJ_AINLINE MSize snap_nextofs(GCtrace *T, SnapShot *snap)
+{
+  if (snap+1 == &T->snap[T->nsnap])
+    return T->nsnapmap;
+  else
+    return (snap+1)->mapofs;
+}
+
 /* Round-robin penalty cache for bytecodes leading to aborted traces. */
 typedef struct HotPenalty {
   MRef pc;              /* Starting bytecode PC. */
diff --git a/src/lj_opt_loop.c b/src/lj_opt_loop.c
index d1512468..8d2935f6 100644
--- a/src/lj_opt_loop.c
+++ b/src/lj_opt_loop.c
@@ -182,7 +182,8 @@ static void loop_subst_snap(jit_State *J, SnapShot *osnap,
                             SnapEntry *loopmap, IRRef1 *subst)
 {
   SnapEntry *nmap, *omap = &J->cur.snapmap[osnap->mapofs];
-  MSize nmapofs, depth;
+  SnapEntry *nextmap = &J->cur.snapmap[snap_nextofs(&J->cur, osnap)];
+  MSize nmapofs;
   MSize on, ln, nn, onent = osnap->nent;
   BCReg nslots = osnap->nslots;
   SnapShot *snap = &J->cur.snap[J->cur.nsnap];
@@ -194,11 +195,9 @@ static void loop_subst_snap(jit_State *J, SnapShot *osnap,
     nmapofs = snap->mapofs;
   }
   J->guardemit.irt = 0;
-  depth = osnap->depth;
   /* Setup new snapshot. */
   snap->mapofs = (uint16_t)nmapofs;
   snap->ref = (IRRef1)J->cur.nins;
-  snap->depth = (uint8_t)depth;
   snap->nslots = nslots;
   snap->count = 0;
   nmap = &J->cur.snapmap[nmapofs];
@@ -220,11 +219,11 @@ static void loop_subst_snap(jit_State *J, SnapShot *osnap,
   while (snap_slot(loopmap[ln]) < nslots)  /* Copy remaining loop slots. */
     nmap[nn++] = loopmap[ln++];
   snap->nent = (uint8_t)nn;
-  J->cur.nsnapmap = (uint16_t)(nmapofs + nn + 1 + depth);
   omap += onent;
   nmap += nn;
-  for (nn = 0; nn <= depth; nn++)  /* Copy PC + frame links. */
-    nmap[nn] = omap[nn];
+  while (omap < nextmap)  /* Copy PC + frame links. */
+    *nmap++ = *omap++;
+  J->cur.nsnapmap = (uint16_t)(nmap - J->cur.snapmap);
 }
 
 /* Unroll loop. */
@@ -335,13 +334,13 @@ static void loop_unroll(jit_State *J)
 }
 
 /* Undo any partial changes made by the loop optimization. */
-static void loop_undo(jit_State *J, IRRef ins, SnapNo nsnap)
+static void loop_undo(jit_State *J, IRRef ins, SnapNo nsnap, MSize nsnapmap)
 {
   ptrdiff_t i;
   SnapShot *snap = &J->cur.snap[nsnap-1];
   SnapEntry *map = J->cur.snapmap;
   map[snap->mapofs + snap->nent] = map[J->cur.snap[0].nent];  /* Restore PC. */
-  J->cur.nsnapmap = (uint16_t)(snap->mapofs + snap->nent + 1 + snap->depth);
+  J->cur.nsnapmap = (uint16_t)nsnapmap;
   J->cur.nsnap = nsnap;
   J->guardemit.irt = 0;
   lj_ir_rollback(J, ins);
@@ -370,6 +369,7 @@ int lj_opt_loop(jit_State *J)
 {
   IRRef nins = J->cur.nins;
   SnapNo nsnap = J->cur.nsnap;
+  MSize nsnapmap = J->cur.nsnapmap;
   int errcode = lj_vm_cpcall(J->L, NULL, J, cploop_opt);
   if (LJ_UNLIKELY(errcode)) {
     lua_State *L = J->L;
@@ -382,7 +382,7 @@ int lj_opt_loop(jit_State *J)
         if (--J->instunroll < 0)  /* But do not unroll forever. */
           break;
         L->top--;  /* Remove error object. */
-        loop_undo(J, nins, nsnap);
+        loop_undo(J, nins, nsnap, nsnapmap);
         return 1;  /* Loop optimization failed, continue recording. */
       default:
         break;
diff --git a/src/lj_record.c b/src/lj_record.c
index 35ec1456..c8502d5d 100644
--- a/src/lj_record.c
+++ b/src/lj_record.c
@@ -2047,6 +2047,7 @@ static void rec_setup_side(jit_State *J, GCtrace *T)
   SnapEntry *map = &T->snapmap[snap->mapofs];
   MSize n, nent = snap->nent;
   BloomFilter seen = 0;
+  J->framedepth = 0;
   /* Emit IR for slots inherited from parent snapshot. */
   for (n = 0; n < nent; n++) {
     SnapEntry sn = map[n];
@@ -2087,12 +2088,12 @@ static void rec_setup_side(jit_State *J, GCtrace *T)
     }
   setslot:
     J->slot[s] = tr | (sn&(SNAP_CONT|SNAP_FRAME));  /* Same as TREF_* flags. */
+    J->framedepth += ((sn & (SNAP_CONT|SNAP_FRAME)) && s);
     if ((sn & SNAP_FRAME))
       J->baseslot = s+1;
   }
   J->base = J->slot + J->baseslot;
   J->maxslot = snap->nslots - J->baseslot;
-  J->framedepth = snap->depth;
   lj_snap_add(J);
 }
 
diff --git a/src/lj_snap.c b/src/lj_snap.c
index e29b5357..a2025d88 100644
--- a/src/lj_snap.c
+++ b/src/lj_snap.c
@@ -118,7 +118,6 @@ static void snapshot_stack(jit_State *J, SnapShot *snap, MSize nsnapmap)
   snap->mapofs = (uint16_t)nsnapmap;
   snap->ref = (IRRef1)J->cur.nins;
   snap->nent = (uint8_t)nent;
-  snap->depth = (uint8_t)J->framedepth;
   snap->nslots = (uint8_t)nslots;
   snap->count = 0;
   J->cur.nsnapmap = (uint16_t)(nsnapmap + nent + 1 + J->framedepth);
@@ -274,7 +273,7 @@ void lj_snap_shrink(jit_State *J)
       map[m++] = map[n];  /* Only copy used slots. */
   }
   snap->nent = (uint8_t)m;
-  nlim = nent + snap->depth;
+  nlim = J->cur.nsnapmap - snap->mapofs - 1;
   while (n <= nlim) map[m++] = map[n++];  /* Move PC + frame links down. */
   J->cur.nsnapmap = (uint16_t)(snap->mapofs + m);  /* Free up space in map. */
 }
@@ -337,7 +336,7 @@ const BCIns *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 + snap->depth;
+  SnapEntry *flinks = &T->snapmap[snap_nextofs(T, snap)-1];
   int32_t ftsz0;
   BCReg nslots = snap->nslots;
   TValue *frame;