From c44a935b04cec66913f23e040f26bcc3e0739a04 Mon Sep 17 00:00:00 2001
From: Yukita Mayako <catgirl@goddess.moe>
Date: Sat, 27 May 2017 14:03:27 -0500
Subject: [PATCH] Rewrote `thinkers.iterate` to handle invalid pointers
 elegantly.

---
 src/lua_thinkerlib.c | 134 +++++++++++++++++++++++++++----------------
 1 file changed, 84 insertions(+), 50 deletions(-)

diff --git a/src/lua_thinkerlib.c b/src/lua_thinkerlib.c
index d5251425..aaa8435e 100644
--- a/src/lua_thinkerlib.c
+++ b/src/lua_thinkerlib.c
@@ -16,84 +16,118 @@
 #include "lua_script.h"
 #include "lua_libs.h"
 
+#define META_ITERATIONSTATE "iteration state"
+
 static const char *const iter_opt[] = {
 	"all",
 	"mobj",
 	NULL};
 
+static const actionf_p1 iter_funcs[] = {
+	NULL,
+	(actionf_p1)P_MobjThinker
+};
+
+struct iterationState {
+	actionf_p1 filter;
+	int next;
+};
+
+static int iterationState_gc(lua_State *L)
+{
+	struct iterationState *it = luaL_checkudata(L, -1, META_ITERATIONSTATE);
+	if (it->next != LUA_REFNIL)
+	{
+		luaL_unref(L, LUA_REGISTRYINDEX, it->next);
+		it->next = LUA_REFNIL;
+	}
+	return 0;
+}
+
+#define push_thinker(th) {\
+	if ((th)->function.acp1 == (actionf_p1)P_MobjThinker) \
+		LUA_PushUserdata(L, (th), META_MOBJ); \
+	else \
+		lua_pushlightuserdata(L, (th)); \
+}
+
 static int lib_iterateThinkers(lua_State *L)
 {
-	int state = luaL_checkoption(L, 1, "mobj", iter_opt);
-
-	thinker_t *th = NULL;
-	actionf_p1 searchFunc;
-	const char *searchMeta;
-
+	thinker_t *th = NULL, *next = NULL;
+	struct iterationState *it = luaL_checkudata(L, 1, META_ITERATIONSTATE);
 	lua_settop(L, 2);
-	lua_remove(L, 1); // remove state now.
 
-	switch(state)
+	if (lua_isnil(L, 2))
+		th = &thinkercap;
+	else if (lua_isuserdata(L, 2))
 	{
-		case 0:
-			searchFunc = NULL;
-			searchMeta = NULL;
-			break;
-		case 1:
-		default:
-			searchFunc = (actionf_p1)P_MobjThinker;
-			searchMeta = META_MOBJ;
-			break;
+		if (lua_islightuserdata(L, 2))
+			th = lua_touserdata(L, 2);
+		else
+		{
+			th = *(thinker_t **)lua_touserdata(L, -1);
+			if (!th)
+			{
+				if (it->next == LUA_REFNIL)
+					return 0;
+
+				lua_rawgeti(L, LUA_REGISTRYINDEX, it->next);
+				if (lua_islightuserdata(L, -1))
+					next = lua_touserdata(L, -1);
+				else
+					next = *(thinker_t **)lua_touserdata(L, -1);
+			}
+		}
 	}
 
-	if (!lua_isnil(L, 1)) {
-		if (lua_islightuserdata(L, 1))
-			th = (thinker_t *)lua_touserdata(L, 1);
-		else if (searchMeta)
-			th = *((thinker_t **)luaL_checkudata(L, 1, searchMeta));
-		else
-			th = *((thinker_t **)lua_touserdata(L, 1));
-	} else
-		th = &thinkercap;
+	luaL_unref(L, LUA_REGISTRYINDEX, it->next);
+	it->next = LUA_REFNIL;
 
-	if (!th) // something got our userdata invalidated!
-		return 0;
+	if (th && !next)
+		next = th->next;
+	if (!next)
+		return luaL_error(L, "next thinker invalidated during iteration");
 
-	if (searchFunc == NULL)
-	{
-		if ((th = th->next) != &thinkercap)
+	for (; next != &thinkercap; next = next->next)
+		if (!it->filter || next->function.acp1 == it->filter)
 		{
-			if (th->function.acp1 == (actionf_p1)P_MobjThinker)
-				LUA_PushUserdata(L, th, META_MOBJ);
-			else
-				lua_pushlightuserdata(L, th);
+			push_thinker(next);
+			if (next->next != &thinkercap)
+			{
+				push_thinker(next->next);
+				it->next = luaL_ref(L, LUA_REGISTRYINDEX);
+			}
 			return 1;
 		}
-		return 0;
-	}
-
-	for (th = th->next; th != &thinkercap; th = th->next)
-	{
-		if (th->function.acp1 != searchFunc)
-			continue;
-
-		LUA_PushUserdata(L, th, searchMeta);
-		return 1;
-	}
 	return 0;
 }
 
 static int lib_startIterate(lua_State *L)
 {
-	luaL_checkoption(L, 1, iter_opt[0], iter_opt);
-	lua_pushcfunction(L, lib_iterateThinkers);
-	lua_pushvalue(L, 1);
+	struct iterationState *it;
+
+	lua_pushvalue(L, lua_upvalueindex(1));
+	it = lua_newuserdata(L, sizeof(struct iterationState));
+	luaL_getmetatable(L, META_ITERATIONSTATE);
+	lua_setmetatable(L, -2);
+
+	it->filter = iter_funcs[luaL_checkoption(L, 1, "mobj", iter_opt)];
+	it->next = LUA_REFNIL;
 	return 2;
 }
 
+#undef push_thinker
+
 int LUA_ThinkerLib(lua_State *L)
 {
+	luaL_newmetatable(L, META_ITERATIONSTATE);
+	lua_pushcfunction(L, iterationState_gc);
+	lua_setfield(L, -2, "__gc");
+	lua_pop(L, 1);
+
 	lua_createtable(L, 0, 1);
-		lua_pushcfunction(L, lib_startIterate);
+		lua_pushcfunction(L, lib_iterateThinkers);
+		lua_pushcclosure(L, lib_startIterate, 1);
 		lua_setfield(L, -2, "iterate");
 	lua_setglobal(L, "thinkers");
 	return 0;