Separate methods per state

Author: mostlynick3

src/LuaEngine/ALETemplate.h

@@ -18,12 +18,32 @@ extern "C"
 #include "ALEUtility.h"
 #include "SharedDefines.h"
 
+enum MethodRegisterState
+{
+    METHOD_REG_MAP   = 0,
+    METHOD_REG_WORLD = 1,
+    METHOD_REG_ALL   = 2
+};
+
+struct ALEGlobalRegister
+{
+    const char* name;
+    int(*func)(lua_State*);
+    MethodRegisterState regState;
+
+    ALEGlobalRegister(const char* name, int(*f)(lua_State*), MethodRegisterState state = METHOD_REG_ALL)
+        : name(name), func(f), regState(state) {}
+
+    ALEGlobalRegister(const char* name, MethodRegisterState state = METHOD_REG_ALL)
+        : name(name), func(nullptr), regState(state) {}
+};
+
 class ALEGlobal
 {
 public:
     static int thunk(lua_State* L)
     {
-        luaL_Reg* l = static_cast<luaL_Reg*>(lua_touserdata(L, lua_upvalueindex(1)));
+        ALEGlobalRegister* l = static_cast<ALEGlobalRegister*>(lua_touserdata(L, lua_upvalueindex(1)));
         int top = lua_gettop(L);
         int expected = l->func(L);
         int args = lua_gettop(L) - top;
@@ -36,15 +56,38 @@ class ALEGlobal
         return expected;
     }
 
-    static void SetMethods(ALE* E, luaL_Reg* methodTable)
+    static int MethodWrongState(lua_State* L)
+    {
+        luaL_error(L, "attempt to call method '%s' that is not available in this state", lua_tostring(L, lua_upvalueindex(1)));
+        return 0;
+    }
+
+    static void SetMethods(ALE* E, ALEGlobalRegister* methodTable)
     {
         ASSERT(E);
         ASSERT(methodTable);
 
         lua_pushglobaltable(E->L);
 
-        for (; methodTable && methodTable->name && methodTable->func; ++methodTable)
+        for (; methodTable && methodTable->name; ++methodTable)
         {
+            if (methodTable->regState != METHOD_REG_ALL)
+            {
+                bool isMapState = (E->GetStateMapId() != ALE_GLOBAL_STATE);
+                if ((!isMapState && methodTable->regState == METHOD_REG_MAP) ||
+                    (isMapState && methodTable->regState == METHOD_REG_WORLD))
+                {
+                    lua_pushstring(E->L, methodTable->name);
+                    lua_pushstring(E->L, methodTable->name);
+                    lua_pushcclosure(E->L, MethodWrongState, 1);
+                    lua_rawset(E->L, -3);
+                    continue;
+                }
+            }
+
+            if (!methodTable->func)
+                continue;
+
             lua_pushstring(E->L, methodTable->name);
             lua_pushlightuserdata(E->L, (void*)methodTable);
             lua_pushcclosure(E->L, thunk, 1);
@@ -117,6 +160,13 @@ struct ALERegister
 {
     const char* name;
     int(*mfunc)(lua_State*, T*);
+    MethodRegisterState regState;
+
+    ALERegister(const char* name, int(*func)(lua_State*, T*), MethodRegisterState state = METHOD_REG_ALL)
+        : name(name), mfunc(func), regState(state) {}
+
+    ALERegister(const char* name, MethodRegisterState state = METHOD_REG_ALL)
+        : name(name), mfunc(nullptr), regState(state) {}
 };
 
 template<typename T>
@@ -241,8 +291,25 @@ class ALETemplate
         lua_rawget(E->L, LUA_REGISTRYINDEX);
         ASSERT(lua_istable(E->L, -1));
 
-        for (; methodTable && methodTable->name && methodTable->mfunc; ++methodTable)
+        for (; methodTable && methodTable->name; ++methodTable)
         {
+            if (methodTable->regState != METHOD_REG_ALL)
+            {
+                bool isMapState = (E->GetStateMapId() != ALE_GLOBAL_STATE);
+                if ((!isMapState && methodTable->regState == METHOD_REG_MAP) ||
+                    (isMapState && methodTable->regState == METHOD_REG_WORLD))
+                {
+                    lua_pushstring(E->L, methodTable->name);
+                    lua_pushstring(E->L, methodTable->name);
+                    lua_pushcclosure(E->L, MethodWrongState, 1);
+                    lua_rawset(E->L, -3);
+                    continue;
+                }
+            }
+
+            if (!methodTable->mfunc)
+                continue;
+
             lua_pushstring(E->L, methodTable->name);
             lua_pushlightuserdata(E->L, (void*)methodTable);
             lua_pushcclosure(E->L, CallMethod, 1);
@@ -322,6 +389,12 @@ class ALETemplate
         return 0;
     }
 
+    static int MethodWrongState(lua_State* L)
+    {
+        luaL_error(L, "attempt to call method '%s' that is not available in this state", lua_tostring(L, lua_upvalueindex(1)));
+        return 0;
+    }
+
     static int CallMethod(lua_State* L)
     {
         T* obj = ALE::CHECKOBJ<T>(L, 1); // get self