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