Allow yields inside '__close' metamethods

Initial implementation to allow yields inside '__close' metamethods.
This current version still does not allow a '__close' metamethod
to yield when called due to an error. '__close' metamethods from
C functions also are not allowed to yield.

7 files changed, 114 insertions, 18 deletions

diff --git a/lapi.c b/lapi.c
index 0f0e31af..3583e9c0 100644
--- a/lapi.c
+++ b/lapi.c
@@ -189,7 +189,7 @@ LUA_API void lua_settop (lua_State *L, int idx) {
   }
 #if defined(LUA_COMPAT_5_4_0)
   if (diff < 0 && hastocloseCfunc(ci->nresults))
-    luaF_close(L, L->top + diff, CLOSEKTOP);
+    luaF_close(L, L->top + diff, CLOSEKTOP, 0);
 #endif
   L->top += diff;
   api_check(L, L->openupval == NULL || uplevel(L->openupval) < L->top,
@@ -205,7 +205,7 @@ LUA_API void lua_closeslot (lua_State *L, int idx) {
   api_check(L, hastocloseCfunc(L->ci->nresults) && L->openupval != NULL &&
                uplevel(L->openupval) == level,
      "no variable to close at given level");
-  luaF_close(L, level, CLOSEKTOP);
+  luaF_close(L, level, CLOSEKTOP, 0);
   setnilvalue(s2v(level));
   lua_unlock(L);
 }
diff --git a/ldo.c b/ldo.c
index ba0c93b8..aa159cf0 100644
--- a/ldo.c
+++ b/ldo.c
@@ -406,7 +406,7 @@ static void moveresults (lua_State *L, StkId res, int nres, int wanted) {
     default:  /* multiple results (or to-be-closed variables) */
       if (hastocloseCfunc(wanted)) {  /* to-be-closed variables? */
         ptrdiff_t savedres = savestack(L, res);
-        luaF_close(L, res, CLOSEKTOP);  /* may change the stack */
+        luaF_close(L, res, CLOSEKTOP, 0);  /* may change the stack */
         res = restorestack(L, savedres);
         wanted = codeNresults(wanted);  /* correct value */
         if (wanted == LUA_MULTRET)
@@ -647,7 +647,7 @@ static void recover (lua_State *L, void *ud) {
   /* "finish" luaD_pcall */
   L->ci = ci;
   L->allowhook = getoah(ci->callstatus);  /* restore original 'allowhook' */
-  luaF_close(L, func, status);  /* may change the stack */
+  luaF_close(L, func, status, 0);  /* may change the stack */
   func = restorestack(L, ci->u2.funcidx);
   luaD_seterrorobj(L, status, func);
   luaD_shrinkstack(L);   /* restore stack size in case of overflow */
@@ -803,7 +803,7 @@ struct CloseP {
 */
 static void closepaux (lua_State *L, void *ud) {
   struct CloseP *pcl = cast(struct CloseP *, ud);
-  luaF_close(L, pcl->level, pcl->status);
+  luaF_close(L, pcl->level, pcl->status, 0);
 }
 
 
diff --git a/lfunc.c b/lfunc.c
index a8030afa..13e44d46 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -101,17 +101,21 @@ UpVal *luaF_findupval (lua_State *L, StkId level) {
 
 
 /*
-** Call closing method for object 'obj' with error message 'err'.
+** Call closing method for object 'obj' with error message 'err'. The
+** boolean 'yy' controls whether the call is yieldable.
 ** (This function assumes EXTRA_STACK.)
 */
-static void callclosemethod (lua_State *L, TValue *obj, TValue *err) {
+static void callclosemethod (lua_State *L, TValue *obj, TValue *err, int yy) {
   StkId top = L->top;
   const TValue *tm = luaT_gettmbyobj(L, obj, TM_CLOSE);
   setobj2s(L, top, tm);  /* will call metamethod... */
   setobj2s(L, top + 1, obj);  /* with 'self' as the 1st argument */
   setobj2s(L, top + 2, err);  /* and error msg. as 2nd argument */
   L->top = top + 3;  /* add function and arguments */
-  luaD_callnoyield(L, top, 0);  /* call method */
+  if (yy)
+    luaD_call(L, top, 0);
+  else
+    luaD_callnoyield(L, top, 0);
 }
 
 
@@ -137,7 +141,7 @@ static void checkclosemth (lua_State *L, StkId level, const TValue *obj) {
 ** the 'level' of the upvalue being closed, as everything after that
 ** won't be used again.
 */
-static void prepcallclosemth (lua_State *L, StkId level, int status) {
+static void prepcallclosemth (lua_State *L, StkId level, int status, int yy) {
   TValue *uv = s2v(level);  /* value being closed */
   TValue *errobj;
   if (status == CLOSEKTOP)
@@ -146,7 +150,7 @@ static void prepcallclosemth (lua_State *L, StkId level, int status) {
     errobj = s2v(level + 1);  /* error object goes after 'uv' */
     luaD_seterrorobj(L, status, level + 1);  /* set error object */
   }
-  callclosemethod(L, uv, errobj);
+  callclosemethod(L, uv, errobj, yy);
 }
 
 
@@ -174,7 +178,7 @@ void luaF_newtbcupval (lua_State *L, StkId level) {
     if (unlikely(status != LUA_OK)) {  /* memory error creating upvalue? */
       lua_assert(status == LUA_ERRMEM);
       luaD_seterrorobj(L, LUA_ERRMEM, level + 1);  /* save error message */
-      callclosemethod(L, s2v(level), s2v(level + 1));
+      callclosemethod(L, s2v(level), s2v(level + 1), 0);
       luaD_throw(L, LUA_ERRMEM);  /* throw memory error */
     }
   }
@@ -194,7 +198,7 @@ void luaF_unlinkupval (UpVal *uv) {
 ** to NOCLOSINGMETH closes upvalues without running any __close
 ** metamethods.
 */
-void luaF_close (lua_State *L, StkId level, int status) {
+void luaF_close (lua_State *L, StkId level, int status, int yy) {
   UpVal *uv;
   StkId upl;  /* stack index pointed by 'uv' */
   while ((uv = L->openupval) != NULL && (upl = uplevel(uv)) >= level) {
@@ -209,7 +213,7 @@ void luaF_close (lua_State *L, StkId level, int status) {
     }
     if (uv->tbc && status != NOCLOSINGMETH) {
       ptrdiff_t levelrel = savestack(L, level);
-      prepcallclosemth(L, upl, status);  /* may change the stack */
+      prepcallclosemth(L, upl, status, yy);  /* may change the stack */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lfunc.h b/lfunc.h
index 40de4636..2e6df535 100644
--- a/lfunc.h
+++ b/lfunc.h
@@ -59,7 +59,7 @@ LUAI_FUNC LClosure *luaF_newLclosure (lua_State *L, int nupvals);
 LUAI_FUNC void luaF_initupvals (lua_State *L, LClosure *cl);
 LUAI_FUNC UpVal *luaF_findupval (lua_State *L, StkId level);
 LUAI_FUNC void luaF_newtbcupval (lua_State *L, StkId level);
-LUAI_FUNC void luaF_close (lua_State *L, StkId level, int status);
+LUAI_FUNC void luaF_close (lua_State *L, StkId level, int status, int yy);
 LUAI_FUNC void luaF_unlinkupval (UpVal *uv);
 LUAI_FUNC void luaF_freeproto (lua_State *L, Proto *f);
 LUAI_FUNC const char *luaF_getlocalname (const Proto *func, int local_number,
diff --git a/lstate.c b/lstate.c
index a6ef82a3..92ccbf9b 100644
--- a/lstate.c
+++ b/lstate.c
@@ -313,7 +313,7 @@ LUA_API lua_State *lua_newthread (lua_State *L) {
 
 void luaE_freethread (lua_State *L, lua_State *L1) {
   LX *l = fromstate(L1);
-  luaF_close(L1, L1->stack, NOCLOSINGMETH);  /* close all upvalues */
+  luaF_close(L1, L1->stack, NOCLOSINGMETH, 0);  /* close all upvalues */
   lua_assert(L1->openupval == NULL);
   luai_userstatefree(L, L1);
   freestack(L1);
diff --git a/lvm.c b/lvm.c
index a6f04606..d6c05bbd 100644
--- a/lvm.c
+++ b/lvm.c
@@ -842,6 +842,10 @@ void luaV_finishOp (lua_State *L) {
       luaV_concat(L, total);  /* concat them (may yield again) */
       break;
     }
+    case OP_CLOSE:  case OP_RETURN: {  /* yielded closing variables */
+      ci->u.l.savedpc--;  /* repeat instruction to close other vars. */
+      break;
+    }
     default: {
       /* only these other opcodes can yield */
       lua_assert(op == OP_TFORCALL || op == OP_CALL ||
@@ -1524,7 +1528,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_CLOSE) {
-        Protect(luaF_close(L, ra, LUA_OK));
+        Protect(luaF_close(L, ra, LUA_OK, 1));
         vmbreak;
       }
       vmcase(OP_TBC) {
@@ -1632,7 +1636,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           /* close upvalues from current call; the compiler ensures
              that there are no to-be-closed variables here, so this
              call cannot change the stack */
-          luaF_close(L, base, NOCLOSINGMETH);
+          luaF_close(L, base, NOCLOSINGMETH, 0);
           lua_assert(base == ci->func + 1);
         }
         while (!ttisfunction(s2v(ra))) {  /* not a function? */
@@ -1662,7 +1666,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         if (TESTARG_k(i)) {  /* may there be open upvalues? */
           if (L->top < ci->top)
             L->top = ci->top;
-          luaF_close(L, base, CLOSEKTOP);
+          luaF_close(L, base, CLOSEKTOP, 1);
           updatetrap(ci);
           updatestack(ci);
         }
diff --git a/testes/locals.lua b/testes/locals.lua
index add023ca..c9c93ccf 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -641,6 +641,94 @@ end
 print "to-be-closed variables in coroutines"
 
 do
+  -- yielding inside closing metamethods
+
+  local function checktable (t1, t2)
+    assert(#t1 == #t2)
+    for i = 1, #t1 do
+      assert(t1[i] == t2[i])
+    end
+  end
+
+  local trace = {}
+  local co = coroutine.wrap(function ()
+
+    trace[#trace + 1] = "nowX"
+
+    -- will be closed after 'y'
+    local x <close> = func2close(function (_, msg)
+      assert(msg == nil)
+      trace[#trace + 1] = "x1"
+      coroutine.yield("x")
+      trace[#trace + 1] = "x2"
+    end)
+
+    return pcall(function ()
+      do   -- 'z' will be closed first
+        local z <close> = func2close(function (_, msg)
+          assert(msg == nil)
+          trace[#trace + 1] = "z1"
+          coroutine.yield("z")
+          trace[#trace + 1] = "z2"
+        end)
+      end
+
+      trace[#trace + 1] = "nowY"
+
+      -- will be closed after 'z'
+      local y <close> = func2close(function(_, msg)
+        assert(msg == nil)
+        trace[#trace + 1] = "y1"
+        coroutine.yield("y")
+        trace[#trace + 1] = "y2"
+      end)
+
+      return 10, 20, 30
+    end)
+  end)
+
+  assert(co() == "z")
+  assert(co() == "y")
+  assert(co() == "x")
+  checktable({co()}, {true, 10, 20, 30})
+  checktable(trace, {"nowX", "z1", "z2", "nowY", "y1", "y2", "x1", "x2"})
+
+end 
+ 
+
+do
+  -- yielding inside closing metamethods after an error:
+  -- not yet implemented; raises an error
+
+  local co = coroutine.wrap(function ()
+
+    local function foo (err)
+
+      local x <close> = func2close(function(_, msg)
+        assert(msg == err)
+        coroutine.yield("x")
+        return 100, 200
+      end)
+
+      if err then error(err) else return 10, 20 end
+    end
+
+    coroutine.yield(pcall(foo, nil))  -- no error
+    return pcall(foo, 10)     -- 'foo' will raise an error
+  end)
+
+  local a, b = co()
+  assert(a == "x" and b == nil)    -- yields inside 'x'; Ok
+
+  local a, b, c = co()
+  assert(a and b == 10 and c == 20)   -- returns from 'pcall(foo, nil)'
+
+  local st, msg = co()    -- error yielding after an error
+  assert(not st and string.find(msg, "attempt to yield"))
+end
+
+
+do
   -- an error in a wrapped coroutine closes variables
   local x = false
   local y = false