New functions 'lua_resetthread' and 'coroutine.kill'

New functions to reset/kill a thread/coroutine, mainly (only?) to
close any pending to-be-closed variable. ('lua_resetthread' also
allows a thread to be reused...)

11 files changed, 195 insertions, 32 deletions

diff --git a/lcorolib.c b/lcorolib.c
index 34462b53..cdb5fedc 100644
--- a/lcorolib.c
+++ b/lcorolib.c
@@ -107,29 +107,40 @@ static int luaB_yield (lua_State *L) {
 }
 
 
-static int luaB_costatus (lua_State *L) {
-  lua_State *co = getco(L);
-  if (L == co) lua_pushliteral(L, "running");
+#define COS_RUN         0
+#define COS_DEAD        1
+#define COS_YIELD       2
+#define COS_NORM        3
+
+
+static const char *statname[] = {"running", "dead", "suspended", "normal"};
+
+
+static int auxstatus (lua_State *L, lua_State *co) {
+  if (L == co) return COS_RUN;
   else {
     switch (lua_status(co)) {
       case LUA_YIELD:
-        lua_pushliteral(L, "suspended");
-        break;
+        return COS_YIELD;
       case LUA_OK: {
         lua_Debug ar;
-        if (lua_getstack(co, 0, &ar) > 0)  /* does it have frames? */
-          lua_pushliteral(L, "normal");  /* it is running */
+        if (lua_getstack(co, 0, &ar))  /* does it have frames? */
+          return COS_NORM;  /* it is running */
         else if (lua_gettop(co) == 0)
-            lua_pushliteral(L, "dead");
+            return COS_DEAD;
         else
-          lua_pushliteral(L, "suspended");  /* initial state */
-        break;
+          return COS_YIELD;  /* initial state */
       }
       default:  /* some error occurred */
-        lua_pushliteral(L, "dead");
-        break;
+        return COS_DEAD;
     }
   }
+}
+
+
+static int luaB_costatus (lua_State *L) {
+  lua_State *co = getco(L);
+  lua_pushstring(L, statname[auxstatus(L, co)]);
   return 1;
 }
 
@@ -147,6 +158,28 @@ static int luaB_corunning (lua_State *L) {
 }
 
 
+static int luaB_kill (lua_State *L) {
+  lua_State *co = getco(L);
+  int status = auxstatus(L, co);
+  switch (status) {
+    case COS_DEAD: case COS_YIELD: {
+      status = lua_resetthread(co);
+      if (status == LUA_OK) {
+        lua_pushboolean(L, 1);
+        return 1;
+      }
+      else {
+        lua_pushboolean(L, 0);
+        lua_xmove(co, L, 1);  /* copy error message */
+        return 2;
+      }
+    }
+    default:  /* normal or running coroutine */
+      return luaL_error(L, "cannot kill a %s coroutine", statname[status]);
+  }
+}
+
+
 static const luaL_Reg co_funcs[] = {
   {"create", luaB_cocreate},
   {"resume", luaB_coresume},
@@ -155,6 +188,7 @@ static const luaL_Reg co_funcs[] = {
   {"wrap", luaB_cowrap},
   {"yield", luaB_yield},
   {"isyieldable", luaB_yieldable},
+  {"kill", luaB_kill},
   {NULL, NULL}
 };
 
diff --git a/ldo.c b/ldo.c
index bdd7fb6d..056fef0c 100644
--- a/ldo.c
+++ b/ldo.c
@@ -98,6 +98,10 @@ void luaD_seterrorobj (lua_State *L, int errcode, StkId oldtop) {
       setsvalue2s(L, oldtop, luaS_newliteral(L, "error in error handling"));
       break;
     }
+    case CLOSEPROTECT: {
+      setnilvalue(s2v(oldtop));  /* no error message */
+      break;
+    }
     default: {
       setobjs2s(L, oldtop, L->top - 1);  /* error message on current top */
       break;
diff --git a/lfunc.c b/lfunc.c
index 11d2850f..bdf3cd25 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -127,17 +127,18 @@ static int prepclosingmethod (lua_State *L, TValue *obj, TValue *err) {
 
 
 /*
-** Prepare and call a closing method. If status is OK, code is
-** still inside the original protected call, and so any error
-** will be handled there. Otherwise, a previous error already
-** activated original protected call, and so the call to the
-** closing method must be protected here.
+** Prepare and call a closing method. If status is OK, code is still
+** inside the original protected call, and so any error will be handled
+** there. Otherwise, a previous error already activated original
+** protected call, and so the call to the closing method must be
+** protected here. (A status = CLOSEPROTECT behaves like a previous
+** error, to also run the closing method in protected mode).
 ** If status is OK, the call to the closing method will be pushed
 ** at the top of the stack. Otherwise, values are pushed after
 ** the 'level' of the upvalue being closed, as everything after
 ** that won't be used again.
 */
-static int closeupval (lua_State *L, TValue *uv, StkId level, int status) {
+static int callclosemth (lua_State *L, TValue *uv, StkId level, int status) {
   if (likely(status == LUA_OK)) {
     if (prepclosingmethod(L, uv, &G(L)->nilvalue))  /* something to call? */
       callclose(L, NULL);  /* call closing method */
@@ -207,9 +208,10 @@ int luaF_close (lua_State *L, StkId level, int status) {
     if (!iswhite(uv))
       gray2black(uv);  /* closed upvalues cannot be gray */
     luaC_barrier(L, uv, slot);
-    if (status >= 0 && uv->tt == LUA_TUPVALTBC) {  /* must be closed? */
+    if (uv->tt == LUA_TUPVALTBC && status != NOCLOSINGMETH) {
+      /* must run closing method */
       ptrdiff_t levelrel = savestack(L, level);
-      status = closeupval(L, uv->v, upl, status);  /* may realloc. the stack */
+      status = callclosemth(L, uv->v, upl, status);  /* may change the stack */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lfunc.h b/lfunc.h
index c9fe1314..0ed79c48 100644
--- a/lfunc.h
+++ b/lfunc.h
@@ -42,6 +42,17 @@
 #define MAXMISS         10
 
 
+/*
+** Special "status" for 'luaF_close'
+*/
+
+/* close upvalues without running their closing methods */
+#define NOCLOSINGMETH   (-1)
+
+/* close upvalues running all closing methods in protected mode */
+#define CLOSEPROTECT    (-2)
+
+
 LUAI_FUNC Proto *luaF_newproto (lua_State *L);
 LUAI_FUNC CClosure *luaF_newCclosure (lua_State *L, int nelems);
 LUAI_FUNC LClosure *luaF_newLclosure (lua_State *L, int nelems);
diff --git a/lstate.c b/lstate.c
index 9d399959..5ee024fc 100644
--- a/lstate.c
+++ b/lstate.c
@@ -258,7 +258,7 @@ static void preinit_thread (lua_State *L, global_State *g) {
 
 static void close_state (lua_State *L) {
   global_State *g = G(L);
-  luaF_close(L, L->stack, -1);  /* close all upvalues for this thread */
+  luaF_close(L, L->stack, CLOSEPROTECT);  /* close all upvalues */
   luaC_freeallobjects(L);  /* collect all objects */
   if (ttisnil(&g->nilvalue))  /* closing a fully built state? */
     luai_userstateclose(L);
@@ -301,7 +301,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, -1);  /* close all upvalues for this thread */
+  luaF_close(L1, L1->stack, NOCLOSINGMETH);  /* close all upvalues */
   lua_assert(L1->openupval == NULL);
   luai_userstatefree(L, L1);
   freestack(L1);
@@ -309,6 +309,29 @@ void luaE_freethread (lua_State *L, lua_State *L1) {
 }
 
 
+int lua_resetthread (lua_State *L) {
+  CallInfo *ci;
+  int status;
+  lua_lock(L);
+  ci = &L->base_ci;
+  status = luaF_close(L, L->stack, CLOSEPROTECT);
+  setnilvalue(s2v(L->stack));  /* 'function' entry for basic 'ci' */
+  if (status != CLOSEPROTECT)  /* real errors? */
+    luaD_seterrorobj(L, status, L->stack + 1);
+  else {
+    status = LUA_OK;
+    L->top = L->stack + 1;
+  }
+  ci->callstatus = CIST_C;
+  ci->func = L->stack;
+  ci->top = L->top + LUA_MINSTACK;
+  L->ci = ci;
+  L->status = status;
+  lua_unlock(L);
+  return status;
+}
+
+
 LUA_API lua_State *lua_newstate (lua_Alloc f, void *ud) {
   int i;
   lua_State *L;
diff --git a/ltests.c b/ltests.c
index 63d423e0..a38a8926 100644
--- a/ltests.c
+++ b/ltests.c
@@ -1366,6 +1366,9 @@ static int runC (lua_State *L, lua_State *L1, const char *pc) {
     else if EQ("newthread") {
       lua_newthread(L1);
     }
+    else if EQ("resetthread") {
+      lua_pushinteger(L1, lua_resetthread(L1));
+    }
     else if EQ("newuserdata") {
       lua_newuserdata(L1, getnum);
     }
diff --git a/lua.h b/lua.h
index 6aa184d1..95ce5a2e 100644
--- a/lua.h
+++ b/lua.h
@@ -147,6 +147,7 @@ extern const char lua_ident[];
 LUA_API lua_State *(lua_newstate) (lua_Alloc f, void *ud);
 LUA_API void       (lua_close) (lua_State *L);
 LUA_API lua_State *(lua_newthread) (lua_State *L);
+LUA_API int        (lua_resetthread) (lua_State *L);
 
 LUA_API lua_CFunction (lua_atpanic) (lua_State *L, lua_CFunction panicf);
 
diff --git a/lvm.c b/lvm.c
index fc8722a8..652095dc 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1565,7 +1565,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           if (nparams1)  /* vararg function? */
             delta = ci->u.l.nextraargs + nparams1;
           /* close upvalues from current call */
-          luaF_close(L, base, -1);  /* (no to-be-closed vars. here) */
+          luaF_close(L, base, LUA_OK);
           updatestack(ci);
         }
         if (!ttisfunction(s2v(ra))) {  /* not a function? */
diff --git a/manual/manual.of b/manual/manual.of
index 0e8e3d72..862d032b 100644
--- a/manual/manual.of
+++ b/manual/manual.of
@@ -3927,6 +3927,19 @@ and then pops the top element.
 
 }
 
+@APIEntry{int lua_resetthread (lua_State *L);|
+@apii{0,?,-}
+
+Resets a thread, cleaning its call stack and closing all pending
+to-be-closed variables.
+Returns a status code:
+@Lid{LUA_OK} for no errors in closing methods,
+or an error status otherwise.
+In case of error,
+leave the error object on the stack,
+
+}
+
 @APIEntry{int lua_resume (lua_State *L, lua_State *from, int nargs,
                           int *nresults);|
 @apii{?,?,-}
@@ -3948,11 +3961,8 @@ or returned by the body function.
 @Lid{LUA_OK} if the coroutine finishes its execution
 without errors,
 or an error code in case of errors @seeC{lua_pcall}.
-
 In case of errors,
-the stack is not unwound,
-so you can use the debug API over it.
-The error object is on the top of the stack.
+the error object is on the top of the stack.
 
 To resume a coroutine,
 you remove all results from the last @Lid{lua_yield},
@@ -6285,6 +6295,17 @@ it is not inside a non-yieldable @N{C function}.
 
 }
 
+@LibEntry{coroutine.kill(co)|
+
+Kills coroutine @id{co},
+closing all its pending to-be-closed variables
+and putting the coroutine in a dead state.
+In case of error closing some variable,
+returns @false plus the error object;
+otherwise returns @true.
+
+}
+
 @LibEntry{coroutine.resume (co [, val1, @Cdots])|
 
 Starts or continues the execution of coroutine @id{co}.
@@ -8648,6 +8669,11 @@ has been removed.
 When needed, this metamethod must be explicitly defined.
 }
 
+@item{
+When a coroutine finishes with an error,
+its stack is unwound (to run any pending closing methods).
+}
+
 }
 
 }
diff --git a/testes/coroutine.lua b/testes/coroutine.lua
index 7d42eadd..5674a4dd 100644
--- a/testes/coroutine.lua
+++ b/testes/coroutine.lua
@@ -119,6 +119,51 @@ end
 assert(#a == 25 and a[#a] == 97)
 x, a = nil
 
+
+-- coroutine kill
+do
+  -- ok to kill a dead coroutine
+  local co = coroutine.create(print)
+  assert(coroutine.resume(co, "testing 'coroutine.kill'"))
+  assert(coroutine.status(co) == "dead")
+  assert(coroutine.kill(co))
+
+  -- cannot kill the running coroutine
+  local st, msg = pcall(coroutine.kill, coroutine.running())
+  assert(not st and string.find(msg, "running"))
+
+  local main = coroutine.running()
+
+  -- cannot kill a "normal" coroutine
+  ;(coroutine.wrap(function ()
+    local st, msg = pcall(coroutine.kill, main)
+    assert(not st and string.find(msg, "normal"))
+  end))()
+
+  -- to-be-closed variables in coroutines
+  local X
+  co = coroutine.create(function ()
+    local *toclose x = function (err) assert(err == nil); X = false end
+    X = true
+    coroutine.yield()
+  end)
+  coroutine.resume(co)
+  assert(X)
+  assert(coroutine.kill(co))
+  assert(not X and coroutine.status(co) == "dead")
+
+  -- error killing a coroutine
+  co = coroutine.create(function()
+    local *toclose x = function (err) assert(err == nil); error(111) end
+    coroutine.yield()
+  end)
+  coroutine.resume(co)
+  local st, msg = coroutine.kill(co)
+  assert(not st and coroutine.status(co) == "dead" and msg == 111)
+
+end
+
+
 -- yielding across C boundaries
 
 co = coroutine.wrap(function()
diff --git a/testes/main.lua b/testes/main.lua
index c7bde0d9..b9dcab1c 100644
--- a/testes/main.lua
+++ b/testes/main.lua
@@ -254,15 +254,15 @@ NoRun("error object is a table value", [[lua %s]], prog)
 
 
 -- chunk broken in many lines
-s = [=[ -- 
-function f ( x ) 
+s = [=[ --
+function f ( x )
   local a = [[
 xuxu
 ]]
   local b = "\
 xuxu\n"
   if x == 11 then return 1 + 12 , 2 + 20 end  --[[ test multiple returns ]]
-  return x + 1 
+  return x + 1
   --\\
 end
 return( f( 100 ) )
@@ -272,10 +272,10 @@ s = string.gsub(s, ' ', '\n\n')   -- change all spaces for newlines
 prepfile(s)
 RUN([[lua -e"_PROMPT='' _PROMPT2=''" -i < %s > %s]], prog, out)
 checkprogout("101\n13\t22\n\n")
-  
+
 prepfile[[#comment in 1st line without \n at the end]]
 RUN('lua %s', prog)
-  
+
 prepfile[[#test line number when file starts with comment line
 debug = require"debug"
 print(debug.getinfo(1).currentline)
@@ -306,6 +306,20 @@ NoRun("", "lua %s", prog)   -- no message
 prepfile("os.exit(false, true)")
 NoRun("", "lua %s", prog)   -- no message
 
+
+-- to-be-closed variables in main chunk
+prepfile[[
+  local *toclose x = function (err)
+    assert(err == 120)
+    print("Ok")
+  end
+  local *toclose e1 = function () error(120) end
+  os.exit(true, true)
+]]
+RUN('lua %s > %s', prog, out)
+checkprogout("Ok")
+
+
 -- remove temporary files
 assert(os.remove(prog))
 assert(os.remove(otherprog))