Implemented support for vararg stream read

1 files changed, 25 insertions, 25 deletions

diff --git a/stream.c b/stream.c
index c633862..2c8f2ed 100644
--- a/stream.c
+++ b/stream.c
@@ -107,12 +107,12 @@ int stream_read(lua_State * L)
 		s->in.offset = 0;
 	}
 
-	return stream_readk(L, LUA_OK, (lua_KContext) s); // Intentionally do not remove arguments from the stack.
+	return stream_readk(L, LUA_OK, 2); // Intentionally do not remove arguments from the stack.
 }
 
 static int grow(lua_State *, struct buffer *);
-static int prepare_at_least(lua_State *, struct stream *, const int);
-static int read_more(lua_State *, struct stream *, const int);
+static int prepare_at_least(lua_State *, struct stream *, const int, lua_KContext);
+static int read_more(lua_State *, struct stream *, const int, lua_KContext);
 static int until(struct buffer *, const char *, const int);
 
 static int grow(lua_State * L, struct buffer * b)
@@ -139,7 +139,7 @@ static int grow(lua_State * L, struct buffer * b)
 	return 0;
 }
 
-static int read_more(lua_State * L, struct stream * s, int minimum_length)
+static int read_more(lua_State * L, struct stream * s, int minimum_length, lua_KContext ctx)
 {
 	const int free_space = s->in.allocated + s->in.offset - s->in.length - 1;
 
@@ -161,7 +161,7 @@ static int read_more(lua_State * L, struct stream * s, int minimum_length)
 	{
 		if (EWOULDBLOCK == errno || EAGAIN == errno)
 		{
-			return lua_yieldk(L, 0, (lua_KContext) s, stream_readk);
+			return lua_yieldk(L, 0, ctx, stream_readk);
 		}
 		else
 		{
@@ -175,13 +175,13 @@ static int read_more(lua_State * L, struct stream * s, int minimum_length)
 	return length;
 }
 
-static int prepare_at_least(lua_State * L, struct stream * s, int minimum_length)
+static int prepare_at_least(lua_State * L, struct stream * s, int minimum_length, lua_KContext ctx)
 {
 	const int remaining_bytes = s->in.length - s->in.offset;
 
 	if (remaining_bytes < minimum_length)
 	{
-		return read_more(L, s, minimum_length);
+		return read_more(L, s, minimum_length, ctx);
 	}
 
 	return 0;
@@ -213,28 +213,28 @@ static int until(struct buffer * b, const char * pattern, int pattern_length)
 /// \param Number of the results pushed to the stack
 int stream_readk(lua_State * L, int status, lua_KContext ctx)
 {
-	struct stream * s = (struct stream *) ctx;
-
-	size_t pattern_length;
-	const char * pattern = lua_tolstring(L, 2, &pattern_length);
-
-	// TODO: NULL check
+	struct stream * s = lua_touserdata(L, 1);
 
-	int offset;
-	do
+	for (; ctx <= lua_gettop(L); ++ctx)
 	{
-		prepare_at_least(L, s, (int) pattern_length);
-		offset = until(&s->in, pattern, (int) pattern_length);
-	}
-	while (-1 == offset);
+		size_t pattern_length;
+		const char * pattern = lua_tolstring(L, ctx, &pattern_length);
 
-	lua_pushlstring(L, &s->in.data[s->in.offset], offset - s->in.offset);
-	s->in.offset = offset + 1;
+		int offset;
+		do
+		{
+			prepare_at_least(L, s, (int) pattern_length, ctx);
+			offset = until(&s->in, pattern, (int) pattern_length);
+		}
+		while (-1 == offset);
 
-	for (int n = -lua_gettop(L); n < -1; ++n)
-	{
-		lua_remove(L, n);
+		lua_pushlstring(L, &s->in.data[s->in.offset], offset - s->in.offset);
+		s->in.offset = offset + 1;
+		lua_remove(L, ctx);
+		lua_insert(L, ctx);
 	}
 
-	return 1;
+	lua_remove(L, 1);
+
+	return lua_gettop(L);
 }