`yield' passes its arguments to `resume'

2 files changed, 47 insertions, 34 deletions

diff --git a/lbaselib.c b/lbaselib.c
index 6cd97d95..b06cc2b3 100644
--- a/lbaselib.c
+++ b/lbaselib.c
@@ -415,6 +415,14 @@ static int luaB_tostring (lua_State *L) {
 static int luaB_resume (lua_State *L) {
   lua_State *co = (lua_State *)lua_touserdata(L, lua_upvalueindex(1));
   lua_resume(L, co);
+  return lua_gettop(L);
+}
+
+
+
+static int gc_coroutine (lua_State *L) {
+  lua_State *co = (lua_State *)lua_touserdata(L, 1);
+  lua_closethread(L, co);
   return 0;
 }
 
@@ -422,22 +430,30 @@ static int luaB_resume (lua_State *L) {
 static int luaB_coroutine (lua_State *L) {
   lua_State *NL;
   int ref;
-  luaL_check_type(L, 1, LUA_TFUNCTION);
+  int i;
+  int n = lua_gettop(L);
+  luaL_arg_check(L, lua_isfunction(L, 1) && !lua_iscfunction(L, 1), 1,
+    "Lua function expected");
   NL = lua_newthread(L, 0);
   if (NL == NULL) lua_error(L, "unable to create new thread");
-  /* move function from L to NL */
-  ref = lua_ref(L, 1);
-  lua_getref(NL, ref);
-  lua_unref(L, ref);
-  lua_cobegin(NL, 0);
+  /* move function and arguments from L to NL */
+  for (i=0; i<n; i++) {
+    ref = lua_ref(L, 1);
+    lua_getref(NL, ref);
+    lua_insert(NL, 1);
+    lua_unref(L, ref);
+  }
+  lua_cobegin(NL, n-1);
   lua_newuserdatabox(L, NL);
+  lua_getstr(L, LUA_REGISTRYINDEX, "Coroutine");
+  lua_seteventtable(L, -2);
   lua_pushcclosure(L, luaB_resume, 1);
   return 1;
 }
 
 
 static int luaB_yield (lua_State *L) {
-  return lua_yield(L, 0);
+  return lua_yield(L, lua_gettop(L));
 }
 
 
@@ -683,6 +699,11 @@ LUALIB_API int lua_baselibopen (lua_State *L) {
   lua_newtable(L);
   lua_pushcclosure(L, luaB_require, 1);
   lua_setglobal(L, "require");
+  /* create metatable for coroutines */
+  lua_newtable(L);
+  lua_pushcfunction(L, gc_coroutine);
+  lua_setstr(L, -2, "gc");
+  lua_setstr(L, LUA_REGISTRYINDEX, "Coroutine");
   return 0;
 }
 
diff --git a/ldo.c b/ldo.c
index c161cfe8..e5fbc5f3 100644
--- a/ldo.c
+++ b/ldo.c
@@ -29,9 +29,6 @@
 #include "lzio.h"
 
 
-/* space to handle stack overflow errors */
-#define EXTRA_STACK   (2*LUA_MINSTACK)
-
 
 static void restore_stack_limit (lua_State *L) {
   StkId limit = L->stack+(L->stacksize-EXTRA_STACK)-1;
@@ -40,21 +37,6 @@ static void restore_stack_limit (lua_State *L) {
 }
 
 
-void luaD_init (lua_State *L, int stacksize) {
-  stacksize += EXTRA_STACK;
-  L->stack = luaM_newvector(L, stacksize, TObject);
-  L->stacksize = stacksize;
-  L->top = L->stack + RESERVED_STACK_PREFIX;
-  restore_stack_limit(L);
-  luaM_reallocvector(L, L->base_ci, 0, 20, CallInfo);
-  L->ci = L->base_ci;
-  L->ci->base = L->top;
-  L->ci->savedpc = NULL;
-  L->size_ci = 20;
-  L->end_ci = L->base_ci + L->size_ci;
-}
-
-
 void luaD_stackerror (lua_State *L) {
   if (L->stack_last == L->stack+L->stacksize-1) {
     /* overflow while handling overflow */
@@ -231,15 +213,21 @@ void luaD_call (lua_State *L, StkId func, int nResults) {
 
 
 LUA_API void lua_cobegin (lua_State *L, int nargs) {
-  StkId func;
   lua_lock(L);
-  func = L->top - (nargs+1);  /* coroutine main function */
-  if (luaD_precall(L, func) != NULL)
-    luaD_error(L, "coroutine started with a C function");
+  luaD_precall(L, L->top - (nargs+1));
   lua_unlock(L);
 }
 
 
+static void resume_results (lua_State *L, lua_State *from, int numresults) {
+  while (numresults) {
+    setobj(L->top, from->top - numresults);
+    numresults--;
+    incr_top;
+  }
+}
+
+
 LUA_API void lua_resume (lua_State *L, lua_State *co) {
   StkId firstResult;
   lua_lock(L);
@@ -248,10 +236,13 @@ LUA_API void lua_resume (lua_State *L, lua_State *co) {
   lua_assert(co->errorJmp == NULL);
   co->errorJmp = L->errorJmp;
   firstResult = luaV_execute(co);
-  if (firstResult != NULL)  /* `return'? */
-    luaD_poscall(co, LUA_MULTRET, firstResult);  /* ends this coroutine */
+  if (firstResult != NULL) {  /* `return'? */
+    resume_results(L, co, co->top - firstResult);
+    luaD_poscall(co, 0, firstResult);  /* ends this coroutine */
+  }
   else {  /* `yield' */
     int nresults = GETARG_C(*((co->ci-1)->savedpc - 1)) - 1;
+    resume_results(L, co, co->ci->yield_results);
     luaD_poscall(co, nresults, co->top);  /* complete it */
     if (nresults >= 0) co->top = co->ci->top;
   }
@@ -264,10 +255,11 @@ LUA_API int lua_yield (lua_State *L, int nresults) {
   CallInfo *ci;
   int ibase;
   lua_lock(L);
-  ci = L->ci - 1;  /* call info of calling function */
-  if (ci->pc == NULL)
+  ci = L->ci;
+  if (ci_func(ci-1)->c.isC)
     luaD_error(L, "cannot `yield' a C function");
-  ibase = L->top - ci->base;
+  ci->yield_results = nresults;  /* very dirty trick! */
+  ibase = L->top - (ci-1)->base;
   lua_unlock(L);
   return ibase;
 }