Restore MULTRES for snapshots pointing to CALLM etc. bytecodes.

3 files changed, 49 insertions, 3 deletions

diff --git a/src/lj_asm.c b/src/lj_asm.c
index cd1c0dbd..c749ada0 100644
--- a/src/lj_asm.c
+++ b/src/lj_asm.c
@@ -3253,6 +3253,16 @@ static void asm_head_side(ASMState *as)
 
 /* -- Tail of trace ------------------------------------------------------- */
 
+/* Set MULTRES in C frame. */
+static void asm_tail_multres(ASMState *as, BCReg mres)
+{
+  /* We don't know spadj yet, so get the C frame from L->cframe. */
+  emit_movmroi(as, RID_RET, CFRAME_OFS_MULTRES, mres);
+  emit_gri(as, XG_ARITHi(XOg_AND), RID_RET|REX_64, CFRAME_RAWMASK);
+  emit_rmro(as, XO_MOV, RID_RET|REX_64, RID_RET, offsetof(lua_State, cframe));
+  emit_getgl(as, RID_RET, jit_L);
+}
+
 /* Link to another trace. */
 static void asm_tail_link(ASMState *as)
 {
@@ -3273,6 +3283,19 @@ static void asm_tail_link(ASMState *as)
     }
     emit_loada(as, RID_DISPATCH, J2GG(as->J)->dispatch);
     emit_loada(as, RID_PC, pc);
+    switch (bc_op(*pc)) {
+    case BC_CALLM: case BC_CALLMT:
+      asm_tail_multres(as, snap->nslots - baseslot - 1 - bc_a(*pc) - bc_c(*pc));
+      break;
+    case BC_RETM:
+      asm_tail_multres(as, snap->nslots - baseslot - bc_a(*pc) - bc_d(*pc));
+      break;
+    case BC_TSETM:
+      asm_tail_multres(as, snap->nslots - baseslot - bc_a(*pc));
+      break;
+    default:
+      break;
+    }
   } else if (baseslot) {
     /* Save modified BASE for linking to trace with higher start frame. */
     emit_setgl(as, RID_BASE, jit_base);
diff --git a/src/lj_snap.c b/src/lj_snap.c
index 86890a26..fe0389fe 100644
--- a/src/lj_snap.c
+++ b/src/lj_snap.c
@@ -309,7 +309,14 @@ const BCIns *lj_snap_restore(jit_State *J, void *exptr)
       }
     }
   }
-  L->top = curr_topL(L);
+  switch (bc_op(*pc)) {
+  case BC_CALLM: case BC_CALLMT: case BC_RETM: case BC_TSETM:
+    L->top = frame + nslots;
+    break;
+  default:
+    L->top = curr_topL(L);
+    break;
+  }
   lua_assert(map + nent == flinks);
   return pc;
 }
diff --git a/src/lj_trace.c b/src/lj_trace.c
index 246dc03c..f7b12068 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -666,6 +666,7 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
   ExitDataCP exd;
   int errcode;
   const BCIns *pc;
+  void *cf;
   exd.J = J;
   exd.exptr = exptr;
   errcode = lj_vm_cpcall(L, NULL, &exd, trace_exit_cp);
@@ -692,7 +693,9 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
 
   pc = exd.pc;
   trace_hotside(J, pc);
-  if (bc_op(*pc) == BC_JLOOP) {
+  cf = cframe_raw(L->cframe);
+  switch (bc_op(*pc)) {
+  case BC_JLOOP: {
     BCIns *retpc = &J->trace[bc_d(*pc)]->startins;
     if (bc_isret(bc_op(*retpc))) {
       if (J->state == LJ_TRACE_RECORD) {
@@ -703,8 +706,21 @@ int LJ_FASTCALL lj_trace_exit(jit_State *J, void *exptr)
         pc = retpc;
       }
     }
+    break;
+    }
+  case BC_CALLM: case BC_CALLMT:
+    cframe_multres(cf) = (BCReg)(L->top - L->base) - bc_a(*pc) - bc_c(*pc);
+    break;
+  case BC_RETM:
+    cframe_multres(cf) = (BCReg)(L->top - L->base) + 1 - bc_a(*pc) - bc_d(*pc);
+    break;
+  case BC_TSETM:
+    cframe_multres(cf) = (BCReg)(L->top - L->base) + 1 - bc_a(*pc);
+    break;
+  default:
+    break;
   }
-  setcframe_pc(cframe_raw(L->cframe), pc);
+  setcframe_pc(cf, pc);
   return 0;
 }