State in generic 'for' acts as a to-be-closed variable

The implicit variable 'state' in a generic 'for' is marked as a
to-be-closed variable, so that the state will be closed as soon
as the loop ends, no matter how.

Taking advantage of this new facility, the call 'io.lines(filename)'
now returns the open file as a second result. Therefore,
an iteraction like 'for l in io.lines(name)...' will close the
file even when the loop ends with a break or an error.

5 files changed, 111 insertions, 24 deletions

diff --git a/liolib.c b/liolib.c
index 3dc509bd..5881b029 100644
--- a/liolib.c
+++ b/liolib.c
@@ -354,12 +354,22 @@ static int io_readline (lua_State *L);
 */
 #define MAXARGLINE	250
 
+/*
+** Auxiliar function to create the iteration function for 'lines'.
+** The iteration function is a closure over 'io_readline', with
+** the following upvalues:
+** 1) The file being read (first value in the stack)
+** 2) the number of arguments to read
+** 3) a boolean, true iff file has to be closed when finished ('toclose')
+** *) a variable number of format arguments (rest of the stack)
+*/
 static void aux_lines (lua_State *L, int toclose) {
   int n = lua_gettop(L) - 1;  /* number of arguments to read */
   luaL_argcheck(L, n <= MAXARGLINE, MAXARGLINE + 2, "too many arguments");
+  lua_pushvalue(L, 1);  /* file */
   lua_pushinteger(L, n);  /* number of arguments to read */
   lua_pushboolean(L, toclose);  /* close/not close file when finished */
-  lua_rotate(L, 2, 2);  /* move 'n' and 'toclose' to their positions */
+  lua_rotate(L, 2, 3);  /* move the three values to their positions */
   lua_pushcclosure(L, io_readline, 3 + n);
 }
 
@@ -371,6 +381,11 @@ static int f_lines (lua_State *L) {
 }
 
 
+/*
+** Return an iteration function for 'io.lines'. If file has to be
+** closed, also returns the file itself as a second result (to be
+** closed as the state at the exit of a generic for).
+*/
 static int io_lines (lua_State *L) {
   int toclose;
   if (lua_isnone(L, 1)) lua_pushnil(L);  /* at least one argument */
@@ -386,8 +401,13 @@ static int io_lines (lua_State *L) {
     lua_replace(L, 1);  /* put file at index 1 */
     toclose = 1;  /* close it after iteration */
   }
-  aux_lines(L, toclose);
-  return 1;
+  aux_lines(L, toclose);  /* push iteration function */
+  if (toclose) {
+    lua_pushvalue(L, 1);  /* file will be second result */
+    return 2;
+  }
+  else
+    return 1;
 }
 
 
@@ -453,7 +473,7 @@ static int readdigits (RN *rn, int hex) {
 /*
 ** Read a number: first reads a valid prefix of a numeral into a buffer.
 ** Then it calls 'lua_stringtonumber' to check whether the format is
-** correct and to convert it to a Lua number
+** correct and to convert it to a Lua number.
 */
 static int read_number (lua_State *L, FILE *f) {
   RN rn;
@@ -604,6 +624,9 @@ static int f_read (lua_State *L) {
 }
 
 
+/*
+** Iteration function for 'lines'.
+*/
 static int io_readline (lua_State *L) {
   LStream *p = (LStream *)lua_touserdata(L, lua_upvalueindex(1));
   int i;
@@ -624,8 +647,8 @@ static int io_readline (lua_State *L) {
       return luaL_error(L, "%s", lua_tostring(L, -n + 1));
     }
     if (lua_toboolean(L, lua_upvalueindex(3))) {  /* generator created file? */
-      lua_settop(L, 0);
-      lua_pushvalue(L, lua_upvalueindex(1));
+      lua_settop(L, 0);  /* clear stack */
+      lua_pushvalue(L, lua_upvalueindex(1));  /* push file at index 1 */
       aux_close(L);  /* close it */
     }
     return 0;
diff --git a/lparser.c b/lparser.c
index 9419f880..a5b84aa1 100644
--- a/lparser.c
+++ b/lparser.c
@@ -1413,6 +1413,7 @@ static void forlist (LexState *ls, TString *indexname) {
   /* create control variables */
   new_localvarliteral(ls, "(for generator)");
   new_localvarliteral(ls, "(for state)");
+  markupval(fs, fs->nactvar);  /* state may create an upvalue */
   new_localvarliteral(ls, "(for control)");
   /* create declared variables */
   new_localvar(ls, indexname);
diff --git a/lvm.c b/lvm.c
index aad965d6..1535700f 100644
--- a/lvm.c
+++ b/lvm.c
@@ -866,7 +866,8 @@ void luaV_finishOp (lua_State *L) {
 #define ProtectNT(exp)  (savepc(L), (exp), updatetrap(ci))
 
 /*
-** Protect code that will finish the loop (returns).
+** Protect code that will finish the loop (returns) or can only raise
+** errors.
 */
 #define halfProtect(exp)  (savepc(L), (exp))
 
@@ -1457,7 +1458,8 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_TBC) {
-        luaF_newtbcupval(L, ra);  /* create new to-be-closed upvalue */
+        /* create new to-be-closed upvalue */
+        halfProtect(luaF_newtbcupval(L, ra));
         vmbreak;
       }
       vmcase(OP_JMP) {
@@ -1745,21 +1747,34 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_TFORPREP) {
+        /* is 'state' a function or has a '__close' metamethod? */
+        if (ttisfunction(s2v(ra + 1)) ||
+            !ttisnil(luaT_gettmbyobj(L, s2v(ra + 1), TM_CLOSE))) {
+          /* create to-be-closed upvalue for it */
+          halfProtect(luaF_newtbcupval(L, ra + 1));
+        }
         pc += GETARG_Bx(i);
-        vmbreak;
+        i = *(pc++);  /* go to next instruction */
+        lua_assert(GET_OPCODE(i) == OP_TFORCALL && ra == RA(i));
+        goto l_tforcall;
       }
       vmcase(OP_TFORCALL) {
-        StkId cb = ra + 3;  /* call base */
-        setobjs2s(L, cb+2, ra+2);
-        setobjs2s(L, cb+1, ra+1);
-        setobjs2s(L, cb, ra);
-        L->top = cb + 3;  /* func. + 2 args (state and index) */
-        Protect(luaD_call(L, cb, GETARG_C(i)));
-        if (trap)  /* keep 'base' correct for next instruction */
-          updatebase(ci);
+       l_tforcall:
+        /* 'ra' has the iterator function, 'ra + 1' has the state,
+           and 'ra + 2' has the control variable. The call will use
+           the stack after these values (starting at 'ra + 3')
+        */
+        /* push function, state, and control variable */
+        memcpy(ra + 3, ra, 3 * sizeof(*ra));
+        L->top = ra + 6;
+        Protect(luaD_call(L, ra + 3, GETARG_C(i)));  /* do the call */
+        if (trap) {  /* stack may have changed? */
+          updatebase(ci);  /* keep 'base' correct */
+          ra = RA(i);  /* keep 'ra' correct for next instruction */
+        }
         i = *(pc++);  /* go to next instruction */
-        ra = RA(i);  /* get its 'ra' */
-        lua_assert(GET_OPCODE(i) == OP_TFORLOOP);
+        ra += 2;  /* adjust for next instruction */
+        lua_assert(GET_OPCODE(i) == OP_TFORLOOP && ra == RA(i));
         goto l_tforloop;
       }
       vmcase(OP_TFORLOOP) {
diff --git a/testes/files.lua b/testes/files.lua
index 9aae5913..a11c5e61 100644
--- a/testes/files.lua
+++ b/testes/files.lua
@@ -462,13 +462,13 @@ X
 -                                   y;
 ]]:close()
 _G.X = 1
-assert(not load(io.lines(file)))
+assert(not load((io.lines(file))))
 collectgarbage()   -- to close file in previous iteration
-load(io.lines(file, "L"))()
+load((io.lines(file, "L")))()
 assert(_G.X == 2)
-load(io.lines(file, 1))()
+load((io.lines(file, 1)))()
 assert(_G.X == 4)
-load(io.lines(file, 3))()
+load((io.lines(file, 3)))()
 assert(_G.X == 8)
 
 print('+')
diff --git a/testes/locals.lua b/testes/locals.lua
index 65b145db..1e0f525b 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -108,7 +108,7 @@ print'+'
 if rawget(_G, "T") then
   -- testing clearing of dead elements from tables
   collectgarbage("stop")   -- stop GC
-  local a = {[{}] = 4, [3] = 0, alo = 1, 
+  local a = {[{}] = 4, [3] = 0, alo = 1,
              a1234567890123456789012345678901234567890 = 10}
 
   local t = T.querytab(a)
@@ -360,6 +360,54 @@ end)
 co()                 -- start coroutine
 assert(co == nil)    -- eventually it will be collected
 
+
+-- to-be-closed variables in generic for loops
+do
+  local numopen = 0
+  local function open (x)
+    numopen = numopen + 1
+    return
+      function ()   -- iteraction function
+        x = x - 1
+        if x > 0 then return x end
+      end,
+      function ()   -- closing function
+        numopen = numopen - 1
+      end
+  end
+
+  local s = 0
+  for i in open(10) do
+     s = s + i
+  end
+  assert(s == 45 and numopen == 0)
+
+  local s = 0
+  for i in open(10) do
+     if i < 5 then break end
+     s = s + i
+  end
+  assert(s == 35 and numopen == 0)
+
+  -- repeat test with '__open' metamethod instead of a function
+  local function open (x)
+    numopen = numopen + 1
+    return
+      function (t)   -- iteraction function
+        t[1] = t[1] - 1
+        if t[1] > 0 then return t[1] end
+      end,
+      setmetatable({x}, {__close = function () numopen = numopen - 1 end})
+  end
+
+  local s = 0
+  for i in open(10) do
+     if (i < 5) then break end
+     s = s + i
+  end
+  assert(s == 35 and numopen == 0)
+end
+
 print('OK')
 
 return 5,f