Lua safe implementation of single userdata and GC

/*
    Title:  Safe lua object implementation with userdata GC and single userdata per object
    Author: Rochet2 https://github.com/Rochet2
    A part of Eluna LuaEngine project: https://github.com/ElunaLuaEngine/Eluna
    Using C++ and lua 5.2

    Decided  to document this piece of work since I did not find this anywhere on the net : |

    They way this works is that lua has a table that has weak values.
    This means that the value can be garbage collected even though referenced by the table.
    The table will have a reference that we use to access it.

    We will use the table to store our userdata and access it from C++.
    The key will be pointer as string, so you will be able to access it just by having the pointer and the table reference.

    When pushing object, we check if the table has the pointer and if the value is valid.
    If so, we will return that value.
    Otherwise we will create a new userdata and push it to the lua table and return it.
    This means that there is only ONE userdata for an object.

    Since we will NOT use a reference to the userdata other than the table reference,
    which is weak and not counted, the userdata can be garbage collected.
    This means that the userdata will not dangle in memory for the whole existence of the actual object once pushed and can be freed if unused.

    Additionally we can remove the value / turn the value to nil in the table once the object is destroyed.
    This way when we check the userdata, we can check if the object still existsm if not -> error
    This is good practice so you wont run into invalid pointer accessing errors and if you do, you get an error safely.

    Pros:
        One userdata per object at a time
        Garbage collecting of unused userdata
        Still allows using lua mem managed objects easily ( just a setting in the regiester )
        Allows one to set up invalid pointer errors, requires adding code to destructors : /
        Does not use hashmap on C++ side

    Cons:
        May possibly be slow since when pushing and checking we always need to get the userdata table and the value from it
        On the other hand an article claimed accessing a variable with a lua ref was fast, so this can be fast as well :)
*/

// HEADER:

// this is a global variable
// this will hold the hidden table reference that /should/ always be valid
int tableref = LUA_REFNIL;
// This holds our lua state and is also a global var
lua_State* L = NULL;

template<typename T>
class ObjectTemplate
{
public:
    // type name
    static const char* tname;
    // if manageMemory == true, lua will handle memory management on GC call
    static bool manageMemory;

    // Called when an object should be converted to a string
    static int tostringT(lua_State* L)
    {
        T* obj = check(L, 1, false); // get self, no error
        if (obj)
        {
            // object exists, push typename and pointer in a string
            lua_pushfstring(L, "%s: (%p)", tname, obj);
            return 1;
        }
        // object does not exist or was invalid, replace object with nil and push it as string
        // (this is a copy paste from to string for nil)
        lua_pushnil(L);
        lua_replace(L, 1);
        luaL_tolstring(L, 1, NULL);
        return 1;
    }

    // this is not needed with this implementation :)
    // The reason is that there is only one userdata for object, so it works by default
    //// Checked when obj == obj is used
    //static int equalT(lua_State* L)
    //{
    //    // Get objects
    //    T* obj1 = check(L, 1);
    //    T* obj2 = check(L, 2);
    //    // Check that they point to same address and push result
    //    lua_pushboolean(L, obj1 == obj2);
    //    return 1;
    //}

    // This function is called on garbage collection for an Object
    static int gcT(lua_State* L)
    {
        // debug
        printf("gcT %s!\n", tname);

        // Check if lua does not handle mem management, return
        if (!manageMemory)
            return 0;

        // Get object pointer and check that has correct metatable
        T** ptrHold = static_cast<T**>(luaL_checkudata(L, -1, tname));
        if (ptrHold) // should always be true, check anyways
            delete *ptrHold; // is userdata, delete
        return 0;
    }

    // Registers a new lua object type
    // name will be used as typename (metatable name and error msg name)
    // If gc is true, lua will handle the memory management for object pushed
    static void Register(lua_State* L, const char* name, bool gc = false)
    {
        // Should register a type only once !!
        ASSERT(!tname);

        // set typename and memory management settings for this typename
        tname = name;
        manageMemory = gc;

        // method table
        lua_newtable(L);
        int methods = lua_gettop(L);

        // store method table in globals so that
        // scripts can add functions in Lua
        lua_pushvalue(L, methods);
        lua_setglobal(L, tname);

        luaL_newmetatable(L, tname);
        int metatable = lua_gettop(L);

        // hide metatable from user
        lua_pushvalue(L, methods);
        lua_setfield(L, metatable, "__metatable");

        // required to access methods IE obj:GetName()
        lua_pushvalue(L, methods);
        lua_setfield(L, metatable, "__index");

        // to string
        lua_pushcfunction(L, tostringT);
        lua_setfield(L, metatable, "__tostring");

        // garbage collector
        lua_pushcfunction(L, gcT);
        lua_setfield(L, metatable, "__gc");

        // equality check
        //lua_pushcfunction(L, equalT);
        //lua_setfield(L, metatable, "__eq");

        // special method to get the object type
        lua_pushcfunction(L, typeT);
        lua_setfield(L, methods, "GetObjectType");

        // set the metatable for the method table
        lua_setmetatable(L, methods);

        // remove the method table from the stack
        lua_remove(L, methods);
    }

    // Adds a new method / function for the object type
    static void SetMethod(lua_State* L, const char* funcname, lua_CFunction func)
    {
        if (!funcname || !func)
            return;

        // Get metatable and check that it is valid
        luaL_getmetatable(L, tname);
        if (!lua_istable(L, -1))
        {
            const char* err = lua_pushfstring(L, "%s missing metatable", tname);
            luaL_argerror(L, -2, err);
            return;
        }

        // Get the method table from the metatable and check that it is valid
        lua_getfield(L, -1, "__metatable");
        lua_remove(L, -2); // remove metatable as it is not needed anymore
        if (!lua_istable(L, -1))
        {
            const char* err = lua_pushfstring(L, "%s missing method table from metatable", tname);
            luaL_argerror(L, -2, err);
            return;
        }

        // add the function to the method table
        lua_pushstring(L, funcname);
        lua_pushcclosure(L, func, 0);
        lua_settable(L, -3);

        // remove the method table from stack
        lua_remove(L, -1);
    }

    // pushes userdata with the given pointer inside
    static int push(lua_State* L, T const* obj)
    {
        // NULL handling
        if (!obj)
        {
            lua_pushnil(L);
            return 1;
        }

        // If an object is mem managed by lua, there will / should always be
        // always be only one userdata for it anyways in our implementation.
        if (!manageMemory)
        {
            // Get our table with userdatas, could have error checking
            lua_rawgeti(L, LUA_REGISTRYINDEX, tableref);
            // Try finding the userdata for this pointer (pointer as string is key)
            lua_pushfstring(L, "%p", obj);
            lua_gettable(L, -2);
            // pop the table
            lua_remove(L, -2);
            // check that the userdata is valid, no error
            if (!lua_isnoneornil(L, -1) && luaL_checkudata(L, -1, tname))
            {
                // debug
                printf("Pushed OLD userdata %s!\n", tname);

                // return valid userdata
                return 1;
            }
            // was not valid, remove the userdata from stack
            lua_remove(L, -1);
        }

        // Create new userdata
        T const** ptrHold = (T const**)lua_newuserdata(L, sizeof(T const**));
        if (!ptrHold)
        {
            // Errored.. not enough memory?
            // remove the userdata value and push nil instead
            lua_remove(L, -1);
            lua_pushnil(L);
            return 1;
        }
        *ptrHold = obj; // set the pointer inside the userdata

        // Get metatable for the userdata and check that it is valid
        luaL_getmetatable(L, tname);
        if (!lua_istable(L, -1))
        {
            // Errored.. no metatable for the type?
            // remove the userdata value and the metatable and push nil instead
            lua_pop(L, 2);
            lua_remove(L, -1);
            lua_pushnil(L);
            return 1;
        }
        // Set the metatable
        lua_setmetatable(L, -2);

        if (!manageMemory)
        {
            // Get hidden lua table, could have error checking
            lua_rawgeti(L, LUA_REGISTRYINDEX, tableref);
            // place userdata in it with the pointer string as key
            lua_pushfstring(L, "%p", obj);
            lua_pushvalue(L, -3);
            lua_settable(L, -3);
            // and pop the table
            lua_remove(L, -1);
        }

        // debug
        printf("Pushed NEW userdata %s!\n", tname);
        return 1;
    }

    // check / get a pointer of this type at given narg / index
    // error will control whether it should error if the object is nil or wrong type
    static T* check(lua_State* L, int narg, bool error = true)
    {
        // debug
        printf("Checking %s!\n", tname);

        // Get the userdata, error if error is true
        T** ptrHold = static_cast<T**>(error ? luaL_checkudata(L, narg, tname) : lua_touserdata(L, narg));
        // if error was false, this can be NULL, so check it and error if necessary
        if (!ptrHold)
        {
            if (error)
            {
                // print an error
                char buff[256];
                snprintf(buff, 256, "%s expected, got %s", tname, luaL_typename(L, narg));
                luaL_argerror(L, narg, buff);
            }
            return NULL;
        }

        // Only memory managed objects need this in this implementation
        // This is only useful if the table element is erased by the destructor of object
        if (!manageMemory)
        {
            // Check pointer validity

            // Get userdata table
            lua_rawgeti(L, LUA_REGISTRYINDEX, tableref);
            // get the userdata from the table for our pointer
            lua_pushfstring(L, "%p", *ptrHold);
            lua_gettable(L, -2);
            // remove the userdata table
            lua_remove(L, -2);
            // check validity and remove the userdata
            // could use luaL_checkudata to error here, but our implementation will have objects that inherit other types.
            // This means that it would error for wrong type even if it was correct
            bool valid = lua_isuserdata(L, -1);
            lua_remove(L, -1);
            if (!valid)
            {
                // Should display error message regardless as a dead pointer is no good!
                char buff[256];
                snprintf(buff, 256, "%s expected, got pointer to nonexisting object (%s). This should never happen", tname, luaL_typename(L, narg));
                if (error)
                {
                    luaL_argerror(L, narg, buff);
                }
                else
                {
                    // this is a silent error since error was false
                    printf("%s\n", buff);
                }
                return NULL;
            }
        }

        // All checks were correct, return the pointer from the userdata
        return *ptrHold;
    }

    // Removes the object userdata from lua if exists
    static void RemoveRef(lua_State* L, T* obj)
    {
        // debug
        printf("RemoveRef %s!\n", tname);

        // Get hidden userdata table, could check errors
        lua_rawgeti(L, LUA_REGISTRYINDEX, tableref);
        // get the userdata for our object
        lua_pushfstring(L, "%p", obj);
        lua_gettable(L, -2);
        // check that it was valid
        if (luaL_checkudata(L, -1))
        {
            lua_pushfstring(L, "%p", obj);
            lua_pushnil(L);
            lua_settable(L, -4);
        }
        lua_pop(L, 2);
    }
};

// CPP:

template<typename T> const char* ObjectTemplate<T>::tname = NULL;
template<typename T> bool ObjectTemplate<T>::manageMemory = false;

class TestObj
{
public:
    TestObj()
    {
        var = 5;
    }

    ~TestObj()
    {
        ObjectTemplate<TestObj>::RemoveRef(L, this);
    }

    int var;
};

static int ObjGetVar(lua_State* L)
{
    // get self ( will error if not valid )
    TestObj* obj = ObjectTemplate<TestObj>::check(L, 1);
    lua_pushnumber(L, obj->var);
    return 1;
}

void LuaMain()
{
    // create lua base (state and libs)
    L = luaL_newstate();
    luaL_openlibs(L);

    // Create table
    lua_newtable(L);
    // Create metatable
    lua_newtable(L);
    // set __mode to v for metatable, which means weak values
    lua_pushstring(L, "v");
    lua_setfield(L, -2, "__mode");
    // set the metatable for the table, pops metatable
    lua_setmetatable(L, -2);
    // create lua reference to the table and save it to the global variable
    tableref = luaL_ref(L, LUA_REGISTRYINDEX);

    // register a new type and add a method for it
    ObjectTemplate<TestObj>::Register(L, "TestObj");
    ObjectTemplate<TestObj>::SetMethod(L, "GetVar", ObjGetVar);

    // make a few tests..
    // Create a new object of the type and push it
    TestObj* obj1 = new TestObj();
    ObjectTemplate<TestObj>::push(L, obj1);

    // test lua garbage collection on the object
    // should collect I think
    lua_gc(L, LUA_GCCOLLECT, 0);

    // push it twice this time to test that it pushes the old again on second time
    ObjectTemplate<TestObj>::push(L, obj1);
    ObjectTemplate<TestObj>::push(L, obj1);

    // check it, should work
    TestObj* obj2 = ObjectTemplate<TestObj>::check(L, -1);
    // delete it, now pointers are invalid :(
    delete obj2;

    // check it again, this will result in an error in this example implementation,
    // since the invalid pointer check is there :)
    TestObj* obj3 = ObjectTemplate<TestObj>::check(L, -1);
}