7 files changed, 68 insertions, 22 deletions

diff --git a/ldump.c b/ldump.c
index e8337f03..e59c4224 100644
--- a/ldump.c
+++ b/ldump.c
@@ -1,5 +1,5 @@
 /*
-** $Id: ldump.c,v 2.10 2008/07/03 14:25:05 roberto Exp roberto $
+** $Id: ldump.c,v 2.11 2009/09/28 16:32:50 roberto Exp roberto $
 ** save precompiled Lua chunks
 ** See Copyright Notice in lua.h
 */
@@ -145,6 +145,7 @@ static void DumpFunction(const Proto* f, const TString* p, DumpState* D)
  DumpChar(f->numparams,D);
  DumpChar(f->is_vararg,D);
  DumpChar(f->maxstacksize,D);
+ DumpChar(f->envreg,D);
  DumpCode(f,D);
  DumpConstants(f,D);
  DumpUpvalues(f,D);
diff --git a/lfunc.c b/lfunc.c
index c6c63f2e..bca286ad 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lfunc.c,v 2.14 2009/04/17 14:40:13 roberto Exp roberto $
+** $Id: lfunc.c,v 2.15 2009/09/28 16:32:50 roberto Exp roberto $
 ** Auxiliary functions to manipulate prototypes and closures
 ** See Copyright Notice in lua.h
 */
@@ -16,6 +16,7 @@
 #include "lgc.h"
 #include "lmem.h"
 #include "lobject.h"
+#include "lopcodes.h"
 #include "lstate.h"
 
 
@@ -133,6 +134,7 @@ Proto *luaF_newproto (lua_State *L) {
   f->linedefined = 0;
   f->lastlinedefined = 0;
   f->source = NULL;
+  f->envreg = NO_REG;
   return f;
 }
 
diff --git a/lobject.h b/lobject.h
index 8b93f70a..581a25a0 100644
--- a/lobject.h
+++ b/lobject.h
@@ -1,5 +1,5 @@
 /*
-** $Id: lobject.h,v 2.28 2009/07/15 18:37:19 roberto Exp roberto $
+** $Id: lobject.h,v 2.29 2009/09/28 16:32:50 roberto Exp roberto $
 ** Type definitions for Lua objects
 ** See Copyright Notice in lua.h
 */
@@ -261,6 +261,7 @@ typedef struct Proto {
   lu_byte numparams;
   lu_byte is_vararg;
   lu_byte maxstacksize;
+  lu_byte envreg;  /* register in outer function with initial environment */
 } Proto;
 
 
diff --git a/lparser.c b/lparser.c
index 935526ce..1ab1cc69 100644
--- a/lparser.c
+++ b/lparser.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lparser.c,v 2.66 2009/09/23 20:14:00 roberto Exp roberto $
+** $Id: lparser.c,v 2.67 2009/09/28 16:32:50 roberto Exp roberto $
 ** Lua Parser
 ** See Copyright Notice in lua.h
 */
@@ -227,10 +227,8 @@ static void markupval (FuncState *fs, int level) {
 
 
 static int singlevaraux (FuncState *fs, TString *n, expdesc *var, int base) {
-  if (fs == NULL) {  /* no more levels? */
-    init_exp(var, VGLOBAL, NO_REG);  /* default is global variable */
-    return VGLOBAL;
-  }
+  if (fs == NULL)  /* no more levels? */
+    return VGLOBAL;  /* default is global variable */
   else {
     int v = searchvar(fs, n);  /* look up at current level */
     if (v >= 0) {
@@ -253,8 +251,16 @@ static int singlevaraux (FuncState *fs, TString *n, expdesc *var, int base) {
 static void singlevar (LexState *ls, expdesc *var) {
   TString *varname = str_checkname(ls);
   FuncState *fs = ls->fs;
-  if (singlevaraux(fs, varname, var, 1) == VGLOBAL)
-    var->u.s.info = luaK_stringK(fs, varname);  /* info points to global name */
+  if (singlevaraux(fs, varname, var, 1) == VGLOBAL) {
+    if (fs->envreg == NO_REG)  /* regular global? */
+      init_exp(var, VGLOBAL, luaK_stringK(fs, varname));
+    else {  /* "globals" are in current lexical environment */
+      expdesc key;
+      init_exp(var, VLOCAL, fs->envreg);  /* current environment */
+      codestring(ls, &key, varname);  /* key is variable name */
+      luaK_indexed(fs, var, &key);  /* env[varname] */
+    }
+  }
 }
 
 
@@ -313,15 +319,17 @@ static void leaveblock (FuncState *fs) {
 }
 
 
-static void pushclosure (LexState *ls, FuncState *func, expdesc *v) {
+static void pushclosure (LexState *ls, Proto *clp, expdesc *v) {
   FuncState *fs = ls->fs->prev;
-  Proto *f = fs->f;
+  Proto *f = fs->f;  /* prototype of function creating new closure */
   int oldsize = f->sizep;
   luaM_growvector(ls->L, f->p, fs->np, f->sizep, Proto *,
                   MAXARG_Bx, "functions");
   while (oldsize < f->sizep) f->p[oldsize++] = NULL;
-  f->p[fs->np++] = func->f;
-  luaC_objbarrier(ls->L, f, func->f);
+  f->p[fs->np++] = clp;
+  /* initial environment for new function is current lexical environment */
+  clp->envreg = fs->envreg;
+  luaC_objbarrier(ls->L, f, clp);
   init_exp(v, VRELOCABLE, luaK_codeABx(fs, OP_CLOSURE, 0, fs->np-1));
 }
 
@@ -342,6 +350,7 @@ static void open_func (LexState *ls, FuncState *fs) {
   fs->nups = 0;
   fs->nlocvars = 0;
   fs->nactvar = 0;
+  fs->envreg = NO_REG;
   fs->bl = NULL;
   fs->h = luaH_new(L);
   /* anchor table of constants (to avoid being collected) */
@@ -589,7 +598,7 @@ static void body (LexState *ls, expdesc *e, int needself, int line) {
   chunk(ls);
   new_fs.f->lastlinedefined = ls->linenumber;
   check_match(ls, TK_END, TK_FUNCTION, line);
-  pushclosure(ls, &new_fs, e);
+  pushclosure(ls, new_fs.f, e);
   close_func(ls);
 }
 
@@ -1035,11 +1044,12 @@ static void repeatstat (LexState *ls, int line) {
 
 static int exp1 (LexState *ls) {
   expdesc e;
-  int k;
+  int reg;
   expr(ls, &e);
-  k = e.k;
   luaK_exp2nextreg(ls->fs, &e);
-  return k;
+  lua_assert(e.k == VNONRELOC);
+  reg = e.u.s.info;
+  return reg;
 }
 
 
@@ -1226,6 +1236,24 @@ static void funcstat (LexState *ls, int line) {
 }
 
 
+static void instat (LexState *ls, int line) {
+  /* instat -> IN exp DO block END */
+  FuncState *fs = ls->fs;
+  int oldenv = fs->envreg;  /* save current environment */
+  BlockCnt bl;
+  luaX_next(ls);  /* skip IN */
+  enterblock(fs, &bl, 0);  /* scope for environment variable */
+  new_localvarliteral(ls, "(environment)", 0);
+  fs->envreg = exp1(ls);  /* new environment */
+  adjustlocalvars(ls, 1);
+  checknext(ls, TK_DO);
+  block(ls);
+  leaveblock(fs);
+  check_match(ls, TK_END, TK_IN, line);
+  fs->envreg = oldenv;  /* restore outer environment */
+}
+
+
 static void exprstat (LexState *ls) {
   /* stat -> func | assignment */
   FuncState *fs = ls->fs;
@@ -1290,6 +1318,10 @@ static int statement (LexState *ls) {
       check_match(ls, TK_END, TK_DO, line);
       return 0;
     }
+    case TK_IN: {
+      instat(ls, line);
+      return 0;
+    }
     case TK_FOR: {  /* stat -> forstat */
       forstat(ls, line);
       return 0;
diff --git a/lparser.h b/lparser.h
index 5e1856e2..850b787d 100644
--- a/lparser.h
+++ b/lparser.h
@@ -1,5 +1,5 @@
 /*
-** $Id: lparser.h,v 1.58 2008/05/08 15:44:51 roberto Exp roberto $
+** $Id: lparser.h,v 1.59 2009/09/28 16:32:50 roberto Exp roberto $
 ** Lua Parser
 ** See Copyright Notice in lua.h
 */
@@ -71,7 +71,8 @@ typedef struct FuncState {
   short nlocvars;  /* number of elements in `locvars' */
   lu_byte nactvar;  /* number of active local variables */
   lu_byte nups;  /* number of upvalues */
-  vardesc actvar[LUAI_MAXVARS];  /* declared-variable stack */
+  lu_byte envreg;  /* register holding current lexical environment */
+  vardesc actvar[LUAI_MAXVARS];  /* stack of active variables */
 } FuncState;
 
 
diff --git a/lundump.c b/lundump.c
index 4ee11ac6..753f9964 100644
--- a/lundump.c
+++ b/lundump.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lundump.c,v 2.10 2009/04/30 17:42:21 roberto Exp roberto $
+** $Id: lundump.c,v 2.11 2009/09/28 16:32:50 roberto Exp roberto $
 ** load precompiled Lua chunks
 ** See Copyright Notice in lua.h
 */
@@ -180,6 +180,7 @@ static Proto* LoadFunction(LoadState* S, TString* p)
  f->numparams=LoadByte(S);
  f->is_vararg=LoadByte(S);
  f->maxstacksize=LoadByte(S);
+ f->envreg=LoadByte(S);
  LoadCode(S,f);
  LoadConstants(S,f);
  LoadUpvalues(S,f);
diff --git a/lvm.c b/lvm.c
index 7473d73b..004e112c 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lvm.c,v 2.97 2009/09/23 20:33:05 roberto Exp roberto $
+** $Id: lvm.c,v 2.98 2009/09/28 16:32:50 roberto Exp roberto $
 ** Lua virtual machine
 ** See Copyright Notice in lua.h
 */
@@ -780,6 +780,14 @@ void luaV_execute (lua_State *L) {
         int j;
         ncl->l.p = p;
         setclvalue(L, ra, ncl);  /* anchor new closure in stack */
+        if (p->envreg != NO_REG) {  /* lexical environment? */
+          StkId env = base + p->envreg;
+          if (!ttistable(env))
+            luaG_runerror(L, "environment is not a table: "
+                             "cannot create closure");
+          else
+            ncl->l.env = hvalue(env);
+        }
         for (j = 0; j < nup; j++) {  /* fill in upvalues */
           if (uv[j].instack)  /* upvalue refers to local variable? */
             ncl->l.upvals[j] = luaF_findupval(L, base + uv[j].idx);