6 files changed, 39 insertions, 12 deletions

diff --git a/src/lj_asm_arm.h b/src/lj_asm_arm.h
index ded63913..326330f4 100644
--- a/src/lj_asm_arm.h
+++ b/src/lj_asm_arm.h
@@ -1250,7 +1250,12 @@ dotypecheck:
       }
     }
     asm_guardcc(as, t == IRT_NUM ? CC_HS : CC_NE);
-    emit_n(as, ARMI_CMN|ARMI_K12|-irt_toitype_(t), type);
+    if ((ir->op2 & IRSLOAD_KEYINDEX)) {
+      emit_n(as, ARMI_CMN|ARMI_K12|1, type);
+      emit_dn(as, ARMI_EOR^emit_isk12(ARMI_EOR, ~LJ_KEYINDEX), type, type);
+    } else {
+      emit_n(as, ARMI_CMN|ARMI_K12|-irt_toitype_(t), type);
+    }
   }
   if (ra_hasreg(dest)) {
 #if !LJ_SOFTFP
diff --git a/src/lj_asm_arm64.h b/src/lj_asm_arm64.h
index eb31b006..1f44d023 100644
--- a/src/lj_asm_arm64.h
+++ b/src/lj_asm_arm64.h
@@ -1209,7 +1209,7 @@ dotypecheck:
       lj_assertA(irt_isinteger(t) || irt_isnum(t),
                  "bad SLOAD type %d", irt_type(t));
       emit_nm(as, A64I_CMPx | A64F_SH(A64SH_LSR, 32),
-              ra_allock(as, LJ_TISNUM << 15, allow), tmp);
+              ra_allock(as, (ir->op2 & IRSLOAD_KEYINDEX) ? LJ_KEYINDEX : (LJ_TISNUM << 15), allow), tmp);
     } else if (irt_isnil(t)) {
       emit_n(as, (A64I_CMNx^A64I_K12) | A64F_U12(1), tmp);
     } else if (irt_ispri(t)) {
diff --git a/src/lj_asm_mips.h b/src/lj_asm_mips.h
index db42b8f3..1686b40f 100644
--- a/src/lj_asm_mips.h
+++ b/src/lj_asm_mips.h
@@ -1577,7 +1577,7 @@ dotypecheck:
       asm_guard(as, MIPSI_BEQ, RID_TMP, RID_ZERO);
       emit_tsi(as, MIPSI_SLTIU, RID_TMP, type, (int32_t)LJ_TISNUM);
     } else {
-      Reg ktype = ra_allock(as, irt_toitype(t), allow);
+      Reg ktype = ra_allock(as, (ir->op2 & IRSLOAD_KEYINDEX) ? LJ_KEYINDEX : irt_toitype(t), allow);
       asm_guard(as, MIPSI_BNE, type, ktype);
     }
   }
@@ -1595,6 +1595,10 @@ dotypecheck:
     if (irt_ispri(t)) {
       asm_guard(as, MIPSI_BNE, type,
                 ra_allock(as, ~((int64_t)~irt_toitype(t) << 47) , allow));
+    } else if ((ir->op2 & IRSLOAD_KEYINDEX)) {
+      asm_guard(as, MIPSI_BNE, RID_TMP,
+                ra_allock(as, (int32_t)LJ_KEYINDEX, allow));
+      emit_dta(as, MIPSI_DSRA32, RID_TMP, type, 0);
     } else {
       if (irt_isnum(t)) {
         asm_guard(as, MIPSI_BEQ, RID_TMP, RID_ZERO);
@@ -2568,7 +2572,22 @@ static void asm_stack_restore(ASMState *as, SnapShot *snap)
       }
       emit_tsi(as, MIPSI_SW, type, RID_BASE, ofs+(LJ_BE?0:4));
 #else
-      asm_tvstore64(as, RID_BASE, ofs, ref);
+      if ((sn & SNAP_KEYINDEX)) {
+        RegSet allow = rset_exclude(RSET_GPR, RID_BASE);
+        int64_t kki = (int64_t)LJ_KEYINDEX << 32;
+        if (irref_isk(ref)) {
+          emit_tsi(as, MIPSI_SD,
+                   ra_allock(as, kki | (int64_t)(uint32_t)ir->i, allow),
+                   RID_BASE, ofs);
+        } else {
+          Reg src = ra_alloc1(as, ref, allow);
+          Reg rki = ra_allock(as, kki, rset_exclude(allow, src));
+          emit_tsi(as, MIPSI_SD, RID_TMP, RID_BASE, ofs);
+          emit_dst(as, MIPSI_DADDU, RID_TMP, src, rki);
+        }
+      } else {
+        asm_tvstore64(as, RID_BASE, ofs, ref);
+      }
 #endif
     }
     checkmclim(as);
diff --git a/src/lj_asm_ppc.h b/src/lj_asm_ppc.h
index 5ea4d47d..546b8e5d 100644
--- a/src/lj_asm_ppc.h
+++ b/src/lj_asm_ppc.h
@@ -1169,7 +1169,12 @@ dotypecheck:
   } else {
     if ((ir->op2 & IRSLOAD_TYPECHECK)) {
       asm_guardcc(as, CC_NE);
-      emit_ai(as, PPCI_CMPWI, RID_TMP, irt_toitype(t));
+      if ((ir->op2 & IRSLOAD_KEYINDEX)) {
+        emit_ai(as, PPCI_CMPWI, RID_TMP, (LJ_KEYINDEX & 0xffff));
+        emit_asi(as, PPCI_XORIS, RID_TMP, RID_TMP, (LJ_KEYINDEX >> 16));
+      } else {
+        emit_ai(as, PPCI_CMPWI, RID_TMP, irt_toitype(t));
+      }
       type = RID_TMP;
     }
     if (ra_hasreg(dest)) emit_tai(as, PPCI_LWZ, dest, base, ofs);
diff --git a/src/lj_asm_x86.h b/src/lj_asm_x86.h
index 38069e1d..4465efa2 100644
--- a/src/lj_asm_x86.h
+++ b/src/lj_asm_x86.h
@@ -1768,14 +1768,11 @@ static void asm_sload(ASMState *as, IRIns *ir)
   if ((ir->op2 & IRSLOAD_TYPECHECK)) {
     /* Need type check, even if the load result is unused. */
     asm_guardcc(as, irt_isnum(t) ? CC_AE : CC_NE);
-    if (LJ_64 && irt_type(t) >= IRT_NUM) {
+    if ((LJ_64 && irt_type(t) >= IRT_NUM) || (ir->op2 & IRSLOAD_KEYINDEX)) {
       lj_assertA(irt_isinteger(t) || irt_isnum(t),
                  "bad SLOAD type %d", irt_type(t));
-#if LJ_GC64
-      emit_u32(as, LJ_TISNUM << 15);
-#else
-      emit_u32(as, LJ_TISNUM);
-#endif
+      emit_u32(as, (ir->op2 & IRSLOAD_KEYINDEX) ? LJ_KEYINDEX :
+                   LJ_GC64 ? (LJ_TISNUM << 15) : LJ_TISNUM);
       emit_rmro(as, XO_ARITHi, XOg_CMP, base, ofs+4);
 #if LJ_GC64
     } else if (irt_isnil(t)) {
diff --git a/src/lj_record.c b/src/lj_record.c
index 59798844..d48908e2 100644
--- a/src/lj_record.c
+++ b/src/lj_record.c
@@ -675,7 +675,8 @@ static LoopEvent rec_itern(jit_State *J, BCReg ra, BCReg rb)
   lj_snap_add(J);  /* Required to make JLOOP the first ins in a side-trace. */
   ix.tab = getslot(J, ra-2);
   ix.key = J->base[ra-1] ? J->base[ra-1] :
-           sloadt(J, (int32_t)(ra-1), IRT_INT, IRSLOAD_KEYINDEX);
+           sloadt(J, (int32_t)(ra-1), IRT_GUARD|IRT_INT,
+                  IRSLOAD_TYPECHECK|IRSLOAD_KEYINDEX);
   copyTV(J->L, &ix.tabv, &J->L->base[ra-2]);
   copyTV(J->L, &ix.keyv, &J->L->base[ra-1]);
   ix.idxchain = (rb < 3);  /* Omit value type check, if unused. */