[PATCH v2 2/6] Add functions to ease using Lua iterators from C

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

Needed for #3276.
---
 src/lua/utils.c | 66 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/lua/utils.h | 28 +++++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/src/lua/utils.c b/src/lua/utils.c
index eefb860ee..4d1eee6ab 100644
--- a/src/lua/utils.c
+++ b/src/lua/utils.c
@@ -969,6 +969,72 @@ 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_fromtable(lua_State *L, int idx)
+{
+	struct luaL_iterator *it = (struct luaL_iterator *) malloc(
+		sizeof(struct luaL_iterator));
+
+	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);
+	lua_call(L, 2, LUA_MULTRET);
+	int nresults = lua_gettop(L) - frame_start;
+	if (nresults == 0) {
+		luaL_error(L, "luaL_iterator_next: gen(param, state) must "
+			      "return at least one result");
+		unreachable();
+		return 0;
+	}
+
+	/* The call above returns nil as the first result. */
+	if (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_free(lua_State *L, struct luaL_iterator *it)
+{
+	luaL_unref(L, LUA_REGISTRYINDEX, it->gen);
+	luaL_unref(L, LUA_REGISTRYINDEX, it->param);
+	luaL_unref(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 bd302d8e9..6ba2e4767 100644
--- a/src/lua/utils.h
+++ b/src/lua/utils.h
@@ -525,6 +525,34 @@ 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 {gen, param, state}.
+ */
+struct luaL_iterator *
+luaL_iterator_new_fromtable(lua_State *L, int idx);
+
+/**
+ * Move iterator to the next value. Push values returned by
+ * gen(param, state) and return its count. Zero means no more
+ * results available.
+ */
+int
+luaL_iterator_next(lua_State *L, struct luaL_iterator *it);
+
+/**
+ * Free all resources held by the iterator.
+ */
+void luaL_iterator_free(lua_State *L, struct luaL_iterator *it);
+
+/* }}} */
+
 int
 tarantool_lua_utils_init(struct lua_State *L);
 
-- 
2.20.1