1 files changed, 14 insertions, 16 deletions

diff --git a/lbaselib.c b/lbaselib.c
index 5260fc35..4536d617 100644
--- a/lbaselib.c
+++ b/lbaselib.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lbaselib.c,v 1.273 2011/11/30 13:03:24 roberto Exp roberto $
+** $Id: lbaselib.c,v 1.274 2012/04/27 14:13:19 roberto Exp roberto $
 ** Basic library
 ** See Copyright Notice in lua.h
 */
@@ -242,10 +242,16 @@ static int luaB_ipairs (lua_State *L) {
 }
 
 
-static int load_aux (lua_State *L, int status) {
-  if (status == LUA_OK)
+static int load_aux (lua_State *L, int status, int envidx) {
+  if (status == LUA_OK) {
+    if (envidx != 0) {  /* 'env' parameter? */
+      lua_pushvalue(L, envidx);  /* environment for loaded function */
+      if (!lua_setupvalue(L, -2, 1))  /* set it as 1st upvalue */
+        lua_pop(L, 1);  /* remove 'env' if not used by previous call */
+    }
     return 1;
-  else {
+  }
+  else {  /* error (message is on top of the stack) */
     lua_pushnil(L);
     lua_insert(L, -2);  /* put before error message */
     return 2;  /* return nil plus error message */
@@ -256,13 +262,9 @@ static int load_aux (lua_State *L, int status) {
 static int luaB_loadfile (lua_State *L) {
   const char *fname = luaL_optstring(L, 1, NULL);
   const char *mode = luaL_optstring(L, 2, NULL);
-  int env = !lua_isnone(L, 3);  /* 'env' parameter? */
+  int env = (!lua_isnone(L, 3) ? 3 : 0);  /* 'env' index or 0 if no 'env' */
   int status = luaL_loadfilex(L, fname, mode);
-  if (status == LUA_OK && env) {  /* 'env' parameter? */
-    lua_pushvalue(L, 3);
-    lua_setupvalue(L, -2, 1);  /* set it as 1st upvalue of loaded chunk */
-  }
-  return load_aux(L, status);
+  return load_aux(L, status, env);
 }
 
 
@@ -307,9 +309,9 @@ static const char *generic_reader (lua_State *L, void *ud, size_t *size) {
 static int luaB_load (lua_State *L) {
   int status;
   size_t l;
-  int top = lua_gettop(L);
   const char *s = lua_tolstring(L, 1, &l);
   const char *mode = luaL_optstring(L, 3, "bt");
+  int env = (!lua_isnone(L, 4) ? 4 : 0);  /* 'env' index or 0 if no 'env' */
   if (s != NULL) {  /* loading a string? */
     const char *chunkname = luaL_optstring(L, 2, s);
     status = luaL_loadbufferx(L, s, l, chunkname, mode);
@@ -320,11 +322,7 @@ static int luaB_load (lua_State *L) {
     lua_settop(L, RESERVEDSLOT);  /* create reserved slot */
     status = lua_load(L, generic_reader, NULL, chunkname, mode);
   }
-  if (status == LUA_OK && top >= 4) {  /* is there an 'env' argument */
-    lua_pushvalue(L, 4);  /* environment for loaded function */
-    lua_setupvalue(L, -2, 1);  /* set it as 1st upvalue */
-  }
-  return load_aux(L, status);
+  return load_aux(L, status, env);
 }
 
 /* }====================================================== */