5 files changed, 118 insertions, 35 deletions

diff --git a/lcode.c b/lcode.c
index b5ef968e..138ee78a 100644
--- a/lcode.c
+++ b/lcode.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lcode.c,v 1.28 2000/04/19 13:41:37 roberto Exp roberto $
+** $Id: lcode.c,v 1.29 2000/05/08 19:32:53 roberto Exp roberto $
 ** Code generator for Lua
 ** See Copyright Notice in lua.h
 */
@@ -475,6 +475,18 @@ int luaK_code2 (FuncState *fs, OpCode o, int arg1, int arg2) {
       mode = iS;
       break;
 
+    case OP_LFORPREP:
+      delta = 3;
+      arg1 = NO_JUMP;
+      mode = iS;
+      break;
+
+    case OP_LFORLOOP:
+      delta = -4;
+      arg1 = NO_JUMP;
+      mode = iS;
+      break;
+
     case OP_END:
     case OP_PUSHNILJMP:
     case OP_NOT:
diff --git a/lopcodes.h b/lopcodes.h
index 4dcec5aa..283dcdb6 100644
--- a/lopcodes.h
+++ b/lopcodes.h
@@ -1,5 +1,5 @@
 /*
-** $Id: lopcodes.h,v 1.59 2000/04/14 17:45:25 roberto Exp roberto $
+** $Id: lopcodes.h,v 1.60 2000/04/27 17:39:15 roberto Exp roberto $
 ** Opcodes for Lua virtual machine
 ** See Copyright Notice in lua.h
 */
@@ -143,6 +143,9 @@ OP_PUSHNILJMP,/* -	-		nil		PC++;		*/
 OP_FORPREP,/*   J                                                       */
 OP_FORLOOP,/*   J                                                       */
 
+OP_LFORPREP,/*  J                                                       */
+OP_LFORLOOP,/*  J                                                       */
+
 OP_CLOSURE,/*   A B     v_b-v_1         closure(KPROTO[a], v_1-v_b)     */
 
 OP_SETLINE/*    U       -               -               LINE=u          */
diff --git a/lparser.c b/lparser.c
index 9691f4d2..3fae5419 100644
--- a/lparser.c
+++ b/lparser.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lparser.c,v 1.85 2000/05/10 16:33:20 roberto Exp roberto $
+** $Id: lparser.c,v 1.86 2000/05/12 18:12:04 roberto Exp roberto $
 ** LL(1) Parser and code generator for Lua
 ** See Copyright Notice in lua.h
 */
@@ -205,7 +205,8 @@ static void store_localvar (LexState *ls, TString *name, int n) {
 }
 
 
-static void adjustlocalvars (LexState *ls, int nvars, int line) {
+static void adjustlocalvars (LexState *ls, int nvars) {
+  int line = ls->fs->lastsetline;
   FuncState *fs = ls->fs;
   int i;
   fs->nlocalvar += nvars;
@@ -214,7 +215,8 @@ static void adjustlocalvars (LexState *ls, int nvars, int line) {
 }
 
 
-static void removelocalvars (LexState *ls, int nvars, int line) {
+static void removelocalvars (LexState *ls, int nvars) {
+  int line = ls->fs->lastsetline;
   ls->fs->nlocalvar -= nvars;
   while (nvars--)
     luaI_unregisterlocalvar(ls, line);
@@ -223,7 +225,7 @@ static void removelocalvars (LexState *ls, int nvars, int line) {
 
 static void add_localvar (LexState *ls, const char *name) {
   store_localvar(ls, luaS_newfixed(ls->L, name), 0);
-  adjustlocalvars(ls, 1, 0);
+  adjustlocalvars(ls, 1);
 }
 
 
@@ -303,7 +305,7 @@ static void adjust_mult_assign (LexState *ls, int nvars, int nexps) {
 
 static void code_args (LexState *ls, int nparams, int dots) {
   FuncState *fs = ls->fs;
-  adjustlocalvars(ls, nparams, 0);
+  adjustlocalvars(ls, nparams);
   checklimit(ls, fs->nlocalvar, MAXPARAMS, "parameters");
   nparams = fs->nlocalvar;  /* `self' could be there already */
   fs->f->numparams = nparams;
@@ -838,7 +840,7 @@ static void block (LexState *ls) {
   int nlocalvar = fs->nlocalvar;
   chunk(ls);
   luaK_adjuststack(fs, fs->nlocalvar - nlocalvar);  /* remove local variables */
-  removelocalvars(ls, fs->nlocalvar - nlocalvar, fs->lastsetline);
+  removelocalvars(ls, fs->nlocalvar - nlocalvar);
 }
 
 
@@ -907,15 +909,20 @@ static void repeatstat (LexState *ls, int line) {
 }
 
 
-static void forstat (LexState *ls, int line) {
-  /* forstat -> FOR NAME '=' expr1 ',' expr1 [',' expr1] DO block END */
+static void forbody (LexState *ls, OpCode prepfor, OpCode loopfor) {
   FuncState *fs = ls->fs;
-  int prep;
-  int blockinit;
-  Breaklabel bl;
-  enterbreak(fs, &bl);
-  setline_and_next(ls);  /* skip for */
-  store_localvar(ls, str_checkname(ls), 0);  /* control variable */
+  int prep = luaK_code0(fs, prepfor);
+  int blockinit = luaK_getlabel(fs);
+  check(ls, TK_DO);
+  block(ls);
+  luaK_patchlist(fs, prep, luaK_getlabel(fs));
+  luaK_patchlist(fs, luaK_code0(fs, loopfor), blockinit);
+}
+
+
+static void fornum (LexState *ls, TString *varname) {
+  FuncState *fs = ls->fs;
+  store_localvar(ls, varname, 0);
   check(ls, '=');
   exp1(ls);  /* initial value */
   check(ls, ',');
@@ -924,18 +931,49 @@ static void forstat (LexState *ls, int line) {
     exp1(ls);  /* optional step */
   else
     luaK_code1(fs, OP_PUSHINT, 1);  /* default step */
-  adjustlocalvars(ls, 1, 0);  /* init scope for control variable */
-  add_localvar(ls, " limit ");
-  add_localvar(ls, " step ");
-  prep = luaK_code0(fs, OP_FORPREP);
-  blockinit = luaK_getlabel(fs);
-  check(ls, TK_DO);
-  block(ls);
-  luaK_patchlist(fs, prep, luaK_getlabel(fs));
-  luaK_patchlist(fs, luaK_code0(fs, OP_FORLOOP), blockinit);
+  adjustlocalvars(ls, 1);  /* scope for control variables */
+  add_localvar(ls, "*limit*");
+  add_localvar(ls, "*step*");
+  forbody(ls, OP_FORPREP, OP_FORLOOP);
+  removelocalvars(ls, 3);
+}
+
+
+static void forlist (LexState *ls, TString *indexname) {
+  TString *valname;
+  check(ls, ',');
+  valname = str_checkname(ls);
+  /* next test is dirty, but avoids `in' being a reserved word */
+  if (ls->token != TK_NAME || ls->seminfo.ts != luaS_new(ls->L, "in"))
+    luaK_error(ls, "`in' expected");
+  next(ls);  /* skip `in' */
+  exp1(ls);  /* table */
+  add_localvar(ls, "*table*");
+  add_localvar(ls, "*counter*");
+  store_localvar(ls, indexname, 0);
+  store_localvar(ls, valname, 1);
+  adjustlocalvars(ls, 2);  /* scope for control variable */
+  forbody(ls, OP_LFORPREP, OP_LFORLOOP);
+  removelocalvars(ls, 4);
+}
+
+
+static void forstat (LexState *ls, int line) {
+  /* forstat -> FOR NAME '=' expr1 ',' expr1 [',' expr1] DO block END */
+  /* forstat -> FOR NAME1, NAME2 IN expr1 DO block END */
+  FuncState *fs = ls->fs;
+  TString *varname;
+  Breaklabel bl;
+  enterbreak(fs, &bl);
+  setline_and_next(ls);  /* skip `for' */
+  varname = str_checkname(ls);  /* first variable name */
+  switch (ls->token) {
+    case '=': fornum(ls, varname); break;
+    case ',': forlist(ls, varname); break;
+    default: luaK_error(ls, "`=' or `,' expected");
+  }
   check_END(ls, TK_FOR, line);
   leavebreak(fs, &bl);
-  removelocalvars(ls, 3, fs->lastsetline);
 }
 
 
@@ -999,13 +1037,12 @@ static int decinit (LexState *ls) {
 
 static void localstat (LexState *ls) {
   /* stat -> LOCAL localnamelist decinit */
-  FuncState *fs = ls->fs;
   int nvars;
   int nexps;
   setline_and_next(ls);  /* skip LOCAL */
   nvars = localnamelist(ls);
   nexps = decinit(ls);
-  adjustlocalvars(ls, nvars, fs->lastsetline);
+  adjustlocalvars(ls, nvars);
   adjust_mult_assign(ls, nvars, nexps);
 }
 
diff --git a/ltests.c b/ltests.c
index 4646c5c9..8502c95a 100644
--- a/ltests.c
+++ b/ltests.c
@@ -1,5 +1,5 @@
 /*
-** $Id: ltests.c,v 1.17 2000/05/08 19:32:53 roberto Exp roberto $
+** $Id: ltests.c,v 1.18 2000/05/10 16:33:20 roberto Exp roberto $
 ** Internal Module for Debugging of the Lua Implementation
 ** See Copyright Notice in lua.h
 */
@@ -103,8 +103,10 @@ static int pushop (lua_State *L, Instruction i) {
     case OP_JMPONF: S("JMPONF"); break;
     case OP_JMP: S("JMP"); break;
     case OP_PUSHNILJMP: O("PUSHNILJMP"); break;
-    case OP_FORPREP: S("OP_FORPREP"); break;
-    case OP_FORLOOP: S("OP_FORLOOP"); break;
+    case OP_FORPREP: S("FORPREP"); break;
+    case OP_FORLOOP: S("FORLOOP"); break;
+    case OP_LFORPREP: S("LFORPREP"); break;
+    case OP_LFORLOOP: S("LFORLOOP"); break;
     case OP_CLOSURE: AB("CLOSURE"); break;
     case OP_SETLINE: U("SETLINE"); break;
   }
diff --git a/lvm.c b/lvm.c
index 9b0ea0c8..ecf2a49c 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lvm.c,v 1.104 2000/04/19 13:36:25 roberto Exp roberto $
+** $Id: lvm.c,v 1.105 2000/05/08 19:32:53 roberto Exp roberto $
 ** Lua virtual machine
 ** See Copyright Notice in lua.h
 */
@@ -11,6 +11,7 @@
 
 #define LUA_REENTRANT
 
+#include "lapi.h"
 #include "lauxlib.h"
 #include "ldebug.h"
 #include "ldo.h"
@@ -634,12 +635,40 @@ StkId luaV_execute (lua_State *L, const Closure *cl, StkId base) {
         if (ttype(top-3) != TAG_NUMBER)
           lua_error(L, "`for' index must be a number");
         index = nvalue(top-3)+step;
-        if ((step>0) ? index<=limit : index>=limit) {
+        if ((step>0) ? index>limit : index<limit)
+          top -= 3;  /* end loop: remove control variables */
+        else {
           nvalue(top-3) = index;
           pc += GETARG_S(i);
         }
-        else  /* end of `for': remove control variables */
-          top -= 3;
+        break;
+      }
+
+      case OP_LFORPREP: {
+        if (ttype(top-1) != TAG_TABLE)
+          lua_error(L, "`for' table must be a table");
+        top += 3;  /* counter + index,value */
+        ttype(top-3) = TAG_NUMBER;
+        nvalue(top-3) = 0.0;  /* counter */
+        ttype(top-2) = ttype(top-1) = TAG_NIL;
+        pc += GETARG_S(i);
+        break;
+      }
+
+      case OP_LFORLOOP: {
+        int n;
+        top -= 2;  /* remove old index,value */
+        LUA_ASSERT(L, ttype(top-2) == TAG_TABLE, "invalid table");
+        LUA_ASSERT(L, ttype(top-1) == TAG_NUMBER, "invalid counter");
+        L->top = top;
+        n = luaA_next(L, avalue(top-2), (int)nvalue(top-1));
+        if (n == 0)  /* end loop? */
+          top -= 2;  /* remove table and counter */
+        else {
+          nvalue(top-1) = (Number)n;
+          top += 2;  /* new index,value */
+          pc += GETARG_S(i);  /* repeat loop */
+        }
         break;
       }