diff --git a/src/lib_base.c b/src/lib_base.c index e5f71d9c..0717ea86 100644 --- a/src/lib_base.c +++ b/src/lib_base.c @@ -163,7 +163,7 @@ LJLIB_CF(unpack) return n; } -LJLIB_CF(select) +LJLIB_CF(select) LJLIB_REC(.) { int32_t n = (int32_t)(L->top - L->base); if (n >= 1 && tvisstr(L->base) && *strVdata(L->base) == '#') { diff --git a/src/lj_record.c b/src/lj_record.c index 739279ad..5c998adf 100644 --- a/src/lj_record.c +++ b/src/lj_record.c @@ -1293,6 +1293,48 @@ static void LJ_FASTCALL recff_rawequal(jit_State *J, RecordFFData *rd) } /* else: Interpreter will throw. */ } +/* Determine mode of select() call. */ +static int32_t select_mode(jit_State *J, TRef tr, TValue *tv) +{ + if (tref_isstr(tr) && *strVdata(tv) == '#') { /* select('#', ...) */ + if (strV(tv)->len == 1) { + emitir(IRT(IR_EQ, IRT_STR), tr, lj_ir_kstr(J, strV(tv))); + } else { + TRef trptr = emitir(IRT(IR_STRREF, IRT_PTR), tr, 0); + TRef trchar = emitir(IRT(IR_XLOAD, IRT_U8), trptr, IRXLOAD_READONLY); + emitir(IRT(IR_EQ, IRT_INT), trchar, lj_ir_kint(J, '#')); + } + return 0; + } else { /* select(n, ...) */ + int32_t start = argv2int(J, tv); + if (start == 0) lj_trace_err(J, LJ_TRERR_BADTYPE); /* A bit misleading. */ + return start; + } +} + +static void LJ_FASTCALL recff_select(jit_State *J, RecordFFData *rd) +{ + TRef tr = J->base[0]; + if (tr) { + ptrdiff_t start = select_mode(J, tr, &rd->argv[0]); + if (start == 0) { /* select('#', ...) */ + J->base[0] = lj_ir_kint(J, J->maxslot - 1); + } else if (tref_isk(tr)) { /* select(k, ...) */ + ptrdiff_t n = (ptrdiff_t)J->maxslot; + if (start < 0) start += n; + else if (start > n) start = n; + rd->nres = n - start; + if (start >= 1) { + ptrdiff_t i; + for (i = 0; i < n - start; i++) + J->base[i] = J->base[start+i]; + } /* else: Interpreter will throw. */ + } else { + recff_nyiu(J); + } + } /* else: Interpreter will throw. */ +} + static void LJ_FASTCALL recff_tonumber(jit_State *J, RecordFFData *rd) { TRef tr = J->base[0];