From b7edf5d2d89ed2ce1e9087de496bcb451e39d131 Mon Sep 17 00:00:00 2001
From: Roberto Ierusalimschy <roberto@inf.puc-rio.br>
Date: Tue, 27 Feb 2018 14:48:28 -0300
Subject: [PATCH] metamethods for 'removekey'/'keyin'

---
 lapi.c     | 19 ++++++++-----------
 lbaselib.c |  8 +++-----
 ltm.c      | 31 ++++++++++++++++++++++++++++++-
 ltm.h      |  6 +++++-
 4 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/lapi.c b/lapi.c
index 92e98def..7ecf68cb 100644
--- a/lapi.c
+++ b/lapi.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lapi.c,v 2.287 2018/02/25 12:48:16 roberto Exp roberto $
+** $Id: lapi.c,v 2.288 2018/02/26 14:16:05 roberto Exp roberto $
 ** Lua API
 ** See Copyright Notice in lua.h
 */
@@ -679,6 +679,7 @@ LUA_API int lua_rawget (lua_State *L, int idx) {
   Table *t;
   const TValue *val;
   lua_lock(L);
+  api_checknelems(L, 1);
   t = gettable(L, idx);
   val = luaH_get(t, s2v(L->top - 1));
   L->top--;  /* remove key */
@@ -704,29 +705,24 @@ LUA_API int lua_rawgetp (lua_State *L, int idx, const void *p) {
 }
 
 
-static int auxkeyman (lua_State *L, int idx, int remove) {
-  Table *t;
-  const TValue *val;
+static int auxkeydef (lua_State *L, int idx, int remove) {
   int res;
   lua_lock(L);
-  t = gettable(L, idx);
-  val = luaH_get(t, s2v(L->top - 1));
+  api_checknelems(L, 1);
+  res = luaT_keydef(L, index2value(L, idx), s2v(L->top - 1), remove);
   L->top--;  /* remove key */
-  res = !isempty(val);
-  if (remove && res)  /* key is present and should be removed? */
-    setempty(cast(TValue*, val));
   lua_unlock(L);
   return res;
 }
 
 
 LUA_API void lua_removekey (lua_State *L, int idx) {
-  auxkeyman(L, idx, 1);
+  auxkeydef(L, idx, 1);
 }
 
 
 LUA_API int lua_keyin (lua_State *L, int idx) {
-  return auxkeyman(L, idx, 0);
+  return auxkeydef(L, idx, 0);
 }
 
 
@@ -1223,6 +1219,7 @@ LUA_API int lua_next (lua_State *L, int idx) {
   Table *t;
   int more;
   lua_lock(L);
+  api_checknelems(L, 1);
   t = gettable(L, idx);
   more = luaH_next(L, t, L->top - 1);
   if (more) {
diff --git a/lbaselib.c b/lbaselib.c
index 443bf066..8e033865 100644
--- a/lbaselib.c
+++ b/lbaselib.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lbaselib.c,v 1.319 2018/02/05 17:10:52 roberto Exp roberto $
+** $Id: lbaselib.c,v 1.320 2018/02/25 12:48:16 roberto Exp roberto $
 ** Basic library
 ** See Copyright Notice in lua.h
 */
@@ -171,8 +171,7 @@ static int luaB_rawset (lua_State *L) {
 
 
 static int luaB_keyin (lua_State *L) {
-  luaL_checktype(L, 1, LUA_TTABLE);
-  luaL_checkany(L, 2);
+  luaL_checkany(L, 2);  /* ensures a first argument too */
   lua_settop(L, 2);
   lua_pushboolean(L, lua_keyin(L, 1));
   return 1;
@@ -180,8 +179,7 @@ static int luaB_keyin (lua_State *L) {
 
 
 static int luaB_removekey (lua_State *L) {
-  luaL_checktype(L, 1, LUA_TTABLE);
-  luaL_checkany(L, 2);
+  luaL_checkany(L, 2);  /* ensures a first argument too */
   lua_settop(L, 2);
   lua_removekey(L, 1);
   return 0;
diff --git a/ltm.c b/ltm.c
index c06c8296..2a9a4cbe 100644
--- a/ltm.c
+++ b/ltm.c
@@ -1,5 +1,5 @@
 /*
-** $Id: ltm.c,v 2.64 2018/02/23 13:13:31 roberto Exp roberto $
+** $Id: ltm.c,v 2.65 2018/02/26 14:16:05 roberto Exp roberto $
 ** Tag methods
 ** See Copyright Notice in lua.h
 */
@@ -38,6 +38,7 @@ LUAI_DDEF const char *const luaT_typenames_[LUA_TOTALTAGS] = {
 void luaT_init (lua_State *L) {
   static const char *const luaT_eventname[] = {  /* ORDER TM */
     "__index", "__newindex",
+    "__undef", "__isdef",
     "__gc", "__mode", "__len", "__eq",
     "__add", "__sub", "__mul", "__mod", "__pow",
     "__div", "__idiv",
@@ -248,3 +249,31 @@ void luaT_getvarargs (lua_State *L, CallInfo *ci, StkId where, int wanted) {
   for (; i < wanted; i++)   /* complete required results with nil */
     setnilvalue(s2v(where + i));
 }
+
+
+int luaT_keydef (lua_State *L, TValue *obj, TValue *key, int remove) {
+  const TValue *tm;
+  TMS event = remove ? TM_UNDEF : TM_ISDEF;
+  if (!ttistable(obj)) {  /* not a table? */
+    tm = luaT_gettmbyobj(L, obj, event);  /* get its metamethod */
+    if (notm(tm)) {  /* no metamethod? */
+      const char *msg = remove ? "remove key from" : "check key from";
+      luaG_typeerror(L, obj, msg);  /* error */
+    }
+    /* else will call metamethod 'tm' */
+  }
+  else {  /* 'obj' is a table */
+    Table *t = hvalue(obj);
+    tm = fasttm(L, t->metatable, event);
+    if (tm == NULL) {  /* no metamethod? */
+      const TValue *val = luaH_get(t, key);  /* get entry */
+      int res = !isempty(val);  /* true if entry is not empty */
+      if (remove && res)  /* key is present and should be removed? */
+        setempty(cast(TValue*, val));  /* remove it */
+      return res;
+    }
+    /* else will call metamethod 'tm' */
+  }
+  luaT_callTMres(L, tm, obj, key, L->top);
+  return !l_isfalse(s2v(L->top));
+}
diff --git a/ltm.h b/ltm.h
index 9b25ef08..6e961d27 100644
--- a/ltm.h
+++ b/ltm.h
@@ -1,5 +1,5 @@
 /*
-** $Id: ltm.h,v 2.32 2018/02/17 19:20:00 roberto Exp roberto $
+** $Id: ltm.h,v 2.33 2018/02/23 13:13:31 roberto Exp roberto $
 ** Tag methods
 ** See Copyright Notice in lua.h
 */
@@ -19,6 +19,8 @@
 typedef enum {
   TM_INDEX,
   TM_NEWINDEX,
+  TM_UNDEF,
+  TM_ISDEF,
   TM_GC,
   TM_MODE,
   TM_LEN,
@@ -89,5 +91,7 @@ LUAI_FUNC void luaT_adjustvarargs (lua_State *L, int nfixparams,
 LUAI_FUNC void luaT_getvarargs (lua_State *L, struct CallInfo *ci,
                                               StkId where, int wanted);
 
+LUAI_FUNC int luaT_keydef (lua_State *L, TValue *obj, TValue *key, int remove);
+
 
 #endif