Refactor table traversal.

Mike Pall · Mike Pall · commit c6f5ef649b64 · 2021-09-19T17:38:49.000+02:00
Sponsored by OpenResty Inc.
diff --git a/src/lib_base.c b/src/lib_base.c
@@ -79,6 +79,7 @@ LJ_STATIC_ASSERT((int)FF_next == FF_next_N);
 LJLIB_ASM(next)
 {
   lj_lib_checktab(L, 1);
+  lj_err_msg(L, LJ_ERR_NEXTIDX);
   return FFH_UNREACHABLE;
 }
 
diff --git a/src/lj_api.c b/src/lj_api.c
@@ -893,11 +893,13 @@ LUA_API int lua_next(lua_State *L, int idx)
   cTValue *t = index2adr(L, idx);
   int more;
   lj_checkapi(tvistab(t), "stack slot %d is not a table", idx);
-  more = lj_tab_next(L, tabV(t), L->top-1);
-  if (more) {
+  more = lj_tab_next(tabV(t), L->top-1, L->top-1);
+  if (more > 0) {
     incr_top(L);  /* Return new key and value slot. */
-  } else {  /* End of traversal. */
+  } else if (!more) {  /* End of traversal. */
     L->top--;  /* Remove key slot. */
+  } else {
+    lj_err_msg(L, LJ_ERR_NEXTIDX);
   }
   return more;
 }
diff --git a/src/lj_obj.h b/src/lj_obj.h
@@ -284,6 +284,9 @@ typedef const TValue cTValue;
 #define LJ_TISGCV		(LJ_TSTR+1)
 #define LJ_TISTABUD		LJ_TTAB
 
+/* Type marker for slot holding a traversal index. Must be lightuserdata. */
+#define LJ_KEYINDEX		0xfffe7fffu
+
 #if LJ_GC64
 #define LJ_GCVMASK		(((uint64_t)1 << 47) - 1)
 #endif
diff --git a/src/lj_tab.c b/src/lj_tab.c
@@ -568,56 +568,66 @@ TValue *lj_tab_set(lua_State *L, GCtab *t, cTValue *key)
 
 /* -- Table traversal ----------------------------------------------------- */
 
-/* Get the traversal index of a key. */
-static uint32_t keyindex(lua_State *L, GCtab *t, cTValue *key)
+/* Table traversal indexes:
+**
+** Array key index: [0 .. t->asize-1]
+** Hash key index:  [t->asize .. t->asize+t->hmask]
+** Invalid key:     ~0
+*/
+
+/* Get the successor traversal index of a key. */
+uint32_t LJ_FASTCALL lj_tab_keyindex(GCtab *t, cTValue *key)
 {
   TValue tmp;
   if (tvisint(key)) {
     int32_t k = intV(key);
     if ((uint32_t)k < t->asize)
-      return (uint32_t)k;  /* Array key indexes: [0..t->asize-1] */
+      return (uint32_t)k + 1;
     setnumV(&tmp, (lua_Number)k);
     key = &tmp;
   } else if (tvisnum(key)) {
     lua_Number nk = numV(key);
     int32_t k = lj_num2int(nk);
     if ((uint32_t)k < t->asize && nk == (lua_Number)k)
-      return (uint32_t)k;  /* Array key indexes: [0..t->asize-1] */
+      return (uint32_t)k + 1;
   }
   if (!tvisnil(key)) {
     Node *n = hashkey(t, key);
     do {
       if (lj_obj_equal(&n->key, key))
-	return t->asize + (uint32_t)(n - noderef(t->node));
-	/* Hash key indexes: [t->asize..t->asize+t->nmask] */
+	return t->asize + (uint32_t)((n+1) - noderef(t->node));
     } while ((n = nextnode(n)));
-    if (key->u32.hi == 0xfffe7fff)  /* ITERN was despecialized while running. */
-      return key->u32.lo - 1;
-    lj_err_msg(L, LJ_ERR_NEXTIDX);
-    return 0;  /* unreachable */
+    if (key->u32.hi == LJ_KEYINDEX)  /* Despecialized ITERN while running. */
+      return key->u32.lo;
+    return ~0u;  /* Invalid key to next. */
   }
-  return ~0u;  /* A nil key starts the traversal. */
+  return 0;  /* A nil key starts the traversal. */
 }
 
-/* Advance to the next step in a table traversal. */
-int lj_tab_next(lua_State *L, GCtab *t, TValue *key)
+/* Get the next key/value pair of a table traversal. */
+int lj_tab_next(GCtab *t, cTValue *key, TValue *o)
 {
-  uint32_t i = keyindex(L, t, key);  /* Find predecessor key index. */
-  for (i++; i < t->asize; i++)  /* First traverse the array keys. */
-    if (!tvisnil(arrayslot(t, i))) {
-      setintV(key, i);
-      copyTV(L, key+1, arrayslot(t, i));
+  uint32_t idx = lj_tab_keyindex(t, key);  /* Find successor index of key. */
+  /* First traverse the array part. */
+  for (; idx < t->asize; idx++) {
+    cTValue *a = arrayslot(t, idx);
+    if (LJ_LIKELY(!tvisnil(a))) {
+      setintV(o, idx);
+      o[1] = *a;
       return 1;
     }
-  for (i -= t->asize; i <= t->hmask; i++) {  /* Then traverse the hash keys. */
-    Node *n = &noderef(t->node)[i];
+  }
+  idx -= t->asize;
+  /* Then traverse the hash part. */
+  for (; idx <= t->hmask; idx++) {
+    Node *n = &noderef(t->node)[idx];
     if (!tvisnil(&n->val)) {
-      copyTV(L, key, &n->key);
-      copyTV(L, key+1, &n->val);
+      o[0] = n->key;
+      o[1] = n->val;
       return 1;
     }
   }
-  return 0;  /* End of traversal. */
+  return (int32_t)idx < 0 ? -1 : 0;  /* Invalid key or end of traversal. */
 }
 
 /* -- Table length calculation -------------------------------------------- */
diff --git a/src/lj_tab.h b/src/lj_tab.h
@@ -86,7 +86,8 @@ LJ_FUNC TValue *lj_tab_set(lua_State *L, GCtab *t, cTValue *key);
 #define lj_tab_setint(L, t, key) \
   (inarray((t), (key)) ? arrayslot((t), (key)) : lj_tab_setinth(L, (t), (key)))
 
-LJ_FUNCA int lj_tab_next(lua_State *L, GCtab *t, TValue *key);
+LJ_FUNC uint32_t LJ_FASTCALL lj_tab_keyindex(GCtab *t, cTValue *key);
+LJ_FUNCA int lj_tab_next(GCtab *t, cTValue *key, TValue *o);
 LJ_FUNCA MSize LJ_FASTCALL lj_tab_len(GCtab *t);
 #if LJ_HASJIT
 LJ_FUNC MSize LJ_FASTCALL lj_tab_len_hint(GCtab *t, size_t hint);
diff --git a/src/vm_arm.dasc b/src/vm_arm.dasc
@@ -1111,24 +1111,18 @@ static void build_subroutines(BuildCtx *ctx)
   |  checktab CARG2, ->fff_fallback
   |   strd CARG34, [BASE, NARGS8:RC]	// Set missing 2nd arg to nil.
   |   ldr PC, [BASE, FRAME_PC]
-  |  mov CARG2, CARG1
-  |    str BASE, L->base		// Add frame since C call can throw.
-  |  mov CARG1, L
-  |    str BASE, L->top			// Dummy frame length is ok.
-  |  add CARG3, BASE, #8
-  |   str PC, SAVE_PC
-  |  bl extern lj_tab_next	// (lua_State *L, GCtab *t, TValue *key)
-  |  // Returns 0 at end of traversal.
+  |  add CARG2, BASE, #8
+  |  sub CARG3, BASE, #8
+  |  bl extern lj_tab_next		// (GCtab *t, cTValue *key, TValue *o)
+  |  // Returns 1=found, 0=end, -1=error.
   |  .IOS ldr BASE, L->base
   |  cmp CRET1, #0
-  |  mvneq CRET2, #~LJ_TNIL
-  |  beq ->fff_restv			// End of traversal: return nil.
-  |  ldrd CARG12, [BASE, #8]		// Copy key and value to results.
-  |   ldrd CARG34, [BASE, #16]
-  |    mov RC, #(2+1)*8
-  |  strd CARG12, [BASE, #-8]
-  |   strd CARG34, [BASE]
-  |  b ->fff_res
+  |   mov RC, #(2+1)*8
+  |  bgt ->fff_res			// Found key/value.
+  |  bmi ->fff_fallback			// Invalid key.
+  |  // End of traversal: return nil.
+  |  mvn CRET2, #~LJ_TNIL
+  |  b ->fff_restv
   |
   |.ffunc_1 pairs
   |  checktab CARG2, ->fff_fallback
@@ -3989,7 +3983,7 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |   ins_next1
     |   ins_next2
     |  mov CARG1, #0
-    |  mvn CARG2, #0x00018000
+    |  mvn CARG2, #~LJ_KEYINDEX
     |  strd CARG1, [RA, #-8]		// Initialize control var.
     |1:
     |   ins_next3
diff --git a/src/vm_arm64.dasc b/src/vm_arm64.dasc
@@ -1086,21 +1086,19 @@ static void build_subroutines(BuildCtx *ctx)
   |//-- Base library: iterators -------------------------------------------
   |
   |.ffunc_1 next
-  |  checktp CARG2, CARG1, LJ_TTAB, ->fff_fallback
+  |  checktp CARG1, LJ_TTAB, ->fff_fallback
   |  str TISNIL, [BASE, NARGS8:RC]	// Set missing 2nd arg to nil.
   |  ldr PC, [BASE, FRAME_PC]
-  |   stp BASE, BASE, L->base		// Add frame since C call can throw.
-  |  mov CARG1, L
-  |  add CARG3, BASE, #8
-  |   str PC, SAVE_PC
-  |  bl extern lj_tab_next	// (lua_State *L, GCtab *t, TValue *key)
-  |  // Returns 0 at end of traversal.
+  |  add CARG2, BASE, #8
+  |  sub CARG3, BASE, #16
+  |  bl extern lj_tab_next		// (GCtab *t, cTValue *key, TValue *o)
+  |  // Returns 1=found, 0=end, -1=error.
+  |   mov RC, #(2+1)*8
+  |  tbnz CRET1w, #31, ->fff_fallback	// Invalid key.
+  |  cbnz CRET1, ->fff_res		// Found key/value.
+  |  // End of traversal: return nil.
   |  str TISNIL, [BASE, #-16]
-  |  cbz CRET1, ->fff_res1		// End of traversal: return nil.
-  |  ldp CARG1, CARG2, [BASE, #8]	// Copy key and value to results.
-  |    mov RC, #(2+1)*8
-  |  stp CARG1, CARG2, [BASE, #-16]
-  |  b ->fff_res
+  |  b ->fff_res1
   |
   |.ffunc_1 pairs
   |  checktp TMP1, CARG1, LJ_TTAB, ->fff_fallback
@@ -3384,7 +3382,7 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |   ccmp CARG4, TISNIL, #0, eq
     |  ccmp TMP1w, #FF_next_N, #0, eq
     |  bne >5
-    |  mov TMP0w, #0xfffe7fff
+    |  mov TMP0w, #0xfffe7fff		// LJ_KEYINDEX
     |  lsl TMP0, TMP0, #32
     |  str TMP0, [RA, #-8]		// Initialize control var.
     |1:
diff --git a/src/vm_mips.dasc b/src/vm_mips.dasc
@@ -1262,35 +1262,27 @@ static void build_subroutines(BuildCtx *ctx)
   |//-- Base library: iterators -------------------------------------------
   |
   |.ffunc next
-  |  lw CARG1, HI(BASE)
-  |   lw TAB:CARG2, LO(BASE)
+  |  lw CARG2, HI(BASE)
+  |   lw TAB:CARG1, LO(BASE)
   |  beqz NARGS8:RC, ->fff_fallback
   |.  addu TMP2, BASE, NARGS8:RC
   |  li AT, LJ_TTAB
   |   sw TISNIL, HI(TMP2)		// Set missing 2nd arg to nil.
-  |  bne CARG1, AT, ->fff_fallback
+  |  bne CARG2, AT, ->fff_fallback
   |.  lw PC, FRAME_PC(BASE)
   |  load_got lj_tab_next
-  |   sw BASE, L->base			// Add frame since C call can throw.
-  |   sw BASE, L->top			// Dummy frame length is ok.
-  |  addiu CARG3, BASE, 8
-  |   sw PC, SAVE_PC
-  |  call_intern lj_tab_next		// (lua_State *L, GCtab *t, TValue *key)
-  |.  move CARG1, L
-  |  // Returns 0 at end of traversal.
+  |  addiu CARG2, BASE, 8
+  |  call_intern lj_tab_next		// (GCtab *t, cTValue *key, TValue *o)
+  |.  addiu CARG3, BASE, -8
+  |  // Returns 1=found, 0=end, -1=error.
+  |   addiu RA, BASE, -8
+  |  bgtz CRET1, ->fff_res		// Found key/value.
+  |.  li RD, (2+1)*8
   |  beqz CRET1, ->fff_restv		// End of traversal: return nil.
   |.  li SFARG1HI, LJ_TNIL
-  |  lw TMP0, 8+HI(BASE)
-  |   lw TMP1, 8+LO(BASE)
-  |    addiu RA, BASE, -8
-  |  lw TMP2, 16+HI(BASE)
-  |   lw TMP3, 16+LO(BASE)
-  |  sw TMP0, HI(RA)
-  |   sw TMP1, LO(RA)
-  |  sw TMP2, 8+HI(RA)
-  |   sw TMP3, 8+LO(RA)
-  |  b ->fff_res
-  |.  li RD, (2+1)*8
+  |   lw CFUNC:RB, FRAME_FUNC(BASE)
+  |  b ->fff_fallback			// Invalid key.
+  |.  li RC, 2*8
   |
   |.ffunc_1 pairs
   |  li AT, LJ_TTAB
@@ -4611,9 +4603,9 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |  addiu CARG2, CARG2, -FF_next_N
     |  or CARG2, CARG2, CARG3
     |  bnez CARG2, >5
-    |.  lui TMP1, 0xfffe
+    |.  lui TMP1, (LJ_KEYINDEX >> 16)
     |  addu PC, TMP0, TMP2
-    |  ori TMP1, TMP1, 0x7fff
+    |  ori TMP1, TMP1, (LJ_KEYINDEX & 0xffff)
     |  sw r0, -8+LO(RA)			// Initialize control var.
     |  sw TMP1, -8+HI(RA)
     |1:
diff --git a/src/vm_mips64.dasc b/src/vm_mips64.dasc
@@ -1322,27 +1322,24 @@ static void build_subroutines(BuildCtx *ctx)
   |//-- Base library: iterators -------------------------------------------
   |
   |.ffunc_1 next
-  |  checktp CARG2, CARG1, -LJ_TTAB, ->fff_fallback
+  |  checktp CARG1, -LJ_TTAB, ->fff_fallback
   |  daddu TMP2, BASE, NARGS8:RC
   |  sd TISNIL, 0(TMP2)			// Set missing 2nd arg to nil.
-  |  ld PC, FRAME_PC(BASE)
   |  load_got lj_tab_next
-  |   sd BASE, L->base			// Add frame since C call can throw.
-  |   sd BASE, L->top			// Dummy frame length is ok.
-  |  daddiu CARG3, BASE, 8
-  |   sd PC, SAVE_PC
-  |  call_intern lj_tab_next		// (lua_State *L, GCtab *t, TValue *key)
-  |.  move CARG1, L
-  |  // Returns 0 at end of traversal.
+  |  ld PC, FRAME_PC(BASE)
+  |  daddiu CARG2, BASE, 8
+  |  call_intern lj_tab_next		// (GCtab *t, cTValue *key, TValue *o)
+  |.  daddiu CARG3, BASE, -16
+  |  // Returns 1=found, 0=end, -1=error.
+  |   daddiu RA, BASE, -16
+  |  bgtz CRET1, ->fff_res		// Found key/value.
+  |.  li RD, (2+1)*8
   |  beqz CRET1, ->fff_restv		// End of traversal: return nil.
   |.  move CARG1, TISNIL
-  |  ld TMP0, 8(BASE)
-  |    daddiu RA, BASE, -16
-  |  ld TMP2, 16(BASE)
-  |  sd TMP0, 0(RA)
-  |  sd TMP2, 8(RA)
-  |  b ->fff_res
-  |.  li RD, (2+1)*8
+  |   ld CFUNC:RB, FRAME_FUNC(BASE)
+  |   cleartp CFUNC:RB
+  |  b ->fff_fallback			// Invalid key.
+  |.  li RC, 2*8
   |
   |.ffunc_1 pairs
   |  checktp TAB:TMP1, CARG1, -LJ_TTAB, ->fff_fallback
@@ -4727,11 +4724,10 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |.  addiu RC, RC, 1
     |   sd TMP2, 0(RA)
     |  sd CARG1, 8(RA)
-    |   or TMP0, RC, CARG3
     |     lui TMP3, (-(BCBIAS_J*4 >> 16) & 65535)
     |     decode_RD4b RD
     |     daddu RD, RD, TMP3
-    |   sw TMP0, -8+LO(RA)		// Update control var.
+    |   sw RC, -8+LO(RA)		// Update control var.
     |     daddu PC, PC, RD
     |3:
     |  ins_next
@@ -4781,9 +4777,9 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |  daddiu TMP1, TMP1, -FF_next_N
     |  or AT, AT, TMP1
     |  bnez AT, >5
-    |.  lui TMP1, 0xfffe
+    |.  lui TMP1, (LJ_KEYINDEX >> 16)
     |  daddu PC, TMP0, TMP2
-    |  ori TMP1, TMP1, 0x7fff
+    |  ori TMP1, TMP1, (LJ_KEYINDEX & 0xffff)
     |  dsll TMP1, TMP1, 32
     |  sd TMP1, -8(RA)
     |1:
diff --git a/src/vm_ppc.dasc b/src/vm_ppc.dasc
diff --git a/src/vm_x64.dasc b/src/vm_x64.dasc
diff --git a/src/vm_x86.dasc b/src/vm_x86.dasc

Original file line number	Diff line number	Diff line change
`@@ -79,6 +79,7 @@ LJ_STATIC_ASSERT((int)FF_next == FF_next_N);`
`79`	`79`	`LJLIB_ASM(next)`
`80`	`80`	`{`
`81`	`81`	`lj_lib_checktab(L, 1);`
	`82`	`+ lj_err_msg(L, LJ_ERR_NEXTIDX);`
`82`	`83`	`return FFH_UNREACHABLE;`
`83`	`84`	`}`
`84`	`85`