# HG changeset patch
# User jbe
# Date 1430569915 -7200
# Node ID e7fac0918f9c95b5613498ded77e1259085ee1ce
# Parent  41da87a681d6ae8fb25d28ecdf3f915bd60edbc6
Added new methods drain_call, drain_yield

diff -r 41da87a681d6 -r e7fac0918f9c moonbridge_io.c
--- a/moonbridge_io.c	Sat May 02 01:48:29 2015 +0200
+++ b/moonbridge_io.c	Sat May 02 14:31:55 2015 +0200
@@ -312,13 +312,13 @@
     } else {
       lua_rawseti(L, 5, lua_rawlen(L, 5) + 1);
     }
-    if (remaining >= 0) {
+    if (strcmp(lua_tostring(L, -1), "block") != 0) break;
+    lua_pop(L, 1);
+    if (remaining >= 0 && len) {
       remaining -= len;
       lua_pushinteger(L, remaining);
       lua_replace(L, 3);
     }
-    if (strcmp(lua_tostring(L, -1), "block") != 0) break;
-    lua_pop(L, 1);
     lua_pushvalue(L, 2);
     lua_pushvalue(L, 1);
     lua_pushliteral(L, "r");
@@ -362,6 +362,65 @@
   return lua_gettop(L);
 }
 
+#if LUA_VERSION_NUM >= 503
+static int moonbr_io_drain_cont(lua_State *L, int status, lua_KContext ctx) {
+#else
+static int moonbr_io_drain_cont(lua_State *L) {
+#endif
+  lua_Integer remaining, len;
+  size_t totallen = 0;
+#if !(LUA_VERSION_NUM >= 503)
+  int ctx = 0;
+  lua_getctx(L, &ctx);
+#endif
+  remaining = lua_tointeger(L, 3);
+  while (1) {
+    lua_pushcfunction(L, moonbr_io_drain_nb);
+    lua_pushvalue(L, 1);
+    lua_pushvalue(L, 3);
+    lua_pushvalue(L, 4);
+    lua_call(L, 3, 2);
+    if (lua_isnil(L, -2)) return 2;
+    lua_insert(L, -2);
+    len = lua_tointeger(L, -1);
+    lua_pop(L, 1);
+    totallen += len;
+    if (strcmp(lua_tostring(L, -1), "block") != 0) break;
+    lua_pop(L, 1);
+    if (remaining >= 0 && len) {
+      remaining -= len;
+      lua_pushinteger(L, remaining);
+      lua_replace(L, 3);
+    }
+    lua_pushvalue(L, 2);
+    lua_pushvalue(L, 1);
+    lua_pushliteral(L, "r");
+    lua_callk(L, 2, 0, ctx, moonbr_io_drain_cont);
+  }
+  lua_pushinteger(L, totallen);
+  lua_pushvalue(L, -2);
+  return 2;
+}
+
+static int moonbr_io_drain_call(lua_State *L) {
+#if LUA_VERSION_NUM >= 503
+  return moonbr_io_drain_cont(L, 0, 0);
+#else
+  return moonbr_io_drain_cont(L);
+#endif
+}
+
+static int moonbr_io_drain_yield(lua_State *L) {
+  int args;
+  lua_pushcfunction(L, moonbr_io_drain_call);
+  lua_insert(L, 1);
+  args = lua_gettop(L);
+  lua_pushcfunction(L, moonbr_io_yield);
+  lua_insert(L, 3);
+  lua_callk(L, args, LUA_MULTRET, 0, moonbr_io_cont_returnall);
+  return lua_gettop(L);
+}
+
 static int moonbr_io_write_impl(lua_State *L, int nonblocking, int flush) {
   moonbr_io_handle_t *handle;
   int i, top;
@@ -1217,6 +1276,8 @@
   {"read_yield", moonbr_io_read_yield},
   {"drain", moonbr_io_drain},
   {"drain_nb", moonbr_io_drain_nb},
+  {"drain_call", moonbr_io_drain_call},
+  {"drain_yield", moonbr_io_drain_yield},
   {"write", moonbr_io_write},
   {"write_nb", moonbr_io_write_nb},
   {"flush", moonbr_io_flush},
diff -r 41da87a681d6 -r e7fac0918f9c reference.txt
--- a/reference.txt	Sat May 02 01:48:29 2015 +0200
+++ b/reference.txt	Sat May 02 14:31:55 2015 +0200
@@ -77,6 +77,12 @@
 second return value) are returned.
 
 
+### socket:drain_call(waitfunc, maxlen, terminator)
+
+Same as socket:drain(maxlen, terminator), but calls waitfunc(socket, "r") (in
+an infinite loop) as long as the reading is blocked.
+
+
 ### socket:drain_nb(maxlen, terminator)
 
 Same as socket:drain(maxlen, terminator), but non-blocking. The status code
@@ -87,6 +93,11 @@
 second return value) are returned.
 
 
+### socket:drain_yield(maxlen, terminator)
+
+Alias for socket:drain_call(coroutine.yield, maxlen, terminator)
+
+
 ### socket:finish()
 
 Sends a TCP FIN packet to indicate EOF on write stream. Subsequent reads are