yields accross metamethods and for iterators (except for __concat)

ldo.c

@@ -1,5 +1,5 @@
 /*
-** $Id: ldo.c,v 2.47 2008/08/13 17:02:42 roberto Exp roberto $
+** $Id: ldo.c,v 2.48 2008/08/26 13:27:42 roberto Exp roberto $
 ** Stack and Call structure of Lua
 ** See Copyright Notice in lua.h
 */
@@ -374,7 +374,6 @@ int luaD_poscall (lua_State *L, StkId firstResult) {
 */
 void luaD_call (lua_State *L, StkId func, int nResults) {
   global_State *g = G(L);
-  lua_assert(g->nCcalls >= L->baseCcalls);
   if (++g->nCcalls >= LUAI_MAXCCALLS) {
     if (g->nCcalls == LUAI_MAXCCALLS)
       luaG_runerror(L, "C stack overflow");
@@ -388,6 +387,55 @@ void luaD_call (lua_State *L, StkId func, int nResults) {
 }
 
 
+static void unroll (lua_State *L) {
+  for (;;) {
+    Instruction inst;
+    luaV_execute(L);  /* execute up to higher C 'boundary' */
+    if (L->ci == L->base_ci)  /* stack is empty? */
+      return;  /* coroutine finished normally */
+    L->baseCcalls--;  /* undo increment that allows yields */
+    inst = *(L->savedpc - 1);  /* interrupted instruction */
+    switch (GET_OPCODE(inst)) {  /* finish its execution */
+      case OP_ADD: case OP_SUB: case OP_MUL: case OP_DIV:
+      case OP_MOD: case OP_POW: case OP_UNM: case OP_LEN:
+      case OP_GETGLOBAL: case OP_GETTABLE: case OP_SELF: {
+        setobjs2s(L, L->base + GETARG_A(inst), --L->top);
+        break;
+      }
+      case OP_LE: case OP_LT: case OP_EQ: {
+        int res;
+        L->top--;
+        res = !l_isfalse(L->top);
+        /* cannot call metamethod with K operand */
+        lua_assert(!ISK(GETARG_B(inst)));
+        if (GET_OPCODE(inst) == OP_LE &&  /* "<=" using "<" instead? */
+            ttisnil(luaT_gettmbyobj(L, L->base + GETARG_B(inst), TM_LE)))
+          res = !res;  /* invert result */
+        lua_assert(GET_OPCODE(*L->savedpc) == OP_JMP);
+        if (res == GETARG_A(inst))
+          L->savedpc += GETARG_sBx(*L->savedpc);  /* jump */
+        L->savedpc++;  /* skip jump instruction */
+        break;
+      }
+      case OP_SETGLOBAL: case OP_SETTABLE:
+        break;  /* nothing to be done */
+      case OP_TFORLOOP: {
+        StkId cb = L->base + GETARG_A(inst) + 3;
+        L->top = L->ci->top;
+        lua_assert(GET_OPCODE(*L->savedpc) == OP_JMP);
+        if (!ttisnil(cb)) {  /* continue loop? */
+          setobjs2s(L, cb - 1, cb);  /* save control variable */
+          L->savedpc += GETARG_sBx(*L->savedpc);  /* jump back */
+        }
+        L->savedpc++;
+        break;
+      }
+      default: lua_assert(0);
+    }
+  }
+}
+
+
 static void resume (lua_State *L, void *ud) {
   StkId firstArg = cast(StkId, ud);
   CallInfo *ci = L->ci;
@@ -399,17 +447,17 @@ static void resume (lua_State *L, void *ud) {
   else {  /* resuming from previous yield */
     lua_assert(L->status == LUA_YIELD);
     L->status = LUA_OK;
-    if (!isLua(ci)) {  /* `common' yield? */
+    if (isLua(ci))  /* yielded inside a hook? */
+      L->base = L->ci->base;  /* just continue its execution */
+    else {  /* 'common' yield */
       /* finish interrupted execution of `OP_CALL' */
       lua_assert(GET_OPCODE(*((ci-1)->savedpc - 1)) == OP_CALL ||
                  GET_OPCODE(*((ci-1)->savedpc - 1)) == OP_TAILCALL);
       if (luaD_poscall(L, firstArg))  /* complete it... */
         L->top = L->ci->top;  /* and correct top if not multiple results */
     }
-    else  /* yielded inside a hook: just continue its execution */
-      L->base = L->ci->base;
   }
-  luaV_execute(L);
+  unroll(L);
 }
 
 
@@ -432,10 +480,11 @@ LUA_API int lua_resume (lua_State *L, int nargs) {
       return resume_error(L, "cannot resume non-suspended coroutine");
   }
   luai_userstateresume(L, nargs);
-  lua_assert(L->errfunc == 0 && L->baseCcalls == 0);
+  lua_assert(L->errfunc == 0); 
   if (G(L)->nCcalls >= LUAI_MAXCCALLS)
       return resume_error(L, "C stack overflow");
-  L->baseCcalls = ++G(L)->nCcalls;
+  ++G(L)->nCcalls;  /* count resume */
+  L->baseCcalls += G(L)->nCcalls;
   status = luaD_rawrunprotected(L, resume, L->top - nargs);
   if (status != LUA_OK && status != LUA_YIELD) {  /* error? */
     L->status = cast_byte(status);  /* mark thread as `dead' */
@@ -443,11 +492,10 @@ LUA_API int lua_resume (lua_State *L, int nargs) {
     L->ci->top = L->top;
   }
   else {
-    lua_assert(L->baseCcalls == G(L)->nCcalls);
     lua_assert(status == L->status);
   }
+  L->baseCcalls -= G(L)->nCcalls;
   --G(L)->nCcalls;
-  L->baseCcalls = 0;
   lua_unlock(L);
   return status;
 }
@@ -456,6 +504,7 @@ LUA_API int lua_resume (lua_State *L, int nargs) {
 LUA_API int lua_yield (lua_State *L, int nresults) {
   luai_userstateyield(L, nresults);
   lua_lock(L);
+/*printf("yield: %d - %d\n", G(L)->nCcalls, L->baseCcalls);*/
   if (G(L)->nCcalls > L->baseCcalls)
     luaG_runerror(L, "attempt to yield across metamethod/C-call boundary");
   L->base = L->top - nresults;  /* protect stack slots below */

lvm.c

@@ -1,5 +1,5 @@
 /*
-** $Id: lvm.c,v 2.76 2008/08/26 13:27:42 roberto Exp roberto $
+** $Id: lvm.c,v 2.77 2008/09/09 13:53:02 roberto Exp roberto $
 ** Lua virtual machine
 ** See Copyright Notice in lua.h
 */
@@ -76,31 +76,24 @@ static void traceexec (lua_State *L) {
 }
 
 
-static void callTMres (lua_State *L, StkId res, const TValue *f,
-                        const TValue *p1, const TValue *p2) {
-  ptrdiff_t result = savestack(L, res);
-  setobj2s(L, L->top, f);  /* push function */
-  setobj2s(L, L->top+1, p1);  /* 1st argument */
-  setobj2s(L, L->top+2, p2);  /* 2nd argument */
-  L->top += 3;
-  luaD_checkstack(L, 0);
-  luaD_call(L, L->top - 3, 1);
-  res = restorestack(L, result);
-  L->top--;
-  setobjs2s(L, res, L->top);
-}
-
-
-
 static void callTM (lua_State *L, const TValue *f, const TValue *p1,
-                    const TValue *p2, const TValue *p3) {
-  setobj2s(L, L->top, f);  /* push function */
-  setobj2s(L, L->top+1, p1);  /* 1st argument */
-  setobj2s(L, L->top+2, p2);  /* 2nd argument */
-  setobj2s(L, L->top+3, p3);  /* 3th argument */
-  L->top += 4;
+                    const TValue *p2, TValue *p3, int hasres) {
+  ptrdiff_t result = savestack(L, p3);
+  int oldbase = L->baseCcalls;
+  setobj2s(L, L->top++, f);  /* push function */
+  setobj2s(L, L->top++, p1);  /* 1st argument */
+  setobj2s(L, L->top++, p2);  /* 2nd argument */
+  if (!hasres)  /* no result? 'p3' is third argument */
+    setobj2s(L, L->top++, p3);  /* 3th argument */
   luaD_checkstack(L, 0);
-  luaD_call(L, L->top - 4, 0);
+  if (isLua(L->ci))  /* metamethod invoked from a Lua function? */
+    L->baseCcalls++;  /* allow it to yield */
+  luaD_call(L, L->top - (4 - hasres), hasres);
+  L->baseCcalls = oldbase;
+  if (hasres) {  /* if has result, move it to its place */
+    p3 = restorestack(L, result);
+    setobjs2s(L, p3, --L->top);
+  }
 }
 
 
@@ -121,7 +114,7 @@ void luaV_gettable (lua_State *L, const TValue *t, TValue *key, StkId val) {
     else if (ttisnil(tm = luaT_gettmbyobj(L, t, TM_INDEX)))
       luaG_typeerror(L, t, "index");
     if (ttisfunction(tm)) {
-      callTMres(L, val, tm, t, key);
+      callTM(L, tm, t, key, val, 1);
       return;
     }
     t = tm;  /* else repeat with `tm' */
@@ -148,7 +141,7 @@ void luaV_settable (lua_State *L, const TValue *t, TValue *key, StkId val) {
     else if (ttisnil(tm = luaT_gettmbyobj(L, t, TM_NEWINDEX)))
       luaG_typeerror(L, t, "index");
     if (ttisfunction(tm)) {
-      callTM(L, tm, t, key, val);
+      callTM(L, tm, t, key, val, 0);
       return;
     }
     t = tm;  /* else repeat with `tm' */
@@ -163,7 +156,7 @@ static int call_binTM (lua_State *L, const TValue *p1, const TValue *p2,
   if (ttisnil(tm))
     tm = luaT_gettmbyobj(L, p2, event);  /* try second operand */
   if (ttisnil(tm)) return 0;
-  callTMres(L, res, tm, p1, p2);
+  callTM(L, tm, p1, p2, res, 1);
   return 1;
 }
 
@@ -190,7 +183,7 @@ static int call_orderTM (lua_State *L, const TValue *p1, const TValue *p2,
   tm2 = luaT_gettmbyobj(L, p2, event);
   if (!luaO_rawequalObj(tm1, tm2))  /* different metamethods? */
     return -1;
-  callTMres(L, L->top, tm1, p1, p2);
+  callTM(L, tm1, p1, p2, L->top, 1);
   return !l_isfalse(L->top);
 }
 
@@ -268,7 +261,7 @@ int luaV_equalval_ (lua_State *L, const TValue *t1, const TValue *t2) {
     default: return gcvalue(t1) == gcvalue(t2);
   }
   if (tm == NULL) return 0;  /* no TM? */
-  callTMres(L, L->top, tm, t1, t2);  /* call TM */
+  callTM(L, tm, t1, t2, L->top, 1);  /* call TM */
   return !l_isfalse(L->top);
 }
 
@@ -336,7 +329,7 @@ static void objlen (lua_State *L, StkId ra, const TValue *rb) {
       break;
     }
   }
-  callTMres(L, ra, tm, rb, luaO_nilobject);
+  callTM(L, tm, rb, luaO_nilobject, ra, 1);
 }
 
 
@@ -680,7 +673,9 @@ void luaV_execute (lua_State *L) {
         setobjs2s(L, cb+1, ra+1);
         setobjs2s(L, cb, ra);
         L->top = cb+3;  /* func. + 2 args (state and index) */
+        L->baseCcalls++;
         Protect(luaD_call(L, cb, GETARG_C(i)));
+        L->baseCcalls--;
         L->top = L->ci->top;
         cb = RA(i) + 3;  /* previous call may change the stack */
         if (!ttisnil(cb)) {  /* continue loop? */