New table API for 'set' functions

5 files changed, 180 insertions, 71 deletions

diff --git a/lapi.c b/lapi.c
index 2e27bc92..97a9f272 100644
--- a/lapi.c
+++ b/lapi.c
@@ -823,17 +823,18 @@ LUA_API int lua_getiuservalue (lua_State *L, int idx, int n) {
 ** t[k] = value at the top of the stack (where 'k' is a string)
 */
 static void auxsetstr (lua_State *L, const TValue *t, const char *k) {
-  const TValue *slot;
+  int aux;
   TString *str = luaS_new(L, k);
   api_checknelems(L, 1);
-  if (luaV_fastget(L, t, str, slot, luaH_getstr)) {
-    luaV_finishfastset(L, t, slot, s2v(L->top.p - 1));
+  luaV_fastset1(t, str, s2v(L->top.p - 1), aux, luaH_setstr1);
+  if (aux == HOK) {
+    luaV_finishfastset1(L, t, s2v(L->top.p - 1));
     L->top.p--;  /* pop value */
   }
   else {
     setsvalue2s(L, L->top.p, str);  /* push 'str' (to make it a TValue) */
     api_incr_top(L);
-    luaV_finishset(L, t, s2v(L->top.p - 1), s2v(L->top.p - 2), slot);
+    luaV_finishset1(L, t, s2v(L->top.p - 1), s2v(L->top.p - 2), aux);
     L->top.p -= 2;  /* pop value and key */
   }
   lua_unlock(L);  /* lock done by caller */
@@ -850,15 +851,16 @@ LUA_API void lua_setglobal (lua_State *L, const char *name) {
 
 LUA_API void lua_settable (lua_State *L, int idx) {
   TValue *t;
-  const TValue *slot;
+  int aux;
   lua_lock(L);
   api_checknelems(L, 2);
   t = index2value(L, idx);
-  if (luaV_fastget(L, t, s2v(L->top.p - 2), slot, luaH_get)) {
-    luaV_finishfastset(L, t, slot, s2v(L->top.p - 1));
+  luaV_fastset1(t, s2v(L->top.p - 2), s2v(L->top.p - 1), aux, luaH_set1);
+  if (aux == HOK) {
+    luaV_finishfastset1(L, t, s2v(L->top.p - 1));
   }
   else
-    luaV_finishset(L, t, s2v(L->top.p - 2), s2v(L->top.p - 1), slot);
+    luaV_finishset1(L, t, s2v(L->top.p - 2), s2v(L->top.p - 1), aux);
   L->top.p -= 2;  /* pop index and value */
   lua_unlock(L);
 }
@@ -872,17 +874,18 @@ LUA_API void lua_setfield (lua_State *L, int idx, const char *k) {
 
 LUA_API void lua_seti (lua_State *L, int idx, lua_Integer n) {
   TValue *t;
-  const TValue *slot;
+  int aux;
   lua_lock(L);
   api_checknelems(L, 1);
   t = index2value(L, idx);
-  if (luaV_fastgeti(L, t, n, slot)) {
-    luaV_finishfastset(L, t, slot, s2v(L->top.p - 1));
+  luaV_fastseti1(t, n, s2v(L->top.p - 1), aux);
+  if (aux == HOK) {
+    luaV_finishfastset1(L, t, s2v(L->top.p - 1));
   }
   else {
-    TValue aux;
-    setivalue(&aux, n);
-    luaV_finishset(L, t, &aux, s2v(L->top.p - 1), slot);
+    TValue temp;
+    setivalue(&temp, n);
+    luaV_finishset1(L, t, &temp, s2v(L->top.p - 1), aux);
   }
   L->top.p--;  /* pop value */
   lua_unlock(L);
diff --git a/ltable.c b/ltable.c
index 8fd83fda..902f05a7 100644
--- a/ltable.c
+++ b/ltable.c
@@ -719,15 +719,7 @@ void luaH_newkey (lua_State *L, Table *t, const TValue *key, TValue *value) {
 }
 
 
-/*
-** Search function for integers. If integer is inside 'alimit', get it
-** directly from the array part. Otherwise, if 'alimit' is not equal to
-** the real size of the array, key still can be in the array part. In
-** this case, try to avoid a call to 'luaH_realasize' when key is just
-** one more than the limit (so that it can be incremented without
-** changing the real size of the array).
-*/
-const TValue *luaH_getint (Table *t, lua_Integer key) {
+static const TValue *getintfromarray (Table *t, lua_Integer key) {
   if (l_castS2U(key) - 1u < t->alimit)  /* 'key' in [1, t->alimit]? */
     return &t->array[key - 1];
   else if (!limitequalsasize(t) &&  /* key still may be in the array part? */
@@ -736,19 +728,40 @@ const TValue *luaH_getint (Table *t, lua_Integer key) {
     t->alimit = cast_uint(key);  /* probably '#t' is here now */
     return &t->array[key - 1];
   }
-  else {
-    Node *n = hashint(t, key);
-    for (;;) {  /* check whether 'key' is somewhere in the chain */
-      if (keyisinteger(n) && keyival(n) == key)
-        return gval(n);  /* that's it */
-      else {
-        int nx = gnext(n);
-        if (nx == 0) break;
-        n += nx;
-      }
+  else return NULL;  /* key is not in the array part */
+}
+
+
+static const TValue *getintfromhash (Table *t, lua_Integer key) {
+  Node *n = hashint(t, key);
+  lua_assert(l_castS2U(key) - 1u >= luaH_realasize(t));
+  for (;;) {  /* check whether 'key' is somewhere in the chain */
+    if (keyisinteger(n) && keyival(n) == key)
+      return gval(n);  /* that's it */
+    else {
+      int nx = gnext(n);
+      if (nx == 0) break;
+      n += nx;
     }
-    return &absentkey;
   }
+  return &absentkey;
+}
+
+
+/*
+** Search function for integers. If integer is inside 'alimit', get it
+** directly from the array part. Otherwise, if 'alimit' is not equal to
+** the real size of the array, key still can be in the array part. In
+** this case, try to avoid a call to 'luaH_realasize' when key is just
+** one more than the limit (so that it can be incremented without
+** changing the real size of the array).
+*/
+const TValue *luaH_getint (Table *t, lua_Integer key) {
+  const TValue *slot = getintfromarray(t, key);
+  if (slot != NULL)
+    return slot;
+  else
+    return getintfromhash(t, key);
 }
 
 
@@ -832,6 +845,58 @@ int luaH_get1 (Table *t, const TValue *key, TValue *res) {
 }
 
 
+static int finishnodeset (Table *t, const TValue *slot, TValue *val) {
+  if (!ttisnil(slot)) {
+    setobj(((lua_State*)NULL), cast(TValue*, slot), val);
+    return HOK;  /* success */
+  }
+  else if (isabstkey(slot))
+    return HNOTFOUND;  /* no slot with that key */
+  else return (cast(Node*, slot) - t->node) + HFIRSTNODE;  /* node encoded */
+}
+
+
+int luaH_setint1 (Table *t, lua_Integer key, TValue *val) {
+  const TValue *slot = getintfromarray(t, key);
+  if (slot != NULL) {
+    if (!ttisnil(slot)) {
+      setobj(((lua_State*)NULL), cast(TValue*, slot), val);
+      return HOK;  /* success */
+    }
+    else
+      return ~cast_int(key);  /* empty slot in the array part */
+  }
+  else
+    return finishnodeset(t, getintfromhash(t, key), val);
+}
+
+
+int luaH_setshortstr1 (Table *t, TString *key, TValue *val) {
+  return finishnodeset(t, luaH_getshortstr(t, key), val);
+}
+
+
+int luaH_setstr1 (Table *t, TString *key, TValue *val) {
+  return finishnodeset(t, luaH_getstr(t, key), val);
+}
+
+
+int luaH_set1 (Table *t, const TValue *key, TValue *val) {
+  switch (ttypetag(key)) {
+    case LUA_VSHRSTR: return luaH_setshortstr1(t, tsvalue(key), val);
+    case LUA_VNUMINT: return luaH_setint1(t, ivalue(key), val);
+    case LUA_VNIL: return HNOTFOUND;
+    case LUA_VNUMFLT: {
+      lua_Integer k;
+      if (luaV_flttointeger(fltvalue(key), &k, F2Ieq)) /* integral index? */
+        return luaH_setint1(t, k, val);  /* use specialized version */
+      /* else... */
+    }  /* FALLTHROUGH */
+    default:
+      return finishnodeset(t, getgeneric(t, key, 0), val);
+  }
+}
+
 /*
 ** Finish a raw "set table" operation, where 'slot' is where the value
 ** should have been (the result of a previous "get table").
@@ -847,6 +912,21 @@ void luaH_finishset (lua_State *L, Table *t, const TValue *key,
 }
 
 
+void luaH_finishset1 (lua_State *L, Table *t, const TValue *key,
+                                    TValue *value, int aux) {
+  if (aux == HNOTFOUND) {
+    luaH_newkey(L, t, key, value);
+  }
+  else if (aux > 0) {  /* regular Node? */
+    setobj2t(L, gval(gnode(t, aux - HFIRSTNODE)), value);
+  }
+  else {  /* array entry */
+    aux = ~aux;  /* real index */
+    val2arr(t, aux, getArrTag(t, aux), value);
+  }
+}
+
+
 /*
 ** beware: when using this function you probably need to check a GC
 ** barrier and invalidate the TM cache.
diff --git a/ltable.h b/ltable.h
index c8a8c5e7..b9bb416a 100644
--- a/ltable.h
+++ b/ltable.h
@@ -39,6 +39,7 @@
 #define HOK             0
 #define HNOTFOUND       1
 #define HNOTATABLE      2
+#define HFIRSTNODE      3
 
 
 /* fast access to components of array values */
@@ -50,12 +51,20 @@
 #define arr2val(h,k,tag,res)  \
   ((res)->tt_ = tag, (res)->value_ = *getArrVal(h,k))
 
+#define val2arr(h,k,tag,val)  \
+  (*tag = (val)->tt_, *getArrVal(h,k) = (val)->value_)
+
 
 LUAI_FUNC int luaH_getshortstr1 (Table *t, TString *key, TValue *res);
 LUAI_FUNC int luaH_getstr1 (Table *t, TString *key, TValue *res);
 LUAI_FUNC int luaH_get1 (Table *t, const TValue *key, TValue *res);
 LUAI_FUNC int luaH_getint1 (Table *t, lua_Integer key, TValue *res);
 
+LUAI_FUNC int luaH_setint1 (Table *t, lua_Integer key, TValue *val);
+LUAI_FUNC int luaH_setshortstr1 (Table *t, TString *key, TValue *val);
+LUAI_FUNC int luaH_setstr1 (Table *t, TString *key, TValue *val);
+LUAI_FUNC int luaH_set1 (Table *t, const TValue *key, TValue *val);
+
 LUAI_FUNC const TValue *luaH_getint (Table *t, lua_Integer key);
 LUAI_FUNC void luaH_setint (lua_State *L, Table *t, lua_Integer key,
                                                     TValue *value);
@@ -68,6 +77,8 @@ LUAI_FUNC void luaH_set (lua_State *L, Table *t, const TValue *key,
                                                  TValue *value);
 LUAI_FUNC void luaH_finishset (lua_State *L, Table *t, const TValue *key,
                                        const TValue *slot, TValue *value);
+LUAI_FUNC void luaH_finishset1 (lua_State *L, Table *t, const TValue *key,
+                                              TValue *value, int aux);
 LUAI_FUNC Table *luaH_new (lua_State *L);
 LUAI_FUNC void luaH_resize (lua_State *L, Table *t, unsigned int nasize,
                                                     unsigned int nhsize);
diff --git a/lvm.c b/lvm.c
index ee386847..f5934025 100644
--- a/lvm.c
+++ b/lvm.c
@@ -325,17 +325,16 @@ void luaV_finishget1 (lua_State *L, const TValue *t, TValue *key, StkId val,
 ** is no such entry.  (The value at 'slot' must be empty, otherwise
 ** 'luaV_fastget' would have done the job.)
 */
-void luaV_finishset (lua_State *L, const TValue *t, TValue *key,
-                     TValue *val, const TValue *slot) {
+void luaV_finishset1 (lua_State *L, const TValue *t, TValue *key,
+                      TValue *val, int aux) {
   int loop;  /* counter to avoid infinite loops */
   for (loop = 0; loop < MAXTAGLOOP; loop++) {
     const TValue *tm;  /* '__newindex' metamethod */
-    if (slot != NULL) {  /* is 't' a table? */
+    if (aux != HNOTATABLE) {  /* is 't' a table? */
       Table *h = hvalue(t);  /* save 't' table */
-      lua_assert(isempty(slot));  /* slot must be empty */
       tm = fasttm(L, h->metatable, TM_NEWINDEX);  /* get metamethod */
       if (tm == NULL) {  /* no metamethod? */
-        luaH_finishset(L, h, key, slot, val);  /* set new value */
+        luaH_finishset1(L, h, key, val, aux);  /* set new value */
         invalidateTMcache(h);
         luaC_barrierback(L, obj2gco(h), val);
         return;
@@ -353,10 +352,9 @@ void luaV_finishset (lua_State *L, const TValue *t, TValue *key,
       return;
     }
     t = tm;  /* else repeat assignment over 'tm' */
-    if (luaV_fastget(L, t, key, slot, luaH_get)) {
-      luaV_finishfastset(L, t, slot, val);
+    luaV_fastset1(t, key, val, aux, luaH_set1);
+    if (aux == HOK)
       return;  /* done */
-    }
     /* else 'return luaV_finishset(L, t, key, val, slot)' (loop) */
   }
   luaG_runerror(L, "'__newindex' chain too long; possible loop");
@@ -1296,59 +1294,61 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_SETTABUP) {
-        const TValue *slot;
+        int aux;
         TValue *upval = cl->upvals[GETARG_A(i)]->v.p;
         TValue *rb = KB(i);
         TValue *rc = RKC(i);
         TString *key = tsvalue(rb);  /* key must be a string */
-        if (luaV_fastget(L, upval, key, slot, luaH_getshortstr)) {
-          luaV_finishfastset(L, upval, slot, rc);
-        }
+        luaV_fastset1(upval, key, rc, aux, luaH_setshortstr1);
+        if (aux == HOK)
+          luaV_finishfastset1(L, upval, rc);
         else
-          Protect(luaV_finishset(L, upval, rb, rc, slot));
+          Protect(luaV_finishset1(L, upval, rb, rc, aux));
         vmbreak;
       }
       vmcase(OP_SETTABLE) {
         StkId ra = RA(i);
-        const TValue *slot;
+        int aux;
         TValue *rb = vRB(i);  /* key (table is in 'ra') */
         TValue *rc = RKC(i);  /* value */
-        lua_Unsigned n;
-        if (ttisinteger(rb)  /* fast track for integers? */
-            ? (cast_void(n = ivalue(rb)), luaV_fastgeti(L, s2v(ra), n, slot))
-            : luaV_fastget(L, s2v(ra), rb, slot, luaH_get)) {
-          luaV_finishfastset(L, s2v(ra), slot, rc);
+        if (ttisinteger(rb)) {  /* fast track for integers? */
+          luaV_fastseti1(s2v(ra), ivalue(rb), rc, aux);
+        }
+        else {
+          luaV_fastset1(s2v(ra), rb, rc, aux, luaH_set1);
         }
+        if (aux == HOK)
+          luaV_finishfastset1(L, s2v(ra), rc);
         else
-          Protect(luaV_finishset(L, s2v(ra), rb, rc, slot));
+          Protect(luaV_finishset1(L, s2v(ra), rb, rc, aux));
         vmbreak;
       }
       vmcase(OP_SETI) {
         StkId ra = RA(i);
-        const TValue *slot;
-        int c = GETARG_B(i);
+        int aux;
+        int b = GETARG_B(i);
         TValue *rc = RKC(i);
-        if (luaV_fastgeti(L, s2v(ra), c, slot)) {
-          luaV_finishfastset(L, s2v(ra), slot, rc);
-        }
+        luaV_fastseti1(s2v(ra), b, rc, aux);
+        if (aux == HOK)
+          luaV_finishfastset1(L, s2v(ra), rc);
         else {
           TValue key;
-          setivalue(&key, c);
-          Protect(luaV_finishset(L, s2v(ra), &key, rc, slot));
+          setivalue(&key, b);
+          Protect(luaV_finishset1(L, s2v(ra), &key, rc, aux));
         }
         vmbreak;
       }
       vmcase(OP_SETFIELD) {
         StkId ra = RA(i);
-        const TValue *slot;
+        int aux;
         TValue *rb = KB(i);
         TValue *rc = RKC(i);
         TString *key = tsvalue(rb);  /* key must be a string */
-        if (luaV_fastget(L, s2v(ra), key, slot, luaH_getshortstr)) {
-          luaV_finishfastset(L, s2v(ra), slot, rc);
-        }
+        luaV_fastset1(s2v(ra), key, rc, aux, luaH_setshortstr1);
+        if (aux == HOK)
+          luaV_finishfastset1(L, s2v(ra), rc);
         else
-          Protect(luaV_finishset(L, s2v(ra), rb, rc, slot));
+          Protect(luaV_finishset1(L, s2v(ra), rb, rc, aux));
         vmbreak;
       }
       vmcase(OP_NEWTABLE) {
diff --git a/lvm.h b/lvm.h
index 704446c2..750a22b2 100644
--- a/lvm.h
+++ b/lvm.h
@@ -104,14 +104,27 @@ typedef enum {
               ? &hvalue(t)->array[k - 1] : luaH_getint(hvalue(t), k), \
       !isempty(slot)))  /* result not empty? */
 
-#define luaV_fastgeti1(t,k,val,aux) \
+#define luaV_fastgeti1(t,k,res,aux) \
   if (!ttistable(t)) aux = HNOTATABLE; \
   else { Table *h = hvalue(t); lua_Unsigned u = l_castS2U(k); \
     if ((u - 1u < h->alimit)) { \
       int tag = *getArrTag(h,u); \
       if (tagisempty(tag)) aux = HNOTFOUND; \
-      else { arr2val(h, u, tag, val); aux = HOK; }} \
-    else { aux = luaH_getint1(h, u, val); }}
+      else { arr2val(h, u, tag, res); aux = HOK; }} \
+    else { aux = luaH_getint1(h, u, res); }}
+
+
+#define luaV_fastset1(t,k,val,aux,f) \
+  (aux = (!ttistable(t) ? HNOTATABLE : f(hvalue(t), k, val)))
+
+#define luaV_fastseti1(t,k,val,aux) \
+  if (!ttistable(t)) aux = HNOTATABLE; \
+  else { Table *h = hvalue(t); lua_Unsigned u = l_castS2U(k); \
+    if ((u - 1u < h->alimit)) { \
+      lu_byte *tag = getArrTag(h,u); \
+      if (tagisempty(*tag)) aux = ~cast_int(u); \
+      else { val2arr(h, u, tag, val); aux = HOK; }} \
+    else { aux = luaH_setint1(h, u, val); }}
 
 
 /*
@@ -122,6 +135,8 @@ typedef enum {
     { setobj2t(L, cast(TValue *,slot), v); \
       luaC_barrierback(L, gcvalue(t), v); }
 
+#define luaV_finishfastset1(L,t,v)      luaC_barrierback(L, gcvalue(t), v)
+
 
 /*
 ** Shift right is the same as shift left with a negative 'y'
@@ -140,8 +155,8 @@ LUAI_FUNC int luaV_tointegerns (const TValue *obj, lua_Integer *p,
 LUAI_FUNC int luaV_flttointeger (lua_Number n, lua_Integer *p, F2Imod mode);
 LUAI_FUNC void luaV_finishget1 (lua_State *L, const TValue *t, TValue *key,
                                               StkId val, int aux);
-LUAI_FUNC void luaV_finishset (lua_State *L, const TValue *t, TValue *key,
-                               TValue *val, const TValue *slot);
+LUAI_FUNC void luaV_finishset1 (lua_State *L, const TValue *t, TValue *key,
+                                              TValue *val, int aux);
 LUAI_FUNC void luaV_finishOp (lua_State *L);
 LUAI_FUNC void luaV_execute (lua_State *L, CallInfo *ci);
 LUAI_FUNC void luaV_concat (lua_State *L, int total);