4 files changed, 91 insertions, 5 deletions

diff --git a/src/lua/common/je_lua.cpp b/src/lua/common/je_lua.cpp
index ad3f756..d53898d 100644
--- a/src/lua/common/je_lua.cpp
+++ b/src/lua/common/je_lua.cpp
@@ -115,5 +115,62 @@ namespace JinEngine
             return 1;
         }
 
+        void luax_getreference(lua_State* L, SharedBase* shared)
+        {
+            luax_getreferencestable(L);
+            luax_pushlightuserdata(L, shared);
+            luax_gettable(L, -2);
+            luax_remove(L, -2);
+        }
+
+        bool luax_addreference(lua_State* L, SharedBase* shared, SharedBase* dep)
+        {
+            luax_getreference(L, shared);
+            // If no dependencies table, add one.
+            if (luax_isnil(L, -1))
+            {
+                luax_pop(L, 1);
+                luax_getreferencestable(L);
+                luax_newtable(L);
+                luax_pushlightuserdata(L, shared);
+                luax_pushvalue(L, -2);
+                luax_settable(L, -4);
+                luax_remove(L, -2); // Remove references table.
+            }
+            luax_pushlightuserdata(L, dep);
+            luax_getobject(L, dep);
+            if (luax_isnil(L, -1))
+            {
+                luax_pop(L, 3); // Pop nil, dep, reftbl
+                return false;
+            }
+            luax_settable(L, -3);
+            luax_pop(L, 1);
+            return true;
+        }
+
+        void luax_removereference(lua_State* L, SharedBase* shared)
+        {
+            luax_getreferencestable(L);
+            luax_pushlightuserdata(L, shared);
+            luax_pushnil(L);
+            luax_settable(L, -3);
+            luax_pop(L, 1);
+        }
+
+        void luax_removereference(lua_State* L, SharedBase* shared, SharedBase* dep)
+        {
+            luax_getreference(L, shared);
+            if (luax_isnil(L, -1))
+            {
+                luax_pop(L, 1);
+                return;
+            }
+            luax_pushlightuserdata(L, dep);
+            luax_pushnil(L);
+            luax_settable(L, -3);
+            luax_pop(L, 1);
+        }
+
     }
 }
\ No newline at end of file
diff --git a/src/lua/common/je_lua.h b/src/lua/common/je_lua.h
index 05cf917..b6fc878 100644
--- a/src/lua/common/je_lua.h
+++ b/src/lua/common/je_lua.h
@@ -26,6 +26,26 @@ namespace JinEngine
         int luax_getobject(lua_State* L, SharedBase* shared);
 
         ///
+        /// Get object's reference table.
+        ///
+        void luax_getreference(lua_State* L, SharedBase* shared);
+
+        ///
+        ///
+        ///
+        bool luax_addreference(lua_State* L, SharedBase* shared, SharedBase* dep);
+
+        ///
+        ///
+        ///
+        void luax_removereference(lua_State* L, SharedBase* shared);
+
+        ///
+        ///
+        ///
+        void luax_removereference(lua_State* L, SharedBase* shared, SharedBase* dep);
+
+        ///
         ///
         ///
         void luax_removeobject(lua_State* L, SharedBase* shared);
diff --git a/src/lua/common/je_lua_shared.cpp b/src/lua/common/je_lua_shared.cpp
index 241aec2..55832fd 100644
--- a/src/lua/common/je_lua_shared.cpp
+++ b/src/lua/common/je_lua_shared.cpp
@@ -15,6 +15,8 @@ namespace JinEngine
         {
             if (--mCount <= 0)
             {
+                // Remove all objects referenced by object.
+                luax_removereference(mL, this);
                 // Remove game object reference.
                 luax_removeobject(mL, this);
                 delete this;
@@ -26,6 +28,8 @@ namespace JinEngine
             removeDependency(key);
             shared->retain();
             mDependencies.insert(std::pair<int, SharedBase*>(key, shared));
+            // Set lua reference.
+            luax_addreference(mL, this, shared);
         }
 
         void SharedBase::removeDependency(int key)
@@ -33,17 +37,23 @@ namespace JinEngine
             if (!isDependOn(key))
                 return;
             std::map<int, SharedBase*>::iterator it = mDependencies.find(key);
-            it->second->release();
+            SharedBase* dep = it->second;
+            // Remove lua reference.
+            luax_removereference(mL, this, dep);
+            dep->release();
             mDependencies.erase(it);
         }
 
-        void SharedBase::removeDependency(SharedBase* dep)
+        void SharedBase::removeDependency(SharedBase* dependency)
         {
             for (std::map<int, SharedBase*>::iterator it = mDependencies.begin(); it != mDependencies.end();)
             {
-                if (it->second == dep)
+                SharedBase* dep = it->second;
+                if (dep == dependency)
                 {
-                    it->second->release();
+                    dep->release();
+                    // Remove lua reference.
+                    luax_removereference(mL, this, dep);
                     mDependencies.erase(it);
                 }
                 else
diff --git a/src/lua/modules/je_lua_modules.h b/src/lua/modules/je_lua_modules.h
index 220e6e5..afe0a82 100644
--- a/src/lua/modules/je_lua_modules.h
+++ b/src/lua/modules/je_lua_modules.h
@@ -16,5 +16,4 @@
 #include "thread/je_lua_thread.h"
 #include "time/je_lua_time.h"
 
-
 #endif
\ No newline at end of file