To-be-closed variables in the C API

6 files changed, 122 insertions, 16 deletions

diff --git a/lapi.c b/lapi.c
index ae6b07ae..4fef43b7 100644
--- a/lapi.c
+++ b/lapi.c
@@ -173,15 +173,17 @@ LUA_API void lua_settop (lua_State *L, int idx) {
   StkId func = L->ci->func;
   lua_lock(L);
   if (idx >= 0) {
+    StkId newtop = (func + 1) + idx;
     api_check(L, idx <= L->stack_last - (func + 1), "new top too large");
-    while (L->top < (func + 1) + idx)
+    while (L->top < newtop)
       setnilvalue(s2v(L->top++));
-    L->top = (func + 1) + idx;
+    L->top = newtop;
   }
   else {
     api_check(L, -(idx+1) <= (L->top - (func + 1)), "invalid new top");
     L->top += idx+1;  /* 'subtract' index (index is negative) */
   }
+  luaF_close(L, L->top, LUA_OK);
   lua_unlock(L);
 }
 
@@ -1205,6 +1207,15 @@ LUA_API int lua_next (lua_State *L, int idx) {
 }
 
 
+LUA_API void lua_tobeclosed (lua_State *L) {
+  int nresults = L->ci->nresults;
+  luaF_newtbcupval(L, L->top - 1);  /* create new to-be-closed upvalue */
+  if (!hastocloseCfunc(nresults))  /* function not marked yet? */
+    L->ci->nresults = codeNresults(nresults);  /* mark it */
+  lua_assert(hastocloseCfunc(L->ci->nresults));
+}
+
+
 LUA_API void lua_concat (lua_State *L, int n) {
   lua_lock(L);
   api_checknelems(L, n);
diff --git a/lapi.h b/lapi.h
index 016f78cc..5a4206f1 100644
--- a/lapi.h
+++ b/lapi.h
@@ -15,10 +15,23 @@
                                 "stack overflow");}
 
 #define adjustresults(L,nres) \
-    { if ((nres) == LUA_MULTRET && L->ci->top < L->top) L->ci->top = L->top; }
+    { if ((nres) <= LUA_MULTRET && L->ci->top < L->top) L->ci->top = L->top; }
 
 #define api_checknelems(L,n)    api_check(L, (n) < (L->top - L->ci->func), \
                                   "not enough elements in the stack")
 
 
+/*
+** To reduce the overhead of returning from C functions, the presence of
+** to-be-closed variables in these functions is coded in the CallInfo's
+** field 'nresults', in a way that functions with no to-be-closed variables
+** with zero, one, or "all" wanted results have no overhead. Functions
+** with other number of wanted results, as well as functions with
+** variables to be closed, have an extra check.
+*/
+
+#define hastocloseCfunc(n)      ((n) < LUA_MULTRET)
+
+#define codeNresults(n)         (-(n) - 3)
+
 #endif
diff --git a/ldo.c b/ldo.c
index 78e4c5c3..b7a76ef6 100644
--- a/ldo.c
+++ b/ldo.c
@@ -366,32 +366,38 @@ void luaD_tryfuncTM (lua_State *L, StkId func) {
 ** separated.
 */
 static void moveresults (lua_State *L, StkId res, int nres, int wanted) {
+  StkId firstresult;
+  int i;
   switch (wanted) {  /* handle typical cases separately */
     case 0:  /* no values needed */
       L->top = res;
-      break;
+      return;
     case 1:  /* one value needed */
       if (nres == 0)   /* no results? */
         setnilvalue(s2v(res));  /* adjust with nil */
       else
         setobjs2s(L, res, L->top - nres);  /* move it to proper place */
       L->top = res + 1;
-      break;
+      return;
     case LUA_MULTRET:
       wanted = nres;  /* we want all results */
-      /* FALLTHROUGH */
-    default: {  /* multiple results */
-      StkId firstresult = L->top - nres;  /* index of first result */
-      int i;
-      /* move all results to correct place */
-      for (i = 0; i < nres && i < wanted; i++)
-        setobjs2s(L, res + i, firstresult + i);
-      for (; i < wanted; i++)  /* complete wanted number of results */
-        setnilvalue(s2v(res + i));
-      L->top = res + wanted;  /* top points after the last result */
       break;
-    }
+    default:  /* multiple results (or to-be-closed variables) */
+      if (hastocloseCfunc(wanted)) {
+        luaF_close(L, res, LUA_OK);
+        wanted = codeNresults(wanted);  /* correct value */
+        if (wanted == LUA_MULTRET)
+          wanted = nres;
+      }
+      break;
   }
+  firstresult = L->top - nres;  /* index of first result */
+  /* move all results to correct place */
+  for (i = 0; i < nres && i < wanted; i++)
+    setobjs2s(L, res + i, firstresult + i);
+  for (; i < wanted; i++)  /* complete wanted number of results */
+    setnilvalue(s2v(res + i));
+  L->top = res + wanted;  /* top points after the last result */
 }
 
 
diff --git a/ltests.c b/ltests.c
index a6968653..fbcb7475 100644
--- a/ltests.c
+++ b/ltests.c
@@ -1554,6 +1554,9 @@ static struct X { int x; } x;
       int i = getindex;
       return lua_yieldk(L1, nres, i, Cfunck);
     }
+    else if EQ("tobeclosed") {
+      lua_tobeclosed(L);
+    }
     else luaL_error(L, "unknown instruction %s", buff);
   }
   return 0;
diff --git a/lua.h b/lua.h
index a014be1f..16d685cc 100644
--- a/lua.h
+++ b/lua.h
@@ -333,6 +333,8 @@ LUA_API size_t   (lua_stringtonumber) (lua_State *L, const char *s);
 LUA_API lua_Alloc (lua_getallocf) (lua_State *L, void **ud);
 LUA_API void      (lua_setallocf) (lua_State *L, lua_Alloc f, void *ud);
 
+LUA_API void  (lua_tobeclosed) (lua_State *L);
+
 
 /*
 ** {==============================================================
diff --git a/testes/api.lua b/testes/api.lua
index 6e11c6ed..988250f7 100644
--- a/testes/api.lua
+++ b/testes/api.lua
@@ -967,6 +967,77 @@ T.closestate(L1)
 L1 = nil
 
 print('+')
+-------------------------------------------------------------------------
+-- testing to-be-closed variables
+-------------------------------------------------------------------------
+print"testing to-be-closed variables"
+
+do
+  local openresource = {}
+
+  local function newresource ()
+    local x = setmetatable({10}, {__close = function(y)
+      assert(openresource[#openresource] == y)
+      openresource[#openresource] = nil
+      y[1] = y[1] + 1
+    end})
+    openresource[#openresource + 1] = x
+    return x
+  end
+
+  local a = T.testC([[
+    call 0 1   # create resource
+    tobeclosed  # mark it to be closed
+    return 1
+  ]], newresource)
+  assert(a[1] == 11)
+  assert(#openresource == 0)    -- was closed
+
+  -- repeat the test, but calling function in a 'multret' context
+  local a = {T.testC([[
+    call 0 1   # create resource
+    tobeclosed  # mark it to be closed
+    return 2
+  ]], newresource)}
+  assert(type(a[1]) == "string" and a[2][1] == 11)
+  assert(#openresource == 0)    -- was closed
+
+  -- error
+  local a, b = pcall(T.testC, [[
+    call 0 1   # create resource
+    tobeclosed  # mark it to be closed
+    error       # resource is the error object
+  ]], newresource)
+  assert(a == false and b[1] == 11)
+  assert(#openresource == 0)    -- was closed
+
+  local function check (n)
+    assert(#openresource == n)
+  end
+
+  -- closing resources with 'settop'
+  local a = T.testC([[
+    pushvalue 2
+    call 0 1   # create resource
+    tobeclosed  # mark it to be closed
+    pushvalue 2
+    call 0 1   # create another resource
+    tobeclosed  # mark it to be closed
+    pushvalue 3
+    pushint 2   # there should be two open resources
+    call 1 0
+    pop 1       # pop second resource from the stack
+    pushvalue 3
+    pushint 1   # there should be one open resource
+    call 1 0
+    pop 1       # pop second resource from the stack
+    pushint *
+    return 1    # return stack size
+  ]], newresource, check)
+  assert(a == 3)   -- no extra items left in the stack
+
+end
+
 
 -------------------------------------------------------------------------
 -- testing memory limits