4 files changed, 98 insertions, 42 deletions

diff --git a/lfunc.c b/lfunc.c
index 4f9362f3..15874f88 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -98,27 +98,29 @@ UpVal *luaF_findupval (lua_State *L, StkId level) {
 
 
 static void callclose (lua_State *L, void *ud) {
-  luaD_callnoyield(L, cast(StkId, ud), 0);
+  UNUSED(ud);
+  luaD_callnoyield(L, L->top - 2, 0);
 }
 
 
 /*
-** Prepare closing method with its argument for object at
-** index 'func' in the stack. Assume there is an error message
-** (or nil) just below the object.
+** Prepare closing method plus its argument for object 'obj' with
+** error message 'err'. (This function assumes EXTRA_STACK.)
 */
-static int prepclosingmethod (lua_State *L, StkId func) {
-  if (ttisfunction(s2v(func))) {  /* object to-be-closed is a function? */
-    setobjs2s(L, func + 1, func - 1);  /* push error msg. as argument */
+static int prepclosingmethod (lua_State *L, TValue *obj, TValue *err) {
+  StkId top = L->top;
+  if (ttisfunction(obj)) {  /* object to-be-closed is a function? */
+    setobj2s(L, top, obj);  /* push function */
+    setobj2s(L, top + 1, err);  /* push error msg. as argument */
   }
   else {  /* try '__close' metamethod */
-    const TValue *tm = luaT_gettmbyobj(L, s2v(func), TM_CLOSE);
+    const TValue *tm = luaT_gettmbyobj(L, obj, TM_CLOSE);
     if (ttisnil(tm))  /* no metamethod? */
       return 0;  /* nothing to call */
-    setobjs2s(L, func + 1, func);  /* 'self' is the argument */
-    setobj2s(L, func, tm);  /* will call metamethod */
+    setobj2s(L, top, tm);  /* will call metamethod... */
+    setobj2s(L, top + 1, obj);  /* with 'self' as the argument */
   }
-  L->top = func + 2;  /* add function and argument */
+  L->top = top + 2;  /* add function and argument */
   return 1;
 }
 
@@ -129,22 +131,24 @@ static int prepclosingmethod (lua_State *L, StkId func) {
 ** 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.
+** 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) {
-  StkId func = level + 1;  /* save slot for old error message */
-  if (unlikely(status != LUA_OK))  /* was there an error? */
-    luaD_seterrorobj(L, status, level);  /* save error message */
-  else
-    setnilvalue(s2v(level));  /* no error message */
-  setobj2s(L, func, uv);  /* put object on top of error message */
-  if (!prepclosingmethod(L, func))
-    return status;  /* nothing to call */
-  if (likely(status == LUA_OK))  /* not in "error mode"? */
-    callclose(L, func);  /* call closing method */
-  else {  /* already inside error handler; cannot raise another error */
-    int newstatus = luaD_pcall(L, callclose, func, savestack(L, level), 0);
-    if (newstatus != LUA_OK)  /* error when closing? */
-      status = newstatus;  /* this will be the new error */
+  if (likely(status == LUA_OK)) {
+    if (prepclosingmethod(L, uv, &G(L)->nilvalue))  /* something to call? */
+      callclose(L, NULL);  /* call closing method */
+  }
+  else {  /* there was an error */
+    /* save error message and set stack top to 'level + 1' */
+    luaD_seterrorobj(L, status, level);
+    if (prepclosingmethod(L, uv, s2v(level))) {  /* something to call? */
+      int newstatus = luaD_pcall(L, callclose, NULL, savestack(L, level), 0);
+      if (newstatus != LUA_OK)  /* another error when closing? */
+        status = newstatus;  /* this will be the new error */
+    }
   }
   return status;
 }
@@ -169,12 +173,10 @@ static void trynewtbcupval (lua_State *L, void *ud) {
 void luaF_newtbcupval (lua_State *L, StkId level) {
   int status = luaD_rawrunprotected(L, trynewtbcupval, level);
   if (unlikely(status != LUA_OK)) {  /* memory error creating upvalue? */
-    StkId func = level + 1;
     lua_assert(status == LUA_ERRMEM);
-    setobjs2s(L, func, level);  /* open space for error message */
-    luaD_seterrorobj(L, status, level);  /* save error message */
-    if (prepclosingmethod(L, func))
-      callclose(L, func);  /* call closing method */
+    luaD_seterrorobj(L, LUA_ERRMEM, level + 1);  /* save error message */
+    if (prepclosingmethod(L, s2v(level), s2v(level + 1)))
+      callclose(L, NULL);  /* call closing method */
     luaD_throw(L, LUA_ERRMEM);  /* throw memory error */
   }
 }
@@ -201,7 +203,7 @@ int luaF_close (lua_State *L, StkId level, int status) {
     luaC_barrier(L, uv, slot);
     if (status >= 0 && uv->tt == LUA_TUPVALTBC) {  /* must be closed? */
       ptrdiff_t levelrel = savestack(L, level);
-      status = closeupval(L, uv->v, upl, status);  /* may reallocate the stack */
+      status = closeupval(L, uv->v, upl, status);  /* may realloc. the stack */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lparser.c b/lparser.c
index 6b14b800..c0c40eae 100644
--- a/lparser.c
+++ b/lparser.c
@@ -561,7 +561,7 @@ static void close_func (LexState *ls) {
   lua_State *L = ls->L;
   FuncState *fs = ls->fs;
   Proto *f = fs->f;
-  luaK_ret(fs, 0, 0);  /* final return */
+  luaK_ret(fs, fs->nactvar, 0);  /* final return */
   leaveblock(fs);
   lua_assert(fs->bl == NULL);
   luaK_finish(fs);
@@ -1602,9 +1602,10 @@ static void retstat (LexState *ls) {
   /* stat -> RETURN [explist] [';'] */
   FuncState *fs = ls->fs;
   expdesc e;
-  int first, nret;  /* registers with returned values */
+  int nret;  /* number of values being returned */
+  int first = fs->nactvar;  /* first slot to be returned */
   if (block_follow(ls, 1) || ls->t.token == ';')
-    first = nret = 0;  /* return no values */
+    nret = 0;  /* return no values */
   else {
     nret = explist(ls, &e);  /* optional return values */
     if (hasmultret(e.k)) {
@@ -1613,15 +1614,13 @@ static void retstat (LexState *ls) {
         SET_OPCODE(getinstruction(fs,&e), OP_TAILCALL);
         lua_assert(GETARG_A(getinstruction(fs,&e)) == fs->nactvar);
       }
-      first = fs->nactvar;
       nret = LUA_MULTRET;  /* return all values */
     }
     else {
       if (nret == 1)  /* only one single value? */
-        first = luaK_exp2anyreg(fs, &e);
-      else {
-        luaK_exp2nextreg(fs, &e);  /* values must go to the stack */
-        first = fs->nactvar;  /* return all active values */
+        first = luaK_exp2anyreg(fs, &e);  /* can use original slot */
+      else {  /* values must go to the top of the stack */
+        luaK_exp2nextreg(fs, &e);
         lua_assert(nret == fs->freereg - first);
       }
     }
diff --git a/lvm.c b/lvm.c
index 0d82756b..2a1ee175 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1452,7 +1452,8 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_CLOSE) {
-        luaF_close(L, ra, LUA_OK);
+        L->top = ra + 1;  /* everything is free after this slot */
+        ProtectNT(luaF_close(L, ra, LUA_OK));
         vmbreak;
       }
       vmcase(OP_TBC) {
@@ -1619,13 +1620,14 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           n = cast_int(L->top - ra);  /* get what is available */
         else
           L->top = ra + n;  /* set call for 'luaD_poscall' */
+        savepc(ci);
         if (TESTARG_k(i)) {
           int nparams1 = GETARG_C(i);
           if (nparams1)  /* vararg function? */
             ci->func -= ci->u.l.nextraargs + nparams1;
           luaF_close(L, base, LUA_OK);  /* there may be open upvalues */
         }
-        halfProtect(luaD_poscall(L, ci, n));
+        luaD_poscall(L, ci, n);
         return;
       }
       vmcase(OP_RETURN0) {
diff --git a/testes/locals.lua b/testes/locals.lua
index f21fa2ec..65b145db 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -175,6 +175,9 @@ assert(x==20)
 
 print"testing to-be-closed variables"
 
+local function stack(n) n = ((n == 0) or stack(n - 1)) end
+
+
 do
   local a = {}
   do
@@ -187,6 +190,57 @@ do
   assert(a[1] == "in" and a[2] == "y" and a[3] == "x" and a[4] == "out")
 end
 
+do
+  local X = false
+
+  local function closescope () stack(10); X = true end
+
+  -- closing functions do not corrupt returning values
+  local function foo (x)
+    local scoped _ = closescope
+    return x, X, 23
+  end
+
+  local a, b, c = foo(1.5)
+  assert(a == 1.5 and b == false and c == 23 and X == true)
+
+  X = false
+  foo = function (x)
+    local scoped _ = closescope
+    local y = 15
+    return y
+  end
+
+  assert(foo() == 15 and X == true)
+
+  X = false
+  foo = function ()
+    local scoped x = closescope
+    return x
+  end
+
+  assert(foo() == closescope and X == true)
+
+end
+
+
+do
+  -- to-be-closed variables must be closed in tail calls
+  local X, Y
+  local function foo ()
+    local scoped _ = function () Y = 10 end
+    assert(X == 20 and Y == nil)
+    return 1,2,3
+  end
+
+  local function bar ()
+    local scoped _ = function () X = 20 end
+    return foo()
+  end
+
+  local a, b, c, d = bar()
+  assert(a == 1 and b == 2 and c == 3 and X == 20 and Y == 10 and d == nil)
+end
 
 do   -- errors in __close
   local log = {}
@@ -211,7 +265,6 @@ do   -- errors in __close
 end
 
 if rawget(_G, "T") then
-  local function stack(n) n = (n == 0) or stack(n - 1); end;
   -- memory error inside closing function
   local function foo ()
     local scoped y = function () T.alloccount() end