feature: implemented the new table.nkeys Lua API.

This Lua builtin returns the number of elements in a Lua table, i.e.,
the number of non-nil array elements plus the number of non-nil key-value
pairs in a single Lua table.

It can be used like this:

  local nkeys = require "table.nkeys"

  print(nkeys({}))  -- 0
  print(nkeys({ "a", nil, "b" }))  -- 2
  print(nkeys({ dog = 3, cat = 4, bird = nil }))  -- 2
  print(nkeys({ "a", dog = 3, cat = 4 }))  -- 3

This Lua API can be JIT compiled.

Signed-off-by: Yichun Zhang (agentzh) <agentzh@gmail.com>

README

@@ -60,6 +60,23 @@ the upstream LuaJIT (https://github.com/LuaJIT/LuaJIT), which are
         print(isarr{"dog" = 3})  -- false
         print(isarr{})  -- true
 
+* feature: implemented the new table.nkeys Lua API.
+
+    This Lua builtin returns the number of elements in a Lua table, i.e.,
+    the number of non-nil array elements plus the number of non-nil key-value
+    pairs in a single Lua table.
+
+    It can be used like this:
+
+      local nkeys = require "table.nkeys"
+
+      print(nkeys({}))  -- 0
+      print(nkeys({ "a", nil, "b" }))  -- 2
+      print(nkeys({ dog = 3, cat = 4, bird = nil }))  -- 2
+      print(nkeys({ "a", dog = 3, cat = 4 }))  -- 3
+
+    This Lua API can be JIT compiled.
+
 * feature: implemented the new Lua and C API functions for thread exdata.
 
     The Lua API can be used like below:

src/lib_table.c

@@ -190,6 +190,16 @@ LJLIB_NOREG LJLIB_CF(table_isarray) LJLIB_REC(.)
   return 1;
 }
 
+LJLIB_NOREG LJLIB_CF(table_nkeys) LJLIB_REC(.)
+{
+  GCtab *src = lj_lib_checktab(L, 1);
+
+  setintV(L->base, lj_tab_nkeys(src));
+  L->top = L->base+1;
+
+  return 1;
+}
+
 LJLIB_NOREG LJLIB_CF(table_isempty) LJLIB_REC(.)
 {
   GCtab *src = lj_lib_checktab(L, 1);
@@ -340,6 +350,11 @@ static int luaopen_table_clone(lua_State *L)
   return lj_lib_postreg(L, lj_cf_table_clone, FF_table_clone, "clone");
 }
 
+static int luaopen_table_nkeys(lua_State *L)
+{
+  return lj_lib_postreg(L, lj_cf_table_nkeys, FF_table_nkeys, "nkeys");
+}
+
 static int luaopen_table_isarray(lua_State *L)
 {
   return lj_lib_postreg(L, lj_cf_table_isarray, FF_table_isarray, "isarray");
@@ -369,6 +384,7 @@ LUALIB_API int luaopen_table(lua_State *L)
   lj_lib_prereg(L, LUA_TABLIBNAME ".new", luaopen_table_new, tabV(L->top-1));
   lj_lib_prereg(L, LUA_TABLIBNAME ".clone", luaopen_table_clone, tabV(L->top-1));
   lj_lib_prereg(L, LUA_TABLIBNAME ".isarray", luaopen_table_isarray, tabV(L->top-1));
+  lj_lib_prereg(L, LUA_TABLIBNAME ".nkeys", luaopen_table_nkeys, tabV(L->top-1));
   lj_lib_prereg(L, LUA_TABLIBNAME ".isempty", luaopen_table_isempty, tabV(L->top-1));
   lj_lib_prereg(L, LUA_TABLIBNAME ".clear", luaopen_table_clear, tabV(L->top-1));
   return 1;

src/lj_ffrecord.c

@@ -1126,6 +1126,14 @@ static void LJ_FASTCALL recff_table_isarray(jit_State *J, RecordFFData *rd)
   }  /* else: Interpreter will throw. */
 }
 
+static void LJ_FASTCALL recff_table_nkeys(jit_State *J, RecordFFData *rd)
+{
+  TRef src = J->base[0];
+  if (LJ_LIKELY(tref_istab(src))) {
+    J->base[0] = lj_ir_call(J, IRCALL_lj_tab_nkeys, src);
+  }  /* else: Interpreter will throw. */
+}
+
 static void LJ_FASTCALL recff_table_isempty(jit_State *J, RecordFFData *rd)
 {
   TRef src = J->base[0];

src/lj_ircall.h

@@ -169,6 +169,7 @@ typedef struct CCallInfo {
   _(ANY,	lj_tab_len,		1,  FL, INT, 0) \
   _(ANY,	lj_tab_clone,		2,  FS, TAB, CCI_L) \
   _(ANY,	lj_tab_isarray,		1,  FL, INT, 0) \
+  _(ANY,	lj_tab_nkeys,		1,  FL, INT, 0) \
   _(ANY,	lj_tab_isempty,		1,  FL, INT, 0) \
   _(ANY,	lj_gc_step_jit,		2,  FS, NIL, CCI_L) \
   _(ANY,	lj_gc_barrieruv,	2,  FS, NIL, 0) \

src/lj_tab.c

@@ -719,6 +719,32 @@ int LJ_FASTCALL lj_tab_isarray(const GCtab *src)
   return 1;
 }
 
+MSize LJ_FASTCALL lj_tab_nkeys(const GCtab *t)
+{
+  MSize narr = (MSize)t->asize;
+  cTValue *e;
+  Node *node;
+  MSize i, cnt = 0;
+
+  e = tvref(t->array);
+  for (i = 0; i < narr; i++)
+    if (LJ_LIKELY(!tvisnil(&e[i])))
+      cnt++;
+
+  if (t->hmask <= 0)
+    return cnt;
+
+  node = noderef(t->node);
+  for (i = 0; i <= (MSize)t->hmask; i++) {
+    Node *n = &node[i];
+    if (LJ_LIKELY(!tvisnil(&n->val))) {
+      cnt++;
+    }
+  }
+
+  return cnt;
+}
+
 int LJ_FASTCALL lj_tab_isempty(const GCtab *t)
 {
   MSize narr = (MSize)t->asize;

src/lj_tab.h

@@ -72,6 +72,7 @@ LJ_FUNCA MSize LJ_FASTCALL lj_tab_len(GCtab *t);
 
 LJ_FUNCA GCtab * LJ_FASTCALL lj_tab_clone(lua_State *L, const GCtab *src);
 LJ_FUNCA int LJ_FASTCALL lj_tab_isarray(const GCtab *src);
+LJ_FUNCA MSize LJ_FASTCALL lj_tab_nkeys(const GCtab *src);
 LJ_FUNCA int LJ_FASTCALL lj_tab_isempty(const GCtab *t);
 
 #endif

t/nkeys.t

@@ -0,0 +1,175 @@
+# vim: set ss=4 ft= sw=4 et sts=4 ts=4:
+
+use lib '.';
+use t::TestLJ;
+
+plan tests => 3 * blocks();
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: hash table, interpreted
+--- lua
+jit.off()
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+print(nkeys(new_tab(0, 4)))
+print(nkeys({}))
+print(nkeys({ cats = 4 }))
+print(nkeys({ dogs = 3, cats = 4 }))
+print(nkeys({ dogs = nil, cats = 4 }))
+--- jv
+--- out
+0
+0
+1
+2
+1
+--- err
+
+
+
+=== TEST 2: hash table, JIT
+--- lua
+jit.on()
+jit.opt.start("minstitch=100000", "hotloop=2")
+
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+
+local list = {
+  new_tab(0, 4),
+  {},
+  { cats = 4 },
+  { dogs = 3, cats = 4 },
+  { dogs = nil, cats = 4 },
+}
+
+for i, t in ipairs(list) do
+  local total = 0
+  for i = 1, 10 do
+    total = total + nkeys(t)
+  end
+  print(total)
+end
+--- jv
+--- out
+0
+0
+10
+20
+10
+--- err
+[TRACE   1 test.lua:18 loop]
+[TRACE   2 test.lua:16 -> 1]
+
+
+
+=== TEST 3: pure arrays, interpreted
+--- lua
+jit.off()
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+print(nkeys(new_tab(5, 0)))
+print(nkeys({}))
+print(nkeys({ "cats" }))
+print(nkeys({ "dogs", 3, "cats", 4 }))
+print(nkeys({ "dogs", nil, "cats", 4 }))
+--- jv
+--- out
+0
+0
+1
+4
+3
+--- err
+
+
+
+=== TEST 4: pure array, JIT
+--- lua
+jit.on()
+jit.opt.start("minstitch=100000", "hotloop=2")
+
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+
+local list = {
+  new_tab(0, 4),
+  {},
+  { 3 },
+  { "cats", 4 },
+  { "dogs", 3, "cats", 4 },
+  { "dogs", nil, "cats", 4 },
+}
+
+for i, t in ipairs(list) do
+  local total = 0
+  for i = 1, 10 do
+    total = total + nkeys(t)
+  end
+  print(total)
+end
+--- jv
+--- out
+0
+0
+10
+20
+40
+30
+--- err
+[TRACE   1 test.lua:19 loop]
+[TRACE   2 test.lua:17 -> 1]
+
+
+
+=== TEST 5: mixing array and hash table, interpreted
+--- lua
+jit.off()
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+print(nkeys({ cats = 4, 5, 6 }))
+print(nkeys({ nil, "foo", dogs = 3, cats = 4 }))
+--- jv
+--- out
+3
+3
+--- err
+
+
+
+=== TEST 6: mixing array & hash, JIT
+--- lua
+jit.on()
+jit.opt.start("minstitch=100000", "hotloop=2")
+
+local new_tab = require "table.new"
+local assert = assert
+local nkeys = require "table.nkeys"
+
+local list = {
+  { cats = 4, 5, 6 },
+  { nil, "foo", dogs = 3, cats = 4 },
+}
+
+for i, t in ipairs(list) do
+  local total = 0
+  for i = 1, 10 do
+    total = total + nkeys(t)
+  end
+  print(total)
+end
+--- jv
+--- out
+30
+30
+--- err
+[TRACE   1 test.lua:15 loop]
+[TRACE   2 test.lua:13 -> 1]