io: multiple arguments for file:write() and stream:write()

2 files changed, 32 insertions, 11 deletions

diff --git a/lem/io/file.c b/lem/io/file.c
index eae8b14..9c5eaef 100644
--- a/lem/io/file.c
+++ b/lem/io/file.c
@@ -27,6 +27,7 @@ struct file {
 		struct {
 			const char *str;
 			size_t len;
+			int idx;
 		} write;
 		struct {
 			off_t val;
@@ -220,14 +221,21 @@ file_write_reap(struct lem_async *a)
 	struct file *f = (struct file *)a;
 	lua_State *T = f->a.T;
 
-	f->a.T = NULL;
 	if (f->ret) {
+		f->a.T = NULL;
 		lem_queue(T, io_strerror(T, f->ret));
 		return;
 	}
 
-	lua_pushboolean(T, 1);
-	lem_queue(T, 1);
+	if (f->write.idx == lua_gettop(T)) {
+		f->a.T = NULL;
+		lua_pushboolean(T, 1);
+		lem_queue(T, 1);
+		return;
+	}
+
+	f->write.str = lua_tolstring(T, ++f->write.idx, &f->write.len);
+	lem_async_put(&f->a);
 }
 
 static int
@@ -236,9 +244,14 @@ file_write(lua_State *T)
 	struct file *f;
 	const char *str;
 	size_t len;
+	int i;
+	int top;
 
 	luaL_checktype(T, 1, LUA_TUSERDATA);
 	str = luaL_checklstring(T, 2, &len);
+	top = lua_gettop(T);
+	for (i = 3; i <= top; i++)
+		(void)luaL_checkstring(T, i);
 
 	f = lua_touserdata(T, 1);
 	if (f->fd < 0)
@@ -248,10 +261,10 @@ file_write(lua_State *T)
 
 	f->write.str = str;
 	f->write.len = len;
+	f->write.idx = 2;
 	lem_async_do(&f->a, T, file_write_work, file_write_reap);
 
-	lua_settop(T, 2);
-	return lua_yield(T, 2);
+	return lua_yield(T, top);
 }
 
 /*
diff --git a/lem/io/stream.c b/lem/io/stream.c
index 4861860..7acccdc 100644
--- a/lem/io/stream.c
+++ b/lem/io/stream.c
@@ -21,6 +21,7 @@ struct stream {
 	struct ev_io w;
 	const char *out;
 	size_t out_len;
+	int idx;
 	struct lem_parser *p;
 	struct lem_inputbuf buf;
 };
@@ -197,12 +198,15 @@ stream__write(lua_State *T, struct stream *s)
 	int err;
 
 	while ((bytes = write(s->w.fd, s->out, s->out_len)) > 0) {
+		s->out += bytes;
 		s->out_len -= bytes;
 		if (s->out_len == 0) {
-			lua_pushboolean(T, 1);
-			return 1;
+			if (s->idx == lua_gettop(T)) {
+				lua_pushboolean(T, 1);
+				return 1;
+			}
+			s->out = lua_tolstring(T, ++s->idx, &s->out_len);
 		}
-		s->out += bytes;
 	}
 	err = errno;
 
@@ -242,10 +246,15 @@ stream_write(lua_State *T)
 	struct stream *s;
 	const char *out;
 	size_t out_len;
+	int i;
+	int top;
 	int ret;
 
 	luaL_checktype(T, 1, LUA_TUSERDATA);
 	out = luaL_checklstring(T, 2, &out_len);
+	top = lua_gettop(T);
+	for (i = 3; i <= top; i++)
+		(void)luaL_checkstring(T, i);
 
 	s = lua_touserdata(T, 1);
 	if (s->w.fd < 0)
@@ -253,10 +262,9 @@ stream_write(lua_State *T)
 	if (s->w.data != NULL)
 		return io_busy(T);
 
-	lua_settop(T, 2);
-
 	s->out = out;
 	s->out_len = out_len;
+	s->idx = 2;
 	ret = stream__write(T, s);
 	if (ret > 0)
 		return ret;
@@ -264,7 +272,7 @@ stream_write(lua_State *T)
 	s->w.data = T;
 	s->w.cb = stream_write_cb;
 	ev_io_start(LEM_ &s->w);
-	return lua_yield(T, 2);
+	return lua_yield(T, top);
 }
 
 #ifndef TCP_CORK