2 files changed, 83 insertions, 8 deletions

diff --git a/src/Bindings/LuaState.cpp b/src/Bindings/LuaState.cpp
index 81893c946..50ce04f7c 100644
--- a/src/Bindings/LuaState.cpp
+++ b/src/Bindings/LuaState.cpp
@@ -549,6 +549,8 @@ void cLuaState::Detach(void)
 
 void cLuaState::AddPackagePath(const AString & a_PathVariable, const AString & a_Path)
 {
+	ASSERT_LUA_STACK_BALANCE(m_LuaState);
+
 	// Get the current path:
 	lua_getfield(m_LuaState, LUA_GLOBALSINDEX, "package");                           // Stk: <package>
 	lua_getfield(m_LuaState, -1, a_PathVariable.c_str());                            // Stk: <package> <package.path>
@@ -574,6 +576,7 @@ void cLuaState::AddPackagePath(const AString & a_PathVariable, const AString & a
 bool cLuaState::LoadFile(const AString & a_FileName, bool a_LogWarnings)
 {
 	ASSERT(IsValid());
+	ASSERT_LUA_STACK_BALANCE(m_LuaState);
 
 	// Load the file:
 	int s = luaL_loadfile(m_LuaState, a_FileName.c_str());
@@ -588,7 +591,7 @@ bool cLuaState::LoadFile(const AString & a_FileName, bool a_LogWarnings)
 	}
 
 	// Execute the globals:
-	s = lua_pcall(m_LuaState, 0, LUA_MULTRET, 0);
+	s = lua_pcall(m_LuaState, 0, 0, 0);
 	if (s != 0)
 	{
 		if (a_LogWarnings)
@@ -609,6 +612,7 @@ bool cLuaState::LoadFile(const AString & a_FileName, bool a_LogWarnings)
 bool cLuaState::LoadString(const AString & a_StringToLoad, const AString & a_FileName, bool a_LogWarnings)
 {
 	ASSERT(IsValid());
+	ASSERT_LUA_STACK_BALANCE(m_LuaState);
 
 	// Load the file:
 	int s = luaL_loadstring(m_LuaState, a_StringToLoad.c_str());
@@ -649,6 +653,7 @@ bool cLuaState::HasFunction(const char * a_FunctionName)
 		return false;
 	}
 
+	ASSERT_LUA_STACK_BALANCE(m_LuaState);
 	lua_getglobal(m_LuaState, a_FunctionName);
 	bool res = (!lua_isnil(m_LuaState, -1) && lua_isfunction(m_LuaState, -1));
 	lua_pop(m_LuaState, 1);
@@ -1475,12 +1480,15 @@ bool cLuaState::CallFunction(int a_NumResults)
 	{
 		// The error has already been printed together with the stacktrace
 		LOGWARNING("Error in %s calling function %s()", m_SubsystemName.c_str(), CurrentFunctionName.c_str());
+
+		// Remove the error handler and error message from the stack:
+		ASSERT(lua_gettop(m_LuaState) == 2);
+		lua_pop(m_LuaState, 2);
 		return false;
 	}
 
 	// Remove the error handler from the stack:
 	lua_remove(m_LuaState, -a_NumResults - 1);
-
 	return true;
 }
 
@@ -2101,6 +2109,7 @@ void cLuaState::LogStackValues(const char * a_Header)
 void cLuaState::LogStackValues(lua_State * a_LuaState, const char * a_Header)
 {
 	// Format string consisting only of %s is used to appease the compiler
+	ASSERT_LUA_STACK_BALANCE(a_LuaState, false);
 	LOG("%s", (a_Header != nullptr) ? a_Header : "Lua C API Stack contents:");
 	for (int i = lua_gettop(a_LuaState); i > 0; i--)
 	{
@@ -2111,9 +2120,22 @@ void cLuaState::LogStackValues(lua_State * a_LuaState, const char * a_Header)
 			case LUA_TBOOLEAN: Value.assign((lua_toboolean(a_LuaState, i) != 0) ? "true" : "false"); break;
 			case LUA_TLIGHTUSERDATA: Printf(Value, "%p", lua_touserdata(a_LuaState, i)); break;
 			case LUA_TNUMBER:        Printf(Value, "%f", static_cast<double>(lua_tonumber(a_LuaState, i))); break;
-			case LUA_TSTRING:        Printf(Value, "%s", lua_tostring(a_LuaState, i)); break;
+			case LUA_TSTRING:
+			{
+				size_t len;
+				const char * txt = lua_tolstring(a_LuaState, i, &len);
+				Value.assign(txt, std::min<size_t>(len, 50));  // Only log up to 50 characters of the string
+				break;
+			}
 			case LUA_TTABLE:         Printf(Value, "%p", lua_topointer(a_LuaState, i)); break;
-			case LUA_TUSERDATA:      Printf(Value, "%p (%s)", lua_touserdata(a_LuaState, i), tolua_typename(a_LuaState, i)); break;
+			case LUA_TFUNCTION:      Printf(Value, "%p", lua_topointer(a_LuaState, i)); break;
+			case LUA_TUSERDATA:
+			{
+				Printf(Value, "%p (%s)", lua_touserdata(a_LuaState, i), tolua_typename(a_LuaState, i));
+				// tolua_typename pushes the string onto Lua stack, pop it off again:
+				lua_pop(a_LuaState, 1);
+				break;
+			}
 			default: break;
 		}
 		LOGD("  Idx %d: type %d (%s) %s", i, Type, lua_typename(a_LuaState, Type), Value.c_str());
@@ -2164,6 +2186,7 @@ int cLuaState::ReportFnCallErrors(lua_State * a_LuaState)
 
 int cLuaState::BreakIntoDebugger(lua_State * a_LuaState)
 {
+	ASSERT_LUA_STACK_BALANCE(a_LuaState);
 	// Call the BreakIntoDebugger function, if available:
 	lua_getglobal(a_LuaState, "BreakIntoDebugger");
 	if (!lua_isfunction(a_LuaState, -1))
@@ -2172,11 +2195,10 @@ int cLuaState::BreakIntoDebugger(lua_State * a_LuaState)
 		lua_pop(a_LuaState, 1);
 		return 1;
 	}
-	lua_insert(a_LuaState, -2);  // Copy the string that has been passed to us
+	lua_pushvalue(a_LuaState, -2);  // Copy the string that has been passed to us
 	LOGD("Calling BreakIntoDebugger()...");
 	lua_call(a_LuaState, 1, 0);
 	LOGD("Returned from BreakIntoDebugger().");
-
 	return 0;
 }
 
diff --git a/src/Bindings/LuaState.h b/src/Bindings/LuaState.h
index 32d346a19..9c97e96d4 100644
--- a/src/Bindings/LuaState.h
+++ b/src/Bindings/LuaState.h
@@ -55,6 +55,54 @@ class cLuaState
 {
 public:
 
+	#ifdef _DEBUG
+		/** Asserts that the Lua stack has the same amount of entries when this object is destructed, as when it was constructed.
+		Used for checking functions that should preserve Lua stack balance. */
+		class cStackBalanceCheck
+		{
+		public:
+			cStackBalanceCheck(const char * a_FileName, int a_LineNum, lua_State * a_LuaState, bool a_ShouldLogStack = true):
+				m_FileName(a_FileName),
+				m_LineNum(a_LineNum),
+				m_LuaState(a_LuaState),
+				m_StackPos(lua_gettop(a_LuaState))
+			{
+				if (a_ShouldLogStack)
+				{
+					// DEBUG: If an unbalanced stack is reported, uncommenting the next line can help debug the imbalance
+					// cLuaState::LogStackValues(a_LuaState, Printf("Started checking Lua stack balance, currently %d items:", m_StackPos).c_str());
+					// Since LogStackValues() itself uses the balance check, we must not call it recursively
+				}
+			}
+
+			~cStackBalanceCheck()
+			{
+				auto currStackPos = lua_gettop(m_LuaState);
+				if (currStackPos != m_StackPos)
+				{
+					LOGD("Lua stack not balanced. Expected %d items, found %d items. Stack watching started in %s:%d",
+						m_StackPos, currStackPos,
+						m_FileName.c_str(), m_LineNum
+					);
+					cLuaState::LogStackValues(m_LuaState);
+					ASSERT(!"Lua stack unbalanced");  // If this assert fires, the Lua stack is inbalanced, check the callstack / m_FileName / m_LineNum
+				}
+			}
+
+		protected:
+			const AString m_FileName;
+			int m_LineNum;
+			lua_State * m_LuaState;
+			int m_StackPos;
+		};
+
+		#define STRINGIFY2(X, Y) X##Y
+		#define STRINGIFY(X, Y) STRINGIFY2(X, Y)
+		#define ASSERT_LUA_STACK_BALANCE(...) cStackBalanceCheck STRINGIFY(Check, __COUNTER__)(__FILE__, __LINE__, __VA_ARGS__)
+	#else
+		#define ASSERT_LUA_STACK_BALANCE(...)
+	#endif
+
 	/** Provides a RAII-style locking for the LuaState.
 	Used mainly by the cPluginLua internals to provide the actual locking for interface operations, such as callbacks. */
 	class cLock
@@ -654,13 +702,15 @@ public:
 	template <typename FnT, typename... Args>
 	bool Call(const FnT & a_Function, Args &&... args)
 	{
+		ASSERT_LUA_STACK_BALANCE(m_LuaState);
 		m_NumCurrentFunctionArgs = -1;
 		if (!PushFunction(std::forward<const FnT &>(a_Function)))
 		{
 			// Pushing the function failed
 			return false;
 		}
-		return PushCallPop(std::forward<Args>(args)...);
+		auto res = PushCallPop(std::forward<Args>(args)...);
+		return res;
 	}
 
 	/** Retrieves a list of values from the Lua stack, starting at the specified index. */
@@ -873,7 +923,10 @@ protected:
 	/**
 	Calls the function that has been pushed onto the stack by PushFunction(),
 	with arguments pushed by PushXXX().
-	Returns true if successful, logs a warning on failure.
+	Returns true if successful, returns false and logs a warning on failure.
+	Pops the function params, the function itself and the error handler off the stack.
+	If successful, leaves a_NumReturnValues new values on Lua stack, corresponding to the return values.
+	On failure, leaves no new values on the Lua stack.
 	*/
 	bool CallFunction(int a_NumReturnValues);