Direct enumeration of collections in Lua for..in loops.

- Support direct for..in enumeration for any type of collection (Lua table and native collection classes), without needing to call `pair` or `ipair`, in the form `for var1, var2, ... in collection do ... end`.
- Simple for..in enumeration can be customized by defining a "__forgen" metamethod that shall return the 3 values needed by the generic for loop: an iterator function, a state, and the initial value of the iterator variable. If no "__forgen" metamethod is defined and the collection is a Lua table, direct for..in enumeration is equivalent to calling the standard library function `pair()`.
- Internally, a new Lua op-code 'OP_TFORPREP' has been added, that replaces the initial JMP instruction when initializing a generic for loop. When executed, this opcode checks for the presence of a "__forgen" metamethod in the loop base object and calls it if present.

lopcodes.c

@@ -92,6 +92,7 @@ LUAI_DDEF const char *const luaP_opnames[NUM_OPCODES+1] = {
   "FORPREP",
   "TFORCALL",
   "TFORLOOP",
+  "TFORPREP",
   "SETLIST",
   "CLOSURE",
   "VARARG",
@@ -177,6 +178,7 @@ LUAI_DDEF const lu_byte luaP_opmodes[NUM_OPCODES] = {
  ,opmode(0, 0, 0, 1, iABx)		/* OP_FORPREP */
  ,opmode(0, 0, 0, 0, iABC)		/* OP_TFORCALL */
  ,opmode(0, 0, 0, 1, iABx)		/* OP_TFORLOOP */
+ ,opmode(0, 0, 0, 1, iABx)      /* OP_TFORPREP */
  ,opmode(0, 1, 0, 0, iABC)		/* OP_SETLIST */
  ,opmode(0, 0, 0, 1, iABx)		/* OP_CLOSURE */
  ,opmode(1, 0, 0, 1, iABC)		/* OP_VARARG */

lopcodes.h

@@ -286,7 +286,8 @@ OP_FORPREP,/*	A Bx	R(A)-=R(A+2); pc+=Bx				*/
 
 OP_TFORCALL,/*	A C	R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2));	*/
 OP_TFORLOOP,/*	A Bx	if R(A+1) ~= nil then { R(A)=R(A+1); pc -= Bx }	*/
-
+OP_TFORPREP,/*  A sBx  if TM(R(A), TM_FORGEN) is function then { R(A), R(A+1), R(A+2) = TM(R(A), TM_FORGEN)(R(A), R(A+1), R(A+2)) }; pc += sBx */
+
 OP_SETLIST,/*	A B C	R(A)[(C-1)*FPF+i] := R(A+i), 1 <= i <= B	*/
 
 OP_CLOSURE,/*	A Bx	R(A) := closure(KPROTO[Bx])			*/

lparser.c

@@ -1364,7 +1364,7 @@ static void forbody (LexState *ls, int base, int line, int nvars, int kind) {
   checknext(ls, TK_DO);
   prep = (kind == 0) ? luaK_codeABx(fs, OP_FORPREP, base, 0)
        : (kind == 1) ? luaK_codeABx(fs, OP_FORPREP1, base, 0)
-       : luaK_jump(fs);
+       : luaK_codeABx(fs, OP_TFORPREP, base, 0);
   enterblock(fs, &bl, 0);  /* scope for declared variables */
   adjustlocalvars(ls, nvars);
   luaK_reserveregs(fs, nvars);

ltm.c

@@ -44,7 +44,7 @@ void luaT_init (lua_State *L) {
     "__div", "__idiv",
     "__band", "__bor", "__bxor", "__shl", "__shr",
     "__unm", "__bnot", "__lt", "__le",
-    "__concat", "__call"
+    "__concat", "__call", "__forgen"
   };
   int i;
   for (i=0; i<TM_N; i++) {

ltm.h

@@ -43,6 +43,7 @@ typedef enum {
   TM_LE,
   TM_CONCAT,
   TM_CALL,
+  TM_FORGEN,
   TM_N		/* number of elements in the enum */
 } TMS;
 

lvm.c

@@ -664,6 +664,46 @@ lua_Integer luaV_shiftl (lua_Integer x, lua_Integer y) {
   }
 }
 
+static int tableiterator (lua_State *L) {
+  luaD_checkstack(L, 2);
+  if (lua_next(L, 1))
+    return 2;
+  else {
+    lua_pushnil(L);
+    return 1;
+  }
+}
+
+static void tryforgenTM(lua_State *L, StkId ra)
+{
+  const TValue *object =  s2v(ra);
+  const TValue *tm = luaT_gettmbyobj(L, object, TM_FORGEN);
+  ptrdiff_t rar = savestack(L, ra);
+  if (ttisfunction(tm)) {
+    StkId func = L->top;
+    setobj2s(L, func, tm);     /* forgen TM function (assume EXTRA_STACK) */
+    setobj2s(L, func + 1, object);  /* object (1st exp) */
+    setobjs2s(L, func + 2, ra + 1); /* 2nd exp */
+    setobjs2s(L, func + 3, ra + 2); /* 3nd exp */
+    L->top += 4;
+    /* metamethod may yield only when called from Lua code */
+    if (isLuacode(L->ci))
+      luaD_call(L, func, 3);
+    else
+      luaD_callnoyield(L, func, 3);
+    ra = restorestack(L, rar);  /* previous call may change stack */
+    setobjs2s(L, ra, L->top - 3);
+    setobjs2s(L, ra + 1, L->top - 2);
+    setobjs2s(L, ra + 2, L->top - 1);
+    L->top -= 3;
+  }
+  else if (ttisnil(tm) && ttistable(object)) {
+    /* table without forgen TM: use the default table iterator */
+    setnilvalue(s2v(ra + 2));
+    setobj2s(L, ra + 1, object)
+    setfvalue(s2v(ra), tableiterator);
+  }
+}
 
 /*
 ** check whether cached closure in prototype 'p' may be reused, that is,
@@ -1753,6 +1793,10 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         }
         vmbreak;
       }
+      vmcase(OP_TFORPREP) {
+        Protect(tryforgenTM(L, ra));
+        ci->u.l.savedpc += GETARG_sBx(i);
+      }
       vmcase(OP_SETLIST) {
         int n = GETARG_B(i);
         int c = GETARG_C(i);