From 2518ff436a272343fb38ae3b6962c365003392fc Mon Sep 17 00:00:00 2001
From: SDraw <sdraw@protonmail.com>
Date: Mon, 5 Jun 2023 22:57:13 +0300
Subject: [PATCH] Critical fix of scripts and resources handling

---
 CVRLua/Lua/LuaVM.cs  | 16 +++++++++++-----
 CVRLua/LuaHandler.cs | 14 ++++++++++++--
 CVRLua/LuaScript.cs  | 16 ++++++++++------
 CVRLua/Main.cs       | 20 +++++++++++++++++---
 4 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/CVRLua/Lua/LuaVM.cs b/CVRLua/Lua/LuaVM.cs
index e228782..5ffcac5 100644
--- a/CVRLua/Lua/LuaVM.cs
+++ b/CVRLua/Lua/LuaVM.cs
@@ -4,7 +4,7 @@
 
 namespace CVRLua.Lua
 {
-    class LuaVM
+    class LuaVM : IDisposable
     {
         class ReferencedObject
         {
@@ -25,7 +25,7 @@ public ReferencedObject(object p_obj)
 
         static readonly Dictionary<IntPtr, LuaVM> ms_VMs = new Dictionary<IntPtr, LuaVM>();
 
-        readonly IntPtr m_state = IntPtr.Zero;
+        IntPtr m_state = IntPtr.Zero;
         readonly Dictionary<long, ReferencedObject> m_objectsMap = null;
 
         internal LuaVM()
@@ -52,10 +52,16 @@ internal LuaVM()
             LuaInterop.lua_setmetatable(m_state, -2); // Combines two previous tables
             LuaInterop.lua_setfield(m_state, LuaInterop.LUA_REGISTRYINDEX, c_objectsPool);
         }
-        ~LuaVM()
+
+        public void Dispose()
         {
-            ms_VMs.Remove(m_state);
-            LuaInterop.lua_close(m_state);
+            if(m_state != IntPtr.Zero)
+            {
+                m_objectsMap.Clear();
+                ms_VMs.Remove(m_state);
+                LuaInterop.lua_close(m_state);
+                m_state = IntPtr.Zero;
+            }
         }
 
         // Generic Lua stuff
diff --git a/CVRLua/LuaHandler.cs b/CVRLua/LuaHandler.cs
index efc5a03..59cc729 100644
--- a/CVRLua/LuaHandler.cs
+++ b/CVRLua/LuaHandler.cs
@@ -3,7 +3,7 @@
 
 namespace CVRLua
 {
-    class LuaHandler
+    class LuaHandler : IDisposable
     {
         public enum ScriptEvent
         {
@@ -35,7 +35,7 @@ public enum ScriptEvent
             OnAttachmentDeattach
         }
 
-        readonly Lua.LuaVM m_vm = null;
+        Lua.LuaVM m_vm = null;
         readonly Dictionary<ScriptEvent, int> m_eventFunctions = null; // Event <-> Reference
 
         internal static void Init()
@@ -167,6 +167,16 @@ internal LuaHandler()
             Lua.LuaDefs.UtilityDefs.RegisterInVM(m_vm);
         }
 
+        public void Dispose()
+        {
+            m_eventFunctions.Clear();
+            if(m_vm != null)
+            {
+                m_vm.Dispose();
+                m_vm = null;
+            }
+        }
+
         public void Execute(string p_code)
         {
             m_vm.Execute(p_code);
diff --git a/CVRLua/LuaScript.cs b/CVRLua/LuaScript.cs
index 2866b72..9476d65 100644
--- a/CVRLua/LuaScript.cs
+++ b/CVRLua/LuaScript.cs
@@ -5,7 +5,7 @@
 
 namespace CVRLua
 {
-    public class LuaScript : MonoBehaviour
+    public class LuaScript : MonoBehaviour, System.IDisposable
     {
         public List<TextAsset> Scripts = new List<TextAsset>();
         public List<string> VariableNames = new List<string>();
@@ -18,11 +18,6 @@ public class LuaScript : MonoBehaviour
         CVRInteractable m_interactable = null;
         CVRAttachment m_attachment = null;
 
-        ~LuaScript()
-        {
-            Core.Instance?.UnregisterScript(this); // Yes, it works
-        }
-
         void Awake()
         {
             if(m_luaHandler == null)
@@ -65,6 +60,15 @@ void Awake()
             }
         }
 
+        public void Dispose()
+        {
+            if(m_luaHandler != null)
+            {
+                m_luaHandler.Dispose();
+                m_luaHandler = null;
+            }
+        }
+
         void Start()
         {
             m_luaHandler?.CallEvent(LuaHandler.ScriptEvent.Start);
diff --git a/CVRLua/Main.cs b/CVRLua/Main.cs
index b159a2e..08812e9 100644
--- a/CVRLua/Main.cs
+++ b/CVRLua/Main.cs
@@ -5,12 +5,13 @@
 using System.Collections.Generic;
 using System.Reflection;
 using UnityEngine;
+using UnityEngine.SceneManagement;
 
 namespace CVRLua
 {
     public class Core : MelonLoader.MelonMod
     {
-        public const int c_modRelease = 20;
+        public const int c_modRelease = 21;
 
         static public Core Instance { get; private set; } = null;
 
@@ -62,6 +63,8 @@ public override void OnInitializeMelon()
                 null,
                 new HarmonyLib.HarmonyMethod(typeof(Core).GetMethod(nameof(OnPuppetMasterStart_Postfix), BindingFlags.NonPublic | BindingFlags.Static))
             );
+
+            SceneManager.sceneLoaded += this.OnSceneWasLoaded;
         }
 
         public override void OnDeinitializeMelon()
@@ -70,14 +73,25 @@ public override void OnDeinitializeMelon()
                 Instance = null;
         }
 
+        // Scrips register and removal
         internal void RegisterScript(LuaScript p_script)
         {
             m_scripts.Add(p_script);
         }
 
-        internal void UnregisterScript(LuaScript p_script)
+        void OnSceneWasLoaded(Scene p_scene, LoadSceneMode p_mode)
         {
-            m_scripts.Remove(p_script);
+            while(true)
+            {
+                int l_index = m_scripts.FindIndex(p => (p == null));
+                if(l_index != -1)
+                {
+                    m_scripts[l_index].Dispose();
+                    m_scripts.RemoveAt(l_index);
+                }
+                else
+                    break;
+            }
         }
 
         // Patches