Allow associating a table with Lua-bound objects, for storing state in Lua.

Design goals are: minimal overhead (one table); minimal extra code to use
it (two lines in the header, to have a space to store a reference to the
table), same API for classes with and without a table; allow switching to
a table-bound class without derived classes having to know about it; keep
zero per-object overhead for classes that don't need it.

stepmania/src/LuaBinding.cpp

@@ -1,5 +1,6 @@
 #include "global.h"
 #include "LuaBinding.h"
+#include "RageUtil.h"
 
 void CreateMethodsTable( lua_State *L, const CString &szName )
 {
@@ -63,6 +64,196 @@ bool CheckLuaObjectType( lua_State *L, int narg, const char *szType )
 	}
 }
 
+bool GetGlobalTable( Lua *L, bool bCreate )
+{
+	lua_pushstring( L, "userdatas" );
+	lua_rawget( L, LUA_REGISTRYINDEX );
+	if( !lua_isnil(L, -1) )
+		return true;
+
+	lua_pop( L, 1 );
+	if( !bCreate )
+		return false;
+
+	/* Save it. */
+	lua_newtable( L );
+	lua_pushstring( L, "userdatas" );
+	lua_pushvalue( L, -2 );
+	lua_rawset( L, LUA_REGISTRYINDEX );
+	return true;
+}
+
+/* The object is on the stack.  It's either a table or a userdata.
+ * If needed, associate the metatable; if a table, also add it to
+ * the userdata table. */
+void ApplyDerivedType( Lua *L, const CString &sClassName, void *pSelf )
+{
+	int iTable = lua_gettop( L );
+
+	int iType = lua_type( L, iTable );
+	ASSERT_M( iType == LUA_TTABLE || iType == LUA_TUSERDATA,
+		ssprintf("%i", iType) );
+
+	if( iType == LUA_TTABLE )
+	{
+		GetGlobalTable( L, true );
+		int iGlobalTable = lua_gettop( L );
+
+		/* If the table is already in the userdata table, then everything
+		 * is already set up. */
+		lua_pushvalue( L, iTable );
+		lua_rawget( L, iGlobalTable );
+		if( !lua_isnil(L, -1) )
+		{
+			void *pData = lua_touserdata( L, -1 );
+			ASSERT( pSelf == pData );
+			ASSERT( pSelf == lua_touserdata(L, -1) );
+
+			lua_settop( L, iTable );
+			return;
+		}
+
+		/* Create the userdata, and add it to the global table. */
+		lua_pushvalue( L, iTable );
+		lua_pushlightuserdata( L, pSelf );
+		lua_rawset( L, iGlobalTable );
+
+		/* Pop everything except the table. */
+		lua_settop( L, iTable );
+	}
+
+	luaL_getmetatable( L, sClassName );
+	lua_setmetatable( L, iTable );
+}
+
+#include "RageUtil_AutoPtr.h"
+REGISTER_CLASS_TRAITS( LuaClass, new LuaClass(*pCopy) )
+
+LuaClass::LuaClass()
+{
+	m_pSelf = NULL;
+}
+
+void *GetUserdataFromGlobalTable( Lua *L, const char *szType, int iArg )
+{
+	if( !GetGlobalTable(L, false) )
+		luaL_error( L, "stale %s referenced (object used but no longer exists)", szType );
+
+	lua_pushvalue( L, iArg );
+	lua_rawget( L, -2 );
+	if( lua_isnil(L, -1) )
+		luaL_error( L, "stale %s referenced (object used but no longer exists)", szType );
+
+	void *pRet = lua_touserdata( L, -1 );
+	lua_pop( L, 2 );
+
+	return pRet;
+}
+
+/* Tricky: when an instance table is copied, we want to do a deep
+ * copy, not a reference copy.  Otherwise, the new higher-level object
+ * will share a table with the original.  Aside from being confusing,
+ * this breaks the global table, which assumes that we have a one-to-
+ * one mapping between tables and objects. */
+LuaClass::LuaClass( const LuaClass &cpy ):
+	LuaTable(cpy)
+{
+	if( !IsSet() )
+		return;
+
+	CString sData = Serialize();
+	LoadFromString( sData );
+}
+
+LuaClass &LuaClass::operator=( const LuaClass &cpy )
+{
+	LuaTable::operator=(cpy);
+
+	if( !IsSet() )
+		return *this;
+
+	CString sData = Serialize();
+	LoadFromString( sData );
+
+	return *this;
+}
+
+
+LuaClass::~LuaClass()
+{
+	if( LUA == NULL )
+		return;
+
+	Lua *L = LUA->Get();
+
+	int iTop = lua_gettop( L );
+
+	/* If we're registered in the global table, unregister. */
+	if( GetGlobalTable(L, false) )
+	{
+		this->PushSelf( L );
+		lua_pushnil( L );
+		lua_rawset( L, -3 );
+	}
+
+	lua_settop( L, iTop );
+
+	LUA->Release( L );
+}
+
+void LuaClass::BeforeReset()
+{
+	LuaTable::BeforeReset();
+
+	/* Read pSelf the name of the class, so we can use them to restore in Register(). */
+	Lua *L = LUA->Get();
+
+	if( !GetGlobalTable(L, false) )
+	{
+		LUA->Release( L );
+		return;
+	}
+
+	this->PushSelf( L );
+	lua_rawget( L, -2 );
+	if( lua_isnil(L, -1) )
+	{
+		/* This table hasn't been pushed yet. */
+		lua_pop( L, 2 );
+		LUA->Release( L );
+		return;
+	}
+
+	m_pSelf = lua_touserdata( L, -1 );
+	lua_pop( L, 2 );
+
+	this->PushSelf( L );
+	lua_getmetatable( L, -1 );
+	lua_rawget( L, LUA_REGISTRYINDEX );
+	ASSERT( !lua_isnil(L, -1) );
+	m_sClassName = lua_tostring( L, -1 );
+	lua_pop( L, 1 );
+
+	LUA->Release( L );
+}
+
+void LuaClass::Register()
+{
+	LuaTable::Register();
+
+	if( m_pSelf != NULL )
+	{
+		Lua *L = LUA->Get();
+		this->PushSelf(L);
+		ApplyDerivedType( L, m_sClassName, m_pSelf );
+		LUA->Release( L );
+
+		/* To conserve memory, clear the class name.  We only need it while restoring. */
+		m_sClassName = CString();
+		m_pSelf = NULL;
+	}
+}
+
 /*
  * (c) 2005 Glenn Maynard
  * All rights reserved.

stepmania/src/LuaBinding.h

@@ -4,9 +4,12 @@
 #define LuaBinding_H
 
 #include "LuaManager.h"
+class LuaReference;
 
 void CreateMethodsTable( lua_State *L, const CString &szName );
 bool CheckLuaObjectType( lua_State *L, int narg, const char *szType );
+void *GetUserdataFromGlobalTable( Lua *L, const char *szType, int iArg );
+void ApplyDerivedType( Lua *L, const CString &sClassname, void *pSelf );
 
 template <typename Type>
 class Luna 
@@ -101,20 +104,29 @@ public:
 				LuaHelpers::TypeError( L, narg, m_sClassName );
 		}
 
-		void **pData = (void **) lua_touserdata( L, narg );
+		return get( L, narg );
-		return (T *) *pData;
 	}
 	
-	// push a userdata containing a pointer to T object
+	static T *get( lua_State *L, int narg )
-	static int Push( Lua *L, T* p )
 	{
-		void **pData = (void **) lua_newuserdata( L, sizeof(void *) );
+		/* The stack has a userdata or a table.  If it's a table, look up the associated userdata. */
-		*pData = p;  // store pointer to object in userdata
+		if( lua_istable(L, narg) )
-		luaL_getmetatable( L, m_sClassName );  // lookup metatable in Lua registry
+		{
-		lua_setmetatable( L, -2 );
+			return (T *) GetUserdataFromGlobalTable( L, m_sClassName, narg );
-		return 1;  // userdata containing pointer to T object
+		}
+		else if( lua_isuserdata(L, narg) )
+		{
+			void **pData = (void **) lua_touserdata( L, narg );
+			return (T *) *pData;
+		}
+		else
+			return NULL;
 	}
 	
+	/* Push a table or userdata for the given object.  This is called on the
+	 * base class, so we pick up the instance of the base class, if any. */ 
+	static void PushObject( Lua *L, T* p );
+
 	static void AddMethod( const char *szName, int (*pFunc)(T *p, lua_State *L) )
 	{
 		if( s_pvMethods == NULL )
@@ -140,9 +152,35 @@ private:
 	/* Two objects are equal if the underlying object is the same. */
 	static int equal( lua_State *L )
 	{
-		void **obj1 = (void **) lua_touserdata( L, 1 );
+		int iType = lua_type( L, 1 );
-		void **obj2 = (void **) lua_touserdata( L, 2 );
+		if( lua_type(L, 2) != iType )
-		lua_pushboolean( L, *obj1 == *obj2 );
+		{
+			lua_pushboolean( L, false );
+			return 1;
+		}
+
+		/* Use the regular method for tables.  If an object's table is
+		 * kept around after the actual object has been destroyed, the
+		 * table is still valid, and the pointer no longer exists. */
+		if( iType == LUA_TTABLE )
+		{
+			int iEqual = lua_rawequal( L, 1, 2 );
+			lua_pushboolean( L, iEqual );
+			return 1;
+		}
+
+		if( !CheckLuaObjectType(L, 1, m_sClassName) ||
+			!CheckLuaObjectType(L, 2, m_sClassName) )
+		{
+			lua_pushboolean( L, false );
+		}
+		else
+		{
+			void *pData1 = get( L, 1 );
+			void *pData2 = get( L, 2 );
+			lua_pushboolean( L, pData1 != NULL && pData1 == pData2 );
+		}
+
 		return 1;
 	}
 
@@ -155,22 +193,58 @@ private:
 	static int tostring_T( lua_State *L )
 	{
 		char buff[32];
-		void *pData = check( L, -1 );
+		void *pData = check( L, 1 );
 		sprintf( buff, "%p", pData );
 		lua_pushfstring( L, "%s (%s)", m_sClassName, buff );
 		return 1;
 	}
 };
 
+/*
+ * Instanced classes have an associated table, which is used as "self"
+ * instead of a raw userdata.  This should be as lightweight as possible.
+ */
+#include "LuaReference.h"
+class LuaClass: public LuaTable
+{
+public:
+	LuaClass();
+	LuaClass( const LuaClass &cpy );
+	virtual ~LuaClass();
+	LuaClass &operator=( const LuaClass &cpy );
+
+protected:
+	virtual void BeforeReset();
+	virtual void Register();
+
+	void *m_pSelf;
+
+	CString m_sClassName;
+};
+
+/* Only a base class has to indicate that it's instanced (has a per-object
+ * Lua table).  Derived classes simply call the base class's Push function,
+ * specifying a different class name, so they don't need to know about it. */
+#define LUA_REGISTER_INSTANCED_BASE_CLASS( T ) \
+	LUA_REGISTER_CLASS_BASIC( T, none ) \
+	void Luna<T>::PushObject( Lua *L, T* p ) { p->m_pLuaInstance->PushSelf( L ); } \
+	void T::PushSelf( lua_State *L ) { Luna<T>::PushObject( L, this ); ApplyDerivedType( L, #T, this ); }
+
 #define LUA_REGISTER_CLASS( T ) \
-	LUA_REGISTER_DERIVED_CLASS( T, none )
+	LUA_REGISTER_CLASS_BASIC( T, none ) \
+	void Luna<T>::PushObject( Lua *L, T* p ) { void **pData = (void **) lua_newuserdata( L, sizeof(void *) ); *pData = p; } \
+	void T::PushSelf( lua_State *L ) { Luna<T>::PushObject( L, this ); ApplyDerivedType( L, #T, this ); }
 
 #define LUA_REGISTER_DERIVED_CLASS( T, B ) \
+	LUA_REGISTER_CLASS_BASIC( T, B ) \
+	void Luna<T>::PushObject( Lua *L, T* p ) { Luna<B>::PushObject( L, p ); } \
+	void T::PushSelf( lua_State *L ) { Luna<B>::PushObject( L, this ); ApplyDerivedType( L, #T, this ); }
+
+#define LUA_REGISTER_CLASS_BASIC( T, B ) \
 	template<> const char *Luna<T>::m_sClassName = #T; \
 	template<> const char *Luna<T>::m_sBaseClassName = #B; \
 	template<> Luna<T>::RegTypeVector* Luna<T>::s_pvMethods = NULL; \
 	static Luna##T registera; \
-	void T::PushSelf( lua_State *L ) { Luna##T::Push( L, this ); } \
 	/* Call PushSelf, so we always call the derived Luna<T>::Push. */ \
 	namespace LuaHelpers { template<> void Push( T *pObject, lua_State *L ) { pObject->PushSelf( L ); } }