2 files changed, 76 insertions, 12 deletions
diff --git a/src/BLua.cc b/src/BLua.cc
index 1ce20e6..62e4010 100644
--- a/src/BLua.cc
+++ b/src/BLua.cc
@@ -908,3 +908,54 @@ void Balau::LuaHelpersBase::validate(const lua_functypes_t & entry, bool method,
         }
     }
 }
+
+static char s_signature = 'B';
+
+void Balau::LuaHelpersBase::pushContext(Lua & L, std::function<int(Lua & L)> context, Events::BaseEvent * evt) {
+    L.push(new LuaHelpersBase(context, evt));
+    L.push((void *) &s_signature);
+}
+
+bool Balau::LuaHelpersBase::resume(Lua & L) {
+    if (L.gettop() < 2)
+        return false;
+
+    if (!L.islightuserdata())
+        return false;
+
+    char * sig = (char *) L.touserdata();
+    if (sig != &s_signature)
+        return false;
+
+    L.remove(-1);
+
+    LuaHelpersBase * b = (LuaHelpersBase *) L.touserdata();
+    L.remove(-1);
+
+    int r = b->m_context(L);
+
+    delete b;
+
+    if (r < 0)
+        return true;
+
+    L.resume(r);
+
+    return true;
+}
+
+Balau::Events::BaseEvent * Balau::LuaHelpersBase::getEvent(Lua & L) {
+    if (L.gettop() < 2)
+        return NULL;
+
+    if (!L.islightuserdata())
+        return NULL;
+
+    char * sig = (char *) L.touserdata();
+    if (sig != &s_signature)
+        return NULL;
+
+    LuaHelpersBase * b = (LuaHelpersBase *) L.touserdata(-2);
+
+    return b->m_evt;
+}
diff --git a/src/LuaTask.cc b/src/LuaTask.cc
index fc92cdf..f94a2b4 100644
--- a/src/LuaTask.cc
+++ b/src/LuaTask.cc
@@ -57,19 +57,32 @@ void Balau::LuaMainTask::Do() {
 }
 
 void Balau::LuaTask::Do() {
-    try {
-        m_cell->run(L);
-    }
-    catch (GeneralException e) {
-        m_cell->m_exception = new GeneralException(e);
-    }
-    catch (...) {
-        m_cell->setError();
+    while(true) {
+        try {
+            if (L.yielded())
+                LuaHelpersBase::resume(L);
+            else
+                m_cell->run(L);
+        }
+        catch (GeneralException e) {
+            m_cell->m_exception = new GeneralException(e);
+        }
+        catch (...) {
+            m_cell->setError();
+        }
+        if (L.yielded()) {
+            Events::BaseEvent * evt = LuaHelpersBase::getEvent(L);
+            IAssert(evt, "We need an event for now here.");
+            waitFor(evt);
+            yield();
+            continue;
+        }
+        if (m_cell->m_detached)
+            delete m_cell;
+        else
+            m_cell->m_event.trigger();
+        break;
     }
-    if (m_cell->m_detached)
-        delete m_cell;
-    else
-        m_cell->m_event.trigger();
 }
 
 void Balau::LuaExecCell::exec(LuaMainTask * mainTask) {