Fix overflow of snapshot map offset.

Thanks to Yichun Zhang.

src/lj_jit.h

@@ -163,7 +163,7 @@ typedef struct MCLink {
 
 /* Stack snapshot header. */
 typedef struct SnapShot {
-  uint16_t mapofs;	/* Offset into snapshot map. */
+  uint32_t mapofs;	/* Offset into snapshot map. */
   IRRef1 ref;		/* First IR ref for this snapshot. */
   uint8_t nslots;	/* Number of valid slots. */
   uint8_t topslot;	/* Maximum frame extent. */
@@ -217,14 +217,12 @@ typedef enum {
 /* Trace object. */
 typedef struct GCtrace {
   GCHeader;
-  uint8_t topslot;	/* Top stack slot already checked to be allocated. */
-  uint8_t linktype;	/* Type of link. */
+  uint16_t nsnap;	/* Number of snapshots. */
   IRRef nins;		/* Next IR instruction. Biased with REF_BIAS. */
   GCRef gclist;
   IRIns *ir;		/* IR instructions/constants. Biased with REF_BIAS. */
   IRRef nk;		/* Lowest IR constant. Biased with REF_BIAS. */
-  uint16_t nsnap;	/* Number of snapshots. */
-  uint16_t nsnapmap;	/* Number of snapshot map elements. */
+  uint32_t nsnapmap;	/* Number of snapshot map elements. */
   SnapShot *snap;	/* Snapshot array. */
   SnapEntry *snapmap;	/* Snapshot map. */
   GCRef startpt;	/* Starting prototype. */
@@ -241,6 +239,8 @@ typedef struct GCtrace {
   TraceNo1 nextroot;	/* Next root trace for same prototype. */
   TraceNo1 nextside;	/* Next side trace of same root trace. */
   uint8_t sinktags;	/* Trace has SINK tags. */
+  uint8_t topslot;	/* Top stack slot already checked to be allocated. */
+  uint8_t linktype;	/* Type of link. */
   uint8_t unused1;
 #ifdef LUAJIT_USE_GDBJIT
   void *gdbjit_entry;	/* GDB JIT entry. */

src/lj_opt_loop.c

@@ -223,7 +223,7 @@ static void loop_subst_snap(jit_State *J, SnapShot *osnap,
   }
   J->guardemit.irt = 0;
   /* Setup new snapshot. */
-  snap->mapofs = (uint16_t)nmapofs;
+  snap->mapofs = (uint32_t)nmapofs;
   snap->ref = (IRRef1)J->cur.nins;
   snap->nslots = nslots;
   snap->topslot = osnap->topslot;
@@ -251,7 +251,7 @@ static void loop_subst_snap(jit_State *J, SnapShot *osnap,
   nmap += nn;
   while (omap < nextmap)  /* Copy PC + frame links. */
     *nmap++ = *omap++;
-  J->cur.nsnapmap = (uint16_t)(nmap - J->cur.snapmap);
+  J->cur.nsnapmap = (uint32_t)(nmap - J->cur.snapmap);
 }
 
 /* Unroll loop. */
@@ -362,7 +362,7 @@ static void loop_unroll(jit_State *J)
     }
   }
   if (!irt_isguard(J->guardemit))  /* Drop redundant snapshot. */
-    J->cur.nsnapmap = (uint16_t)J->cur.snap[--J->cur.nsnap].mapofs;
+    J->cur.nsnapmap = (uint32_t)J->cur.snap[--J->cur.nsnap].mapofs;
   lua_assert(J->cur.nsnapmap <= J->sizesnapmap);
   *psentinel = J->cur.snapmap[J->cur.snap[0].nent];  /* Restore PC. */
 
@@ -376,7 +376,7 @@ static void loop_undo(jit_State *J, IRRef ins, SnapNo nsnap, MSize nsnapmap)
   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)nsnapmap;
+  J->cur.nsnapmap = (uint32_t)nsnapmap;
   J->cur.nsnap = nsnap;
   J->guardemit.irt = 0;
   lj_ir_rollback(J, ins);

src/lj_snap.c

@@ -129,12 +129,12 @@ static void snapshot_stack(jit_State *J, SnapShot *snap, MSize nsnapmap)
   p = &J->cur.snapmap[nsnapmap];
   nent = snapshot_slots(J, p, nslots);
   snap->topslot = (uint8_t)snapshot_framelinks(J, p + nent);
-  snap->mapofs = (uint16_t)nsnapmap;
+  snap->mapofs = (uint32_t)nsnapmap;
   snap->ref = (IRRef1)J->cur.nins;
   snap->nent = (uint8_t)nent;
   snap->nslots = (uint8_t)nslots;
   snap->count = 0;
-  J->cur.nsnapmap = (uint16_t)(nsnapmap + nent + 1 + J->framedepth);
+  J->cur.nsnapmap = (uint32_t)(nsnapmap + nent + 1 + J->framedepth);
 }
 
 /* Add or merge a snapshot. */
@@ -294,7 +294,7 @@ void lj_snap_shrink(jit_State *J)
   snap->nent = (uint8_t)m;
   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. */
+  J->cur.nsnapmap = (uint32_t)(snap->mapofs + m);  /* Free up space in map. */
 }
 
 /* -- Snapshot access ----------------------------------------------------- */