4 files changed, 23 insertions, 13 deletions

diff --git a/lapi.c b/lapi.c
index 9201aac4..b82706fd 100644
--- a/lapi.c
+++ b/lapi.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lapi.c,v 2.162 2012/05/22 17:50:39 roberto Exp roberto $
+** $Id: lapi.c,v 2.163 2012/05/23 15:42:27 roberto Exp roberto $
 ** Lua API
 ** See Copyright Notice in lua.h
 */
@@ -950,7 +950,7 @@ LUA_API int lua_pcallk (lua_State *L, int nargs, int nresults, int errfunc,
     ci->u.c.k = k;  /* save continuation */
     ci->u.c.ctx = ctx;  /* save context */
     /* save information for error recovery */
-    ci->u.c.extra = savestack(L, c.func);
+    ci->extra = savestack(L, c.func);
     ci->u.c.old_allowhook = L->allowhook;
     ci->u.c.old_errfunc = L->errfunc;
     L->errfunc = func;
diff --git a/ldo.c b/ldo.c
index f3121fd0..9e1a5b00 100644
--- a/ldo.c
+++ b/ldo.c
@@ -1,5 +1,5 @@
 /*
-** $Id: ldo.c,v 2.103 2012/04/26 20:41:18 roberto Exp roberto $
+** $Id: ldo.c,v 2.104 2012/05/08 13:53:33 roberto Exp roberto $
 ** Stack and Call structure of Lua
 ** See Copyright Notice in lua.h
 */
@@ -451,7 +451,7 @@ static int recover (lua_State *L, int status) {
   CallInfo *ci = findpcall(L);
   if (ci == NULL) return 0;  /* no recovery point */
   /* "finish" luaD_pcall */
-  oldtop = restorestack(L, ci->u.c.extra);
+  oldtop = restorestack(L, ci->extra);
   luaF_close(L, oldtop);
   seterrorobj(L, status, oldtop);
   L->ci = ci;
@@ -498,10 +498,10 @@ static void resume (lua_State *L, void *ud) {
     resume_error(L, "cannot resume dead coroutine", firstArg);
   else {  /* resuming from previous yield */
     L->status = LUA_OK;
+    ci->func = restorestack(L, ci->extra);
     if (isLua(ci))  /* yielded inside a hook? */
       luaV_execute(L);  /* just continue running Lua code */
     else {  /* 'common' yield */
-      ci->func = restorestack(L, ci->u.c.extra);
       if (ci->u.c.k != NULL) {  /* does it have a continuation? */
         int n;
         ci->u.c.status = LUA_YIELD;  /* 'default' status */
@@ -563,13 +563,13 @@ LUA_API int lua_yieldk (lua_State *L, int nresults, int ctx, lua_CFunction k) {
       luaG_runerror(L, "attempt to yield from outside a coroutine");
   }
   L->status = LUA_YIELD;
+  ci->extra = savestack(L, ci->func);  /* save current 'func' */
   if (isLua(ci)) {  /* inside a hook? */
     api_check(L, k == NULL, "hooks cannot continue after yielding");
   }
   else {
     if ((ci->u.c.k = k) != NULL)  /* is there a continuation? */
       ci->u.c.ctx = ctx;  /* save context */
-    ci->u.c.extra = savestack(L, ci->func);  /* save current 'func' */
     ci->func = L->top - nresults - 1;  /* protect stack below results */
     luaD_throw(L, LUA_YIELD);
   }
diff --git a/lstate.h b/lstate.h
index 665cf160..5c6bd67b 100644
--- a/lstate.h
+++ b/lstate.h
@@ -1,5 +1,5 @@
 /*
-** $Id: lstate.h,v 2.79 2012/05/22 17:32:25 roberto Exp roberto $
+** $Id: lstate.h,v 2.80 2012/05/22 17:50:39 roberto Exp roberto $
 ** Global State
 ** See Copyright Notice in lua.h
 */
@@ -72,6 +72,7 @@ typedef struct CallInfo {
   struct CallInfo *previous, *next;  /* dynamic call link */
   short nresults;  /* expected number of results from this function */
   lu_byte callstatus;
+  ptrdiff_t extra;
   union {
     struct {  /* only for Lua functions */
       StkId base;  /* base for this function */
@@ -81,7 +82,6 @@ typedef struct CallInfo {
       int ctx;  /* context info. in case of yields */
       lua_CFunction k;  /* continuation in case of yields */
       ptrdiff_t old_errfunc;
-      ptrdiff_t extra;
       lu_byte old_allowhook;
       lu_byte status;
     } c;
@@ -100,6 +100,7 @@ typedef struct CallInfo {
 #define CIST_YPCALL     (1<<4)  /* call is a yieldable protected call */
 #define CIST_STAT       (1<<5)  /* call has an error status (pcall) */
 #define CIST_TAIL       (1<<6)  /* call was tail called */
+#define CIST_HOOKYIELD  (1<<7)  /* last hook called yielded */
 
 
 #define isLua(ci)       ((ci)->callstatus & CIST_LUA)
diff --git a/lvm.c b/lvm.c
index cd7642eb..2d27d512 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1,5 +1,5 @@
 /*
-** $Id: lvm.c,v 2.150 2012/05/08 13:53:33 roberto Exp roberto $
+** $Id: lvm.c,v 2.151 2012/05/14 17:50:49 roberto Exp roberto $
 ** Lua virtual machine
 ** See Copyright Notice in lua.h
 */
@@ -60,10 +60,15 @@ int luaV_tostring (lua_State *L, StkId obj) {
 static void traceexec (lua_State *L) {
   CallInfo *ci = L->ci;
   lu_byte mask = L->hookmask;
-  if ((mask & LUA_MASKCOUNT) && L->hookcount == 0) {
-    resethookcount(L);
-    luaD_hook(L, LUA_HOOKCOUNT, -1);
+  int counthook = ((mask & LUA_MASKCOUNT) && L->hookcount == 0);
+  if (counthook)
+    resethookcount(L);  /* reset count */
+  if (ci->callstatus & CIST_HOOKYIELD) {  /* called hook last time? */
+    ci->callstatus &= ~CIST_HOOKYIELD;  /* erase mark */
+    return;  /* do not call hook again (VM yielded, so it did not move) */
   }
+  if (counthook)
+    luaD_hook(L, LUA_HOOKCOUNT, -1);  /* call count hook */
   if (mask & LUA_MASKLINE) {
     Proto *p = ci_func(ci)->p;
     int npc = pcRel(ci->u.l.savedpc, p);
@@ -71,11 +76,15 @@ static void traceexec (lua_State *L) {
     if (npc == 0 ||  /* call linehook when enter a new function, */
         ci->u.l.savedpc <= L->oldpc ||  /* when jump back (loop), or when */
         newline != getfuncline(p, pcRel(L->oldpc, p)))  /* enter a new line */
-      luaD_hook(L, LUA_HOOKLINE, newline);
+      luaD_hook(L, LUA_HOOKLINE, newline);  /* call line hook */
   }
   L->oldpc = ci->u.l.savedpc;
   if (L->status == LUA_YIELD) {  /* did hook yield? */
+    if (counthook)
+      L->hookcount = 1;  /* undo decrement to zero */
     ci->u.l.savedpc--;  /* undo increment (resume will increment it again) */
+    ci->callstatus |= CIST_HOOKYIELD;  /* mark that it yieled */
+    ci->func = L->top - 1;  /* protect stack below results */
     luaD_throw(L, LUA_YIELD);
   }
 }