10 files changed, 103 insertions, 66 deletions

diff --git a/lapi.c b/lapi.c
index 41e6b86d..a9cf2fdb 100644
--- a/lapi.c
+++ b/lapi.c
@@ -192,9 +192,8 @@ LUA_API void lua_settop (lua_State *L, int idx) {
   if (diff < 0 && hastocloseCfunc(ci->nresults))
     luaF_close(L, L->top + diff, CLOSEKTOP, 0);
 #endif
+  api_check(L, L->tbclist < L->top + diff, "cannot pop an unclosed slot");
   L->top += diff;
-  api_check(L, L->openupval == NULL || uplevel(L->openupval) < L->top,
-               "cannot pop an unclosed slot");
   lua_unlock(L);
 }
 
@@ -203,8 +202,7 @@ LUA_API void lua_closeslot (lua_State *L, int idx) {
   StkId level;
   lua_lock(L);
   level = index2stack(L, idx);
-  api_check(L, hastocloseCfunc(L->ci->nresults) && L->openupval != NULL &&
-               uplevel(L->openupval) == level,
+  api_check(L, hastocloseCfunc(L->ci->nresults) && L->tbclist == level,
      "no variable to close at given level");
   luaF_close(L, level, CLOSEKTOP, 0);
   level = index2stack(L, idx);  /* stack may be moved */
@@ -1266,8 +1264,7 @@ LUA_API void lua_toclose (lua_State *L, int idx) {
   lua_lock(L);
   o = index2stack(L, idx);
   nresults = L->ci->nresults;
-  api_check(L, L->openupval == NULL || uplevel(L->openupval) <= o,
-               "marked index below or equal new one");
+  api_check(L, L->tbclist < o, "given index below or equal a marked one");
   luaF_newtbcupval(L, o);  /* create new to-be-closed upvalue */
   if (!hastocloseCfunc(nresults))  /* function not marked yet? */
     L->ci->nresults = codeNresults(nresults);  /* mark it */
diff --git a/ldo.c b/ldo.c
index 65f0a7b9..bc7212c6 100644
--- a/ldo.c
+++ b/ldo.c
@@ -163,7 +163,7 @@ static void correctstack (lua_State *L, StkId oldstack, StkId newstack) {
   if (oldstack == newstack)
     return;  /* stack address did not change */
   L->top = (L->top - oldstack) + newstack;
-  lua_assert(L->ptbc == NULL);
+  L->tbclist = (L->tbclist - oldstack) + newstack;
   for (up = L->openupval; up != NULL; up = up->u.open.next)
     up->v = s2v((uplevel(up) - oldstack) + newstack);
   for (ci = L->ci; ci != NULL; ci = ci->previous) {
diff --git a/lfunc.c b/lfunc.c
index 105590fc..b4c04bd0 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -120,11 +120,11 @@ static void callclosemethod (lua_State *L, TValue *obj, TValue *err, int yy) {
 
 
 /*
-** Check whether 'obj' has a close metamethod and raise an error
-** if not.
+** Check whether object at given level has a close metamethod and raise
+** an error if not.
 */
-static void checkclosemth (lua_State *L, StkId level, const TValue *obj) {
-  const TValue *tm = luaT_gettmbyobj(L, obj, TM_CLOSE);
+static void checkclosemth (lua_State *L, StkId level) {
+  const TValue *tm = luaT_gettmbyobj(L, s2v(level), TM_CLOSE);
   if (ttisnil(tm)) {  /* no metamethod? */
     int idx = cast_int(level - L->ci->func);  /* variable index */
     const char *vname = luaG_findlocal(L, L->ci, idx, NULL);
@@ -155,20 +155,21 @@ static void prepcallclosemth (lua_State *L, StkId level, int status, int yy) {
 
 
 /*
-** Create a to-be-closed upvalue. If there is a memory allocation error,
-** 'ptbc' keeps the object so it can be closed as soon as possible.
-** (Since memory errors have no handler, that will happen before any
-** stack reallocation.)
+** Insert a variable in the list of to-be-closed variables.
 */
 void luaF_newtbcupval (lua_State *L, StkId level) {
-  TValue *obj = s2v(level);
-  lua_assert(L->openupval == NULL || uplevel(L->openupval) < level);
-  if (!l_isfalse(obj)) {  /* false doesn't need to be closed */
-    checkclosemth(L, level, obj);
-    L->ptbc = level;  /* in case of allocation error */
-    newupval(L, 1, level, &L->openupval);
-    L->ptbc = NULL;  /* no errors */
+  lua_assert(level > L->tbclist);
+  if (l_isfalse(s2v(level)))
+    return;  /* false doesn't need to be closed */
+  checkclosemth(L, level);  /* value must have a close method */
+  while (level - L->tbclist > USHRT_MAX) {  /* is delta too large? */
+    L->tbclist += USHRT_MAX;  /* create a dummy node at maximum delta */
+    L->tbclist->tbclist.delta = USHRT_MAX;
+    L->tbclist->tbclist.isdummy = 1;
   }
+  level->tbclist.delta = level - L->tbclist;
+  level->tbclist.isdummy = 0;
+  L->tbclist = level;
 }
 
 
@@ -181,23 +182,11 @@ void luaF_unlinkupval (UpVal *uv) {
 
 
 /*
-** Close all upvalues up to the given stack level. A 'status' equal
-** to NOCLOSINGMETH closes upvalues without running any __close
-** metamethods. If there is a pending to-be-closed value, close
-** it before anything else.
+** Close all upvalues up to the given stack level.
 */
-void luaF_close (lua_State *L, StkId level, int status, int yy) {
+void luaF_closeupval (lua_State *L, StkId level) {
   UpVal *uv;
   StkId upl;  /* stack index pointed by 'uv' */
-  if (unlikely(status == LUA_ERRMEM && L->ptbc != NULL)) {
-    ptrdiff_t levelrel = savestack(L, level);
-    upl = L->ptbc;
-    L->ptbc = NULL;  /* remove from "list" before closing */
-    prepcallclosemth(L, upl, status, yy);
-    level = restorestack(L, levelrel);
-  }
-  else
-    lua_assert(L->ptbc == NULL);  /* must be empty for other status */
   while ((uv = L->openupval) != NULL && (upl = uplevel(uv)) >= level) {
     TValue *slot = &uv->u.value;  /* new position for value */
     lua_assert(uplevel(uv) < L->top);
@@ -208,9 +197,22 @@ void luaF_close (lua_State *L, StkId level, int status, int yy) {
       nw2black(uv);  /* closed upvalues cannot be gray */
       luaC_barrier(L, uv, slot);
     }
-    if (uv->tbc && status != NOCLOSINGMETH) {
-      ptrdiff_t levelrel = savestack(L, level);
-      prepcallclosemth(L, upl, status, yy);  /* may change the stack */
+  }
+}
+
+
+/*
+** Close all upvalues and to-be-closed variables up to the given stack
+** level.
+*/
+void luaF_close (lua_State *L, StkId level, int status, int yy) {
+  ptrdiff_t levelrel = savestack(L, level);
+  luaF_closeupval(L, level);  /* first, close the upvalues */
+  while (L->tbclist >= level) {  /* traverse tbc's down to that level */
+    StkId tbc = L->tbclist;  /* get variable index */
+    L->tbclist -= tbc->tbclist.delta;  /* remove it from list */
+    if (!tbc->tbclist.isdummy) {  /* not a dummy entry? */
+      prepcallclosemth(L, tbc, status, yy);  /* close variable */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lfunc.h b/lfunc.h
index 2e6df535..dc1cebcc 100644
--- a/lfunc.h
+++ b/lfunc.h
@@ -42,15 +42,9 @@
 #define MAXMISS         10
 
 
-/*
-** Special "status" for 'luaF_close'
-*/
-
-/* close upvalues without running their closing methods */
-#define NOCLOSINGMETH   (-1)
 
 /* special status to close upvalues preserving the top of the stack */
-#define CLOSEKTOP       (-2)
+#define CLOSEKTOP       (-1)
 
 
 LUAI_FUNC Proto *luaF_newproto (lua_State *L);
@@ -59,6 +53,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_closeupval (lua_State *L, StkId level);
 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);
diff --git a/lobject.h b/lobject.h
index 470b17d5..1a7a7372 100644
--- a/lobject.h
+++ b/lobject.h
@@ -136,10 +136,18 @@ typedef struct TValue {
 
 
 /*
-** Entries in the Lua stack
+** Entries in a Lua stack. Field 'tbclist' forms a list of all
+** to-be-closed variables active in this stack. Dummy entries are
+** used when the distance between two tbc variables does not fit
+** in an unsigned short.
 */
 typedef union StackValue {
   TValue val;
+  struct {
+    TValuefields;
+    lu_byte isdummy;
+    unsigned short delta;
+  } tbclist;
 } StackValue;
 
 
diff --git a/lstate.c b/lstate.c
index 52336f44..38078521 100644
--- a/lstate.c
+++ b/lstate.c
@@ -181,6 +181,7 @@ static void stack_init (lua_State *L1, lua_State *L) {
   int i; CallInfo *ci;
   /* initialize stack array */
   L1->stack = luaM_newvector(L, BASIC_STACK_SIZE + EXTRA_STACK, StackValue);
+  L1->tbclist = L1->stack;
   for (i = 0; i < BASIC_STACK_SIZE + EXTRA_STACK; i++)
     setnilvalue(s2v(L1->stack + i));  /* erase new stack */
   L1->top = L1->stack;
@@ -262,16 +263,18 @@ static void preinit_thread (lua_State *L, global_State *g) {
   L->status = LUA_OK;
   L->errfunc = 0;
   L->oldpc = 0;
-  L->ptbc = NULL;
 }
 
 
 static void close_state (lua_State *L) {
   global_State *g = G(L);
-  luaD_closeprotected(L, 0, LUA_OK);  /* close all upvalues */
-  luaC_freeallobjects(L);  /* collect all objects */
-  if (completestate(g))  /* closing a fully built state? */
+  if (!completestate(g))  /* closing a partially built state? */
+    luaC_freeallobjects(L);  /* jucst collect its objects */
+  else {  /* closing a fully built state */
+    luaD_closeprotected(L, 1, LUA_OK);  /* close all upvalues */
+    luaC_freeallobjects(L);  /* collect all objects */
     luai_userstateclose(L);
+  }
   luaM_freearray(L, G(L)->strt.hash, G(L)->strt.size);
   freestack(L);
   lua_assert(gettotalbytes(g) == sizeof(LG));
@@ -312,7 +315,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, 0);  /* close all upvalues */
+  luaF_closeupval(L1, L1->stack);  /* close all upvalues */
   lua_assert(L1->openupval == NULL);
   luai_userstatefree(L, L1);
   freestack(L1);
@@ -327,7 +330,7 @@ int luaE_resetthread (lua_State *L, int status) {
   ci->callstatus = CIST_C;
   if (status == LUA_YIELD)
     status = LUA_OK;
-  status = luaD_closeprotected(L, 0, status);
+  status = luaD_closeprotected(L, 1, status);
   if (status != LUA_OK)  /* errors? */
     luaD_seterrorobj(L, status, L->stack + 1);
   else
diff --git a/lstate.h b/lstate.h
index 5ef55355..b6ade7c7 100644
--- a/lstate.h
+++ b/lstate.h
@@ -307,6 +307,7 @@ struct lua_State {
   StkId stack_last;  /* end of stack (last element + 1) */
   StkId stack;  /* stack base */
   UpVal *openupval;  /* list of open upvalues in this stack */
+  StkId tbclist;  /* list of to-be-closed variables */
   GCObject *gclist;
   struct lua_State *twups;  /* list of threads with open upvalues */
   struct lua_longjmp *errorJmp;  /* current error recover point */
@@ -318,7 +319,6 @@ struct lua_State {
   int basehookcount;
   int hookcount;
   volatile l_signalT hookmask;
-  StkId ptbc;  /* pending to-be-closed variable */
 };
 
 
diff --git a/ltests.c b/ltests.c
index 9c13338a..da95d027 100644
--- a/ltests.c
+++ b/ltests.c
@@ -446,6 +446,7 @@ static void checkstack (global_State *g, lua_State *L1) {
   for (uv = L1->openupval; uv != NULL; uv = uv->u.open.next)
     assert(upisopen(uv));  /* must be open */
   assert(L1->top <= L1->stack_last);
+  assert(L1->tbclist <= L1->top);
   for (ci = L1->ci; ci != NULL; ci = ci->previous) {
     assert(ci->top <= L1->stack_last);
     assert(lua_checkpc(ci));
diff --git a/lvm.c b/lvm.c
index e9b1dcdd..1252ecbf 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1635,10 +1635,8 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           b = cast_int(L->top - ra);
         savepc(ci);  /* several calls here can raise errors */
         if (TESTARG_k(i)) {
-          /* 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, 0);
+          luaF_closeupval(L, base);  /* close upvalues from current call */
+          lua_assert(L->tbclist < base);  /* no pending tbc variables */
           lua_assert(base == ci->func + 1);
         }
         while (!ttisfunction(s2v(ra))) {  /* not a function? */
diff --git a/testes/locals.lua b/testes/locals.lua
index a25b2b9f..446ec13a 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -529,6 +529,40 @@ local function checktable (t1, t2)
 end
 
 
+do    -- test for tbc variable high in the stack
+
+   -- function to force a stack overflow
+  local function overflow (n)
+    overflow(n + 1)
+  end
+
+  -- error handler will create tbc variable handling a stack overflow,
+  -- high in the stack
+  local function errorh (m)
+    assert(string.find(m, "stack overflow"))
+    local x <close> = func2close(function (o) o[1] = 10 end)
+    return x
+  end
+
+  local flag
+  local st, obj
+  -- run test in a coroutine so as not to swell the main stack
+  local co = coroutine.wrap(function ()
+    -- tbc variable down the stack
+    local y <close> = func2close(function (obj, msg)
+      assert(msg == nil)
+      obj[1] = 100
+      flag = obj
+    end)
+    collectgarbage("stop")
+    st, obj = xpcall(overflow, errorh, 0)
+    collectgarbage("restart")
+  end)
+  co()
+  assert(not st and obj[1] == 10 and flag[1] == 100)
+end
+
+
 if rawget(_G, "T") then
 
   -- memory error inside closing function
@@ -563,13 +597,13 @@ if rawget(_G, "T") then
 
   local function test ()
     local x <close> = enter(0)   -- set a memory limit
-    -- creation of previous upvalue will raise a memory error
-    assert(false)    -- should not run
+    local y = {}    -- raise a memory error
   end
 
   local _, msg = pcall(test)
   assert(msg == "not enough memory" and closemsg == "not enough memory")
 
+
   -- repeat test with extra closing upvalues
   local function test ()
     local xxx <close> = func2close(function (self, msg)
@@ -580,8 +614,7 @@ if rawget(_G, "T") then
       assert(msg == "not enough memory");
     end)
     local x <close> = enter(0)   -- set a memory limit
-    -- creation of previous upvalue will raise a memory error
-    os.exit(false)    -- should not run
+    local y = {}   -- raise a memory error
   end
 
   local _, msg = pcall(test)
@@ -607,7 +640,7 @@ if rawget(_G, "T") then
     -- concat this table needs two buffer resizes (one for each 's')
     local a = {s, s}
 
-    collectgarbage()
+    collectgarbage(); collectgarbage()
 
     m = T.totalmem()
     collectgarbage("stop")
@@ -630,7 +663,7 @@ if rawget(_G, "T") then
     -- second buffer was released by 'toclose'
     assert(T.totalmem() - m <= extra)
 
-    -- userdata, upvalue, buffer, buffer, final string
+    -- userdata, buffer, buffer, final string
     T.totalmem(m + 4*lim + extra)
     assert(#table.concat(a) == 2*lim)
 
@@ -753,8 +786,8 @@ do
   checktable({co()}, {true, 10, 20, 30})
   checktable(trace, {"nowX", "z1", "z2", "nowY", "y1", "y2", "x1", "x2"})
 
-end 
- 
+end
+
 
 do
   -- yielding inside closing metamethods after an error