4 files changed, 70 insertions, 18 deletions

diff --git a/src/Bindings/LuaState.cpp b/src/Bindings/LuaState.cpp
index 50ce04f7c..61d206109 100644
--- a/src/Bindings/LuaState.cpp
+++ b/src/Bindings/LuaState.cpp
@@ -44,9 +44,69 @@ extern "C"
 const cLuaState::cRet cLuaState::Return = {};
 const cLuaState::cNil cLuaState::Nil = {};
 
-/** Each Lua state stores a pointer to its creating cLuaState in Lua globals, under this name.
-This way any cLuaState can reference the main cLuaState's TrackedCallbacks, mutex etc. */
-static const char * g_CanonLuaStateGlobalName = "_CuberiteInternal_CanonLuaState";
+
+
+
+
+
+////////////////////////////////////////////////////////////////////////////////
+// cCanonLuaStates:
+
+/** Tracks the canon cLuaState instances for each lua_State pointer.
+Used for tracked refs - the ref needs to be tracked by a single cLuaState (the canon state), despite being created from a different (attached) cLuaState.
+The canon state must be available without accessing the Lua state itself (so it cannot be stored within Lua). */
+class cCanonLuaStates
+{
+public:
+	/** Returns the canon Lua state for the specified lua_State pointer. */
+	static cLuaState * GetCanonState(lua_State * a_LuaState)
+	{
+		auto & inst = GetInstance();
+		cCSLock lock(inst.m_CS);
+		auto itr = inst.m_CanonStates.find(a_LuaState);
+		if (itr == inst.m_CanonStates.end())
+		{
+			return nullptr;
+		}
+		return itr->second;
+	}
+
+	/** Adds a new canon cLuaState instance to the map.
+	Used when a new Lua state is created, this informs the map that a new canon Lua state should be tracked. */
+	static void Add(cLuaState & a_LuaState)
+	{
+		auto & inst = GetInstance();
+		cCSLock lock(inst.m_CS);
+		ASSERT(inst.m_CanonStates.find(a_LuaState) == inst.m_CanonStates.end());
+		inst.m_CanonStates[a_LuaState.operator lua_State *()] = &a_LuaState;
+	}
+
+	/** Removes the bindings between the specified canon state and its lua_State pointer.
+	Used when a Lua state is being closed. */
+	static void Remove(cLuaState & a_LuaState)
+	{
+		auto & inst = GetInstance();
+		cCSLock lock(inst.m_CS);
+		auto itr = inst.m_CanonStates.find(a_LuaState);
+		ASSERT(itr != inst.m_CanonStates.end());
+		inst.m_CanonStates.erase(itr);
+	}
+
+protected:
+	/** The mutex protecting m_CanonStates against multithreaded access. */
+	cCriticalSection m_CS;
+
+	/** Map of lua_State pointers to their canon cLuaState instances. */
+	std::map<lua_State *, cLuaState *> m_CanonStates;
+
+
+	/** Returns the singleton instance of this class. */
+	static cCanonLuaStates & GetInstance(void)
+	{
+		static cCanonLuaStates canonLuaStates;
+		return canonLuaStates;
+	}
+};
 
 
 
@@ -426,10 +486,7 @@ void cLuaState::Create(void)
 	luaL_openlibs(m_LuaState);
 	m_IsOwned = true;
 	cLuaStateTracker::Add(*this);
-
-	// Add the CanonLuaState value into the Lua state, so that we can get it from anywhere:
-	lua_pushlightuserdata(m_LuaState, reinterpret_cast<void *>(this));
-	lua_setglobal(m_LuaState, g_CanonLuaStateGlobalName);
+	cCanonLuaStates::Add(*this);
 }
 
 
@@ -493,6 +550,7 @@ void cLuaState::Close(void)
 		}
 	}
 
+	cCanonLuaStates::Remove(*this);
 	cLuaStateTracker::Del(*this);
 	lua_close(m_LuaState);
 	m_LuaState = nullptr;
@@ -2146,15 +2204,9 @@ void cLuaState::LogStackValues(lua_State * a_LuaState, const char * a_Header)
 
 
 
-cLuaState * cLuaState::QueryCanonLuaState(void)
+cLuaState * cLuaState::QueryCanonLuaState(void) const
 {
-	// Get the CanonLuaState global from Lua:
-	auto cb = WalkToNamedGlobal(g_CanonLuaStateGlobalName);
-	if (!cb.IsValid())
-	{
-		return nullptr;
-	}
-	return reinterpret_cast<cLuaState *>(lua_touserdata(m_LuaState, -1));
+	return cCanonLuaStates::GetCanonState(m_LuaState);
 }
 
 
diff --git a/src/Bindings/LuaState.h b/src/Bindings/LuaState.h
index 9c97e96d4..01faf2c26 100644
--- a/src/Bindings/LuaState.h
+++ b/src/Bindings/LuaState.h
@@ -809,7 +809,7 @@ public:
 
 	/** Returns the canon Lua state (the primary cLuaState instance that was used to create, rather than attach, the lua_State structure).
 	Returns nullptr if the canon Lua state cannot be queried. */
-	cLuaState * QueryCanonLuaState(void);
+	cLuaState * QueryCanonLuaState(void) const;
 
 	/** Outputs to log a warning about API call being unable to read its parameters from the stack,
 	logs the stack trace and stack values. */
diff --git a/src/Bindings/LuaWindow.cpp b/src/Bindings/LuaWindow.cpp
index 6fffa7aac..83463520c 100644
--- a/src/Bindings/LuaWindow.cpp
+++ b/src/Bindings/LuaWindow.cpp
@@ -116,7 +116,7 @@ bool cLuaWindow::ClosedByPlayer(cPlayer & a_Player, bool a_CanRefuse)
 		bool res;
 		if (
 			m_OnClosing->Call(this, &a_Player, a_CanRefuse, cLuaState::Return, res) &&  // The callback succeeded
-			res                                                                        // The callback says not to close the window
+			res                                                                         // The callback says not to close the window
 		)
 		{
 			// The callback disagrees (the higher levels check the CanRefuse flag compliance)
diff --git a/src/Bindings/LuaWindow.h b/src/Bindings/LuaWindow.h
index 5f987d4c6..d90d52a94 100644
--- a/src/Bindings/LuaWindow.h
+++ b/src/Bindings/LuaWindow.h
@@ -60,7 +60,7 @@ protected:
 	/** Contents of the non-inventory part */
 	cItemGrid m_Contents;
 
-	/** The Lua state that has opened the window and owns the m_LuaRef */
+	/** The canon Lua state that has opened the window and owns the m_LuaRef */
 	cLuaState * m_LuaState;
 
 	/** The Lua callback to call when the window is closing for any player */