[PATCH v3 2/7] Add functions to ease using Lua iterators from C

[Alexander Turenko <alexander.turenko@tarantool.org>]

Needed for #3276.
---
 src/lua/utils.c                |  92 +++++++++++++++
 src/lua/utils.h                |  37 ++++++
 test/unit/CMakeLists.txt       |   4 +
 test/unit/luaL_iterator.c      | 208 +++++++++++++++++++++++++++++++++
 test/unit/luaL_iterator.result |  89 ++++++++++++++
 5 files changed, 430 insertions(+)
 create mode 100644 test/unit/luaL_iterator.c
 create mode 100644 test/unit/luaL_iterator.result

diff --git a/src/lua/utils.c b/src/lua/utils.c
index 83bc1695e..192912ab8 100644
--- a/src/lua/utils.c
+++ b/src/lua/utils.c
@@ -1065,6 +1065,98 @@ luaT_state(void)
 	return tarantool_L;
 }
 
+/* {{{ Helper functions to interact with a Lua iterator from C */
+
+struct luaL_iterator {
+	int gen;
+	int param;
+	int state;
+};
+
+struct luaL_iterator *
+luaL_iterator_new(lua_State *L, int idx)
+{
+	struct luaL_iterator *it = malloc(sizeof(struct luaL_iterator));
+	if (it == NULL) {
+		diag_set(OutOfMemory, sizeof(struct luaL_iterator),
+			 "malloc", "luaL_iterator");
+		return NULL;
+	}
+
+	if (idx == 0) {
+		/* gen, param, state are on top of a Lua stack. */
+		lua_pushvalue(L, -3); /* Popped by luaL_ref(). */
+		it->gen = luaL_ref(L, LUA_REGISTRYINDEX);
+		lua_pushvalue(L, -2); /* Popped by luaL_ref(). */
+		it->param = luaL_ref(L, LUA_REGISTRYINDEX);
+		lua_pushvalue(L, -1); /* Popped by luaL_ref(). */
+		it->state = luaL_ref(L, LUA_REGISTRYINDEX);
+	} else {
+		/*
+		 * {gen, param, state} table is at idx in a Lua
+		 * stack.
+		 */
+		lua_rawgeti(L, idx, 1); /* Popped by luaL_ref(). */
+		it->gen = luaL_ref(L, LUA_REGISTRYINDEX);
+		lua_rawgeti(L, idx, 2); /* Popped by luaL_ref(). */
+		it->param = luaL_ref(L, LUA_REGISTRYINDEX);
+		lua_rawgeti(L, idx, 3); /* Popped by luaL_ref(). */
+		it->state = luaL_ref(L, LUA_REGISTRYINDEX);
+	}
+
+	return it;
+}
+
+int
+luaL_iterator_next(lua_State *L, struct luaL_iterator *it)
+{
+	int frame_start = lua_gettop(L);
+
+	/* Call gen(param, state). */
+	lua_rawgeti(L, LUA_REGISTRYINDEX, it->gen);
+	lua_rawgeti(L, LUA_REGISTRYINDEX, it->param);
+	lua_rawgeti(L, LUA_REGISTRYINDEX, it->state);
+	if (luaT_call(L, 2, LUA_MULTRET) != 0) {
+		/*
+		 * Pop garbage from the call (a gen function
+		 * likely will not leave the stack even when raise
+		 * an error), pop a returned error.
+		 */
+		lua_settop(L, frame_start);
+		return -1;
+	}
+	int nresults = lua_gettop(L) - frame_start;
+
+	/*
+	 * gen() function can either return nil when the iterator
+	 * ends or return zero count of values.
+	 *
+	 * In LuaJIT pairs() returns nil, but ipairs() returns
+	 * nothing when ends.
+	 */
+	if (nresults == 0 || lua_isnil(L, frame_start + 1)) {
+		lua_settop(L, frame_start);
+		return 0;
+	}
+
+	/* Save the first result to it->state. */
+	luaL_unref(L, LUA_REGISTRYINDEX, it->state);
+	lua_pushvalue(L, frame_start + 1); /* Popped by luaL_ref(). */
+	it->state = luaL_ref(L, LUA_REGISTRYINDEX);
+
+	return nresults;
+}
+
+void luaL_iterator_delete(struct luaL_iterator *it)
+{
+	luaL_unref(tarantool_L, LUA_REGISTRYINDEX, it->gen);
+	luaL_unref(tarantool_L, LUA_REGISTRYINDEX, it->param);
+	luaL_unref(tarantool_L, LUA_REGISTRYINDEX, it->state);
+	free(it);
+}
+
+/* }}} */
+
 int
 tarantool_lua_utils_init(struct lua_State *L)
 {
diff --git a/src/lua/utils.h b/src/lua/utils.h
index cf8c8d09d..a22492227 100644
--- a/src/lua/utils.h
+++ b/src/lua/utils.h
@@ -536,6 +536,43 @@ luaL_checkfinite(struct lua_State *L, struct luaL_serializer *cfg,
 		luaL_error(L, "number must not be NaN or Inf");
 }
 
+/* {{{ Helper functions to interact with a Lua iterator from C */
+
+/**
+ * Holds iterator state (references to Lua objects).
+ */
+struct luaL_iterator;
+
+/**
+ * Create a Lua iterator from a gen, param, state triplet.
+ *
+ * If idx == 0, then three top stack values are used as the
+ * triplet. Note: they are not popped.
+ *
+ * Otherwise idx is index on Lua stack points to a
+ * {gen, param, state} table.
+ */
+struct luaL_iterator *
+luaL_iterator_new(lua_State *L, int idx);
+
+/**
+ * Move iterator to the next value. Push values returned by
+ * gen(param, state).
+ *
+ * Return count of pushed values. Zero means no more results
+ * available. In case of a Lua error in a gen function return -1
+ * and set a diag.
+ */
+int
+luaL_iterator_next(lua_State *L, struct luaL_iterator *it);
+
+/**
+ * Free all resources held by the iterator.
+ */
+void luaL_iterator_delete(struct luaL_iterator *it);
+
+/* }}} */
+
 int
 tarantool_lua_utils_init(struct lua_State *L);
 
diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt
index 2bcb6e0a8..64739ab09 100644
--- a/test/unit/CMakeLists.txt
+++ b/test/unit/CMakeLists.txt
@@ -144,6 +144,10 @@ add_executable(luaT_tuple_new.test luaT_tuple_new.c)
 target_link_libraries(luaT_tuple_new.test unit box server core misc
     ${CURL_LIBRARIES} ${LIBYAML_LIBRARIES} ${READLINE_LIBRARIES}
     ${ICU_LIBRARIES} ${LUAJIT_LIBRARIES})
+add_executable(luaL_iterator.test luaL_iterator.c)
+target_link_libraries(luaL_iterator.test unit server coll core misc
+    ${CURL_LIBRARIES} ${LIBYAML_LIBRARIES} ${READLINE_LIBRARIES}
+    ${ICU_LIBRARIES} ${LUAJIT_LIBRARIES} dl)
 
 add_executable(say.test say.c)
 target_link_libraries(say.test core unit)
diff --git a/test/unit/luaL_iterator.c b/test/unit/luaL_iterator.c
new file mode 100644
index 000000000..8d25a0062
--- /dev/null
+++ b/test/unit/luaL_iterator.c
@@ -0,0 +1,208 @@
+#include <lua.h>       /* lua_*() */
+#include <lauxlib.h>   /* luaL_*() */
+#include <lualib.h>    /* luaL_openlibs() */
+#include "unit.h"      /* plan, header, footer, is */
+#include "memory.h"    /* memory_init() */
+#include "fiber.h"     /* fiber_init() */
+#include "diag.h"      /* struct error, diag_*() */
+#include "exception.h" /* type_LuajitError */
+#include "lua/utils.h" /* luaL_iterator_*() */
+#include "lua/error.h" /* tarantool_lua_error_init() */
+
+extern char fun_lua[];
+
+int
+main()
+{
+	struct {
+		/* A string to output with a test case. */
+		const char *description;
+		/* A string with Lua code to push an iterator. */
+		const char *init;
+		/*
+		 * How much values are pushed by the Lua code
+		 * above.
+		 */
+		int init_retvals;
+		/*
+		 * Start values from this number to distinguish
+		 * them from its ordinal.
+		 */
+		int first_value;
+		/*
+		 * Lua stack index where {gen, param, state} is
+		 * placed or zero.
+		 */
+		int idx;
+		/* How much values are in the iterator. */
+		int iterations;
+		/* Expected error (if any). */
+		const char *exp_err;
+	} cases[] = {
+		{
+			.description = "pairs, zero idx",
+			.init = "return pairs({42})",
+			.init_retvals = 3,
+			.first_value = 42,
+			.idx = 0,
+			.iterations = 1,
+			.exp_err = NULL,
+		},
+		{
+			.description = "ipairs, zero idx",
+			.init = "return ipairs({42, 43, 44})",
+			.init_retvals = 3,
+			.first_value = 42,
+			.idx = 0,
+			.iterations = 3,
+			.exp_err = NULL,
+		},
+		{
+			.description = "luafun iterator, zero idx",
+			.init = "return fun.wrap(ipairs({42, 43, 44}))",
+			.init_retvals = 3,
+			.first_value = 42,
+			.idx = 0,
+			.iterations = 3,
+			.exp_err = NULL,
+		},
+		{
+			.description = "pairs, from table",
+			.init = "return {pairs({42})}",
+			.init_retvals = 1,
+			.first_value = 42,
+			.idx = -1,
+			.iterations = 1,
+			.exp_err = NULL,
+		},
+		{
+			.description = "ipairs, from table",
+			.init = "return {ipairs({42, 43, 44})}",
+			.init_retvals = 1,
+			.first_value = 42,
+			.idx = -1,
+			.iterations = 3,
+			.exp_err = NULL,
+		},
+		{
+			.description = "luafun iterator, from table",
+			.init = "return {fun.wrap(ipairs({42, 43, 44}))}",
+			.init_retvals = 1,
+			.first_value = 42,
+			.idx = -1,
+			.iterations = 3,
+			.exp_err = NULL,
+		},
+		{
+			.description = "lua error",
+			.init = "return error, 'I am the error', 0",
+			.init_retvals = 3,
+			.first_value = 0,
+			.idx = 0,
+			.iterations = 0,
+			.exp_err = "I am the error",
+		},
+	};
+
+	int cases_cnt = (int) (sizeof(cases) / sizeof(cases[0]));
+	/*
+	 * * 4 checks per iteration.
+	 * * 3 checks of a stack size.
+	 * * 1 check that values ends (for success cases).
+	 * * 1 check for an iterator error (for error cases).
+	 * * 1 check for an error type (for error cases).
+	 * * 1 check for an error message (for error cases).
+	 */
+	int planned = 0;
+	for (int i = 0; i < cases_cnt; ++i) {
+		planned += cases[i].iterations * 4 + 4;
+		if (cases[i].exp_err != NULL)
+			planned += 2;
+	}
+
+	plan(planned);
+	header();
+
+	struct lua_State *L = luaL_newstate();
+	luaL_openlibs(L);
+	tarantool_L = L;
+
+	memory_init();
+	fiber_init(fiber_c_invoke);
+	tarantool_lua_error_init(L);
+
+	/*
+	 * Check that everything works fine in a thread (a fiber)
+	 * other then the main one.
+	 */
+	L = lua_newthread(L);
+
+	/*
+	 * Expose luafun.
+	 *
+	 * Don't register it in package.loaded for simplicity.
+	 */
+	luaL_loadstring(L, fun_lua);
+	lua_call(L, 0, 1);
+	lua_setglobal(L, "fun");
+
+	for (int i = 0; i < cases_cnt; ++i) {
+		const char *description = cases[i].description;
+		int top = lua_gettop(L);
+
+		/* Push an iterator to the Lua stack. */
+		luaL_loadstring(L, cases[i].init);
+		lua_call(L, 0, cases[i].init_retvals);
+
+		/* Create the luaL_iterator structure. */
+		struct luaL_iterator *it = luaL_iterator_new(L, cases[i].idx);
+		lua_pop(L, cases[i].init_retvals);
+
+		/* Check stack size. */
+		is(lua_gettop(L) - top, 0, "%s: stack size", description);
+
+		/* Iterate over values and check them. */
+		for (int j = 0; j < cases[i].iterations; ++j) {
+			int top = lua_gettop(L);
+			int rc = luaL_iterator_next(L, it);
+			is(rc, 2, "%s: iter %d: gen() retval count",
+			   description, j);
+			is(luaL_checkinteger(L, -2), j + 1,
+			   "%s: iter %d: gen() 1st retval",
+			   description, j);
+			is(luaL_checkinteger(L, -1), j + cases[i].first_value,
+			   "%s: iter %d: gen() 2nd retval",
+			   description, j);
+			lua_pop(L, 2);
+			is(lua_gettop(L) - top, 0, "%s: iter: %d: stack size",
+			   description, j);
+		}
+
+		if (cases[i].exp_err == NULL) {
+			/* Check the iterator ends when expected. */
+			int rc = luaL_iterator_next(L, it);
+			is(rc, 0, "%s: iterator ends", description);
+		} else {
+			/* Check expected error. */
+			int rc = luaL_iterator_next(L, it);
+			is(rc, -1, "%s: iterator error", description);
+			struct error *e = diag_last_error(diag_get());
+			is(e->type, &type_LuajitError, "%s: check error type",
+			   description);
+			ok(!strcmp(e->errmsg, cases[i].exp_err),
+			   "%s: check error message", description);
+		}
+
+		/* Check stack size. */
+		is(lua_gettop(L) - top, 0, "%s: stack size", description);
+
+		/* Free the luaL_iterator structure. */
+		luaL_iterator_delete(it);
+
+		/* Check stack size. */
+		is(lua_gettop(L) - top, 0, "%s: stack size", description);
+	}
+
+	footer();
+	return check_plan();
+}
diff --git a/test/unit/luaL_iterator.result b/test/unit/luaL_iterator.result
new file mode 100644
index 000000000..2472eedcf
--- /dev/null
+++ b/test/unit/luaL_iterator.result
@@ -0,0 +1,89 @@
+1..86
+	*** main ***
+ok 1 - pairs, zero idx: stack size
+ok 2 - pairs, zero idx: iter 0: gen() retval count
+ok 3 - pairs, zero idx: iter 0: gen() 1st retval
+ok 4 - pairs, zero idx: iter 0: gen() 2nd retval
+ok 5 - pairs, zero idx: iter: 0: stack size
+ok 6 - pairs, zero idx: iterator ends
+ok 7 - pairs, zero idx: stack size
+ok 8 - pairs, zero idx: stack size
+ok 9 - ipairs, zero idx: stack size
+ok 10 - ipairs, zero idx: iter 0: gen() retval count
+ok 11 - ipairs, zero idx: iter 0: gen() 1st retval
+ok 12 - ipairs, zero idx: iter 0: gen() 2nd retval
+ok 13 - ipairs, zero idx: iter: 0: stack size
+ok 14 - ipairs, zero idx: iter 1: gen() retval count
+ok 15 - ipairs, zero idx: iter 1: gen() 1st retval
+ok 16 - ipairs, zero idx: iter 1: gen() 2nd retval
+ok 17 - ipairs, zero idx: iter: 1: stack size
+ok 18 - ipairs, zero idx: iter 2: gen() retval count
+ok 19 - ipairs, zero idx: iter 2: gen() 1st retval
+ok 20 - ipairs, zero idx: iter 2: gen() 2nd retval
+ok 21 - ipairs, zero idx: iter: 2: stack size
+ok 22 - ipairs, zero idx: iterator ends
+ok 23 - ipairs, zero idx: stack size
+ok 24 - ipairs, zero idx: stack size
+ok 25 - luafun iterator, zero idx: stack size
+ok 26 - luafun iterator, zero idx: iter 0: gen() retval count
+ok 27 - luafun iterator, zero idx: iter 0: gen() 1st retval
+ok 28 - luafun iterator, zero idx: iter 0: gen() 2nd retval
+ok 29 - luafun iterator, zero idx: iter: 0: stack size
+ok 30 - luafun iterator, zero idx: iter 1: gen() retval count
+ok 31 - luafun iterator, zero idx: iter 1: gen() 1st retval
+ok 32 - luafun iterator, zero idx: iter 1: gen() 2nd retval
+ok 33 - luafun iterator, zero idx: iter: 1: stack size
+ok 34 - luafun iterator, zero idx: iter 2: gen() retval count
+ok 35 - luafun iterator, zero idx: iter 2: gen() 1st retval
+ok 36 - luafun iterator, zero idx: iter 2: gen() 2nd retval
+ok 37 - luafun iterator, zero idx: iter: 2: stack size
+ok 38 - luafun iterator, zero idx: iterator ends
+ok 39 - luafun iterator, zero idx: stack size
+ok 40 - luafun iterator, zero idx: stack size
+ok 41 - pairs, from table: stack size
+ok 42 - pairs, from table: iter 0: gen() retval count
+ok 43 - pairs, from table: iter 0: gen() 1st retval
+ok 44 - pairs, from table: iter 0: gen() 2nd retval
+ok 45 - pairs, from table: iter: 0: stack size
+ok 46 - pairs, from table: iterator ends
+ok 47 - pairs, from table: stack size
+ok 48 - pairs, from table: stack size
+ok 49 - ipairs, from table: stack size
+ok 50 - ipairs, from table: iter 0: gen() retval count
+ok 51 - ipairs, from table: iter 0: gen() 1st retval
+ok 52 - ipairs, from table: iter 0: gen() 2nd retval
+ok 53 - ipairs, from table: iter: 0: stack size
+ok 54 - ipairs, from table: iter 1: gen() retval count
+ok 55 - ipairs, from table: iter 1: gen() 1st retval
+ok 56 - ipairs, from table: iter 1: gen() 2nd retval
+ok 57 - ipairs, from table: iter: 1: stack size
+ok 58 - ipairs, from table: iter 2: gen() retval count
+ok 59 - ipairs, from table: iter 2: gen() 1st retval
+ok 60 - ipairs, from table: iter 2: gen() 2nd retval
+ok 61 - ipairs, from table: iter: 2: stack size
+ok 62 - ipairs, from table: iterator ends
+ok 63 - ipairs, from table: stack size
+ok 64 - ipairs, from table: stack size
+ok 65 - luafun iterator, from table: stack size
+ok 66 - luafun iterator, from table: iter 0: gen() retval count
+ok 67 - luafun iterator, from table: iter 0: gen() 1st retval
+ok 68 - luafun iterator, from table: iter 0: gen() 2nd retval
+ok 69 - luafun iterator, from table: iter: 0: stack size
+ok 70 - luafun iterator, from table: iter 1: gen() retval count
+ok 71 - luafun iterator, from table: iter 1: gen() 1st retval
+ok 72 - luafun iterator, from table: iter 1: gen() 2nd retval
+ok 73 - luafun iterator, from table: iter: 1: stack size
+ok 74 - luafun iterator, from table: iter 2: gen() retval count
+ok 75 - luafun iterator, from table: iter 2: gen() 1st retval
+ok 76 - luafun iterator, from table: iter 2: gen() 2nd retval
+ok 77 - luafun iterator, from table: iter: 2: stack size
+ok 78 - luafun iterator, from table: iterator ends
+ok 79 - luafun iterator, from table: stack size
+ok 80 - luafun iterator, from table: stack size
+ok 81 - lua error: stack size
+ok 82 - lua error: iterator error
+ok 83 - lua error: check error type
+ok 84 - lua error: check error message
+ok 85 - lua error: stack size
+ok 86 - lua error: stack size
+	*** main: done ***
-- 
2.20.1