From eaea27659753d89fe6f51d18af9e85db4a5a74f1 Mon Sep 17 00:00:00 2001
From: praydog <praydog@praydog.com>
Date: Tue, 5 Mar 2024 21:11:07 -0800
Subject: [PATCH] Lua: Stability improvements on script reset

---
 lua-api/lib/include/ScriptContext.hpp |  8 +++++---
 lua-api/lib/src/ScriptContext.cpp     | 27 +++++++++++++++++----------
 src/mods/LuaLoader.cpp                |  4 +++-
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/lua-api/lib/include/ScriptContext.hpp b/lua-api/lib/include/ScriptContext.hpp
index 16dbbe24..ebe18e07 100644
--- a/lua-api/lib/include/ScriptContext.hpp
+++ b/lua-api/lib/include/ScriptContext.hpp
@@ -12,11 +12,10 @@ namespace uevr {
 class ScriptContext : public std::enable_shared_from_this<ScriptContext> {
 public:
     static std::shared_ptr<ScriptContext> create(lua_State* l, UEVR_PluginInitializeParam* param = nullptr) {
-        return std::make_shared<ScriptContext>(l, param);
+        return std::shared_ptr<ScriptContext>(new ScriptContext(l, param));
     }
 
-    ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param = nullptr);
-
+    ScriptContext() = delete;
     virtual ~ScriptContext();
 
     int setup_bindings();
@@ -85,6 +84,9 @@ class ScriptContext : public std::enable_shared_from_this<ScriptContext> {
     }
 
 private:
+    // Private constructor to prevent direct instantiation
+    ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param = nullptr);
+
     std::vector<void*> m_callbacks_to_remove{};
 
     sol::state_view m_lua;
diff --git a/lua-api/lib/src/ScriptContext.cpp b/lua-api/lib/src/ScriptContext.cpp
index 78e6604b..9a585115 100644
--- a/lua-api/lib/src/ScriptContext.cpp
+++ b/lua-api/lib/src/ScriptContext.cpp
@@ -11,11 +11,11 @@
 namespace uevr {
 class ScriptContexts {
 public:
-    void add(std::shared_ptr<ScriptContext> ctx) {
+    void add(ScriptContext* ctx) {
         std::scoped_lock _{mtx};
 
         // Check if the context is already in the list
-        for (auto& c : list) {
+        for (ScriptContext* c : list) {
             if (c == ctx) {
                 return;
             }
@@ -24,29 +24,36 @@ class ScriptContexts {
         list.push_back(ctx);
     }
 
-    void remove(std::shared_ptr<ScriptContext> ctx) {
+    void remove(ScriptContext* ctx) {
         std::scoped_lock _{mtx};
-        list.erase(std::remove(list.begin(), list.end(), ctx), list.end());
+
+        ScriptContext::log("Removing context from list");
+        std::erase_if(list, [ctx](ScriptContext* c) {
+            return c == ctx;
+        });
+        ScriptContext::log(std::format("New context count: {}", list.size()));
     }
 
     template<typename T>
     void for_each(T&& fn) {
         std::scoped_lock _{mtx};
         for (auto& ctx : list) {
-            fn(ctx);
+            fn(ctx->shared_from_this());
         }
     }
 
 private:
-    std::vector<std::shared_ptr<ScriptContext>> list{};
+    std::vector<ScriptContext*> list{};
     std::mutex mtx{};
 } g_contexts{};
 
 ScriptContext::ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param) 
-    : m_lua{l} 
+    : m_lua{l}
 {
     std::scoped_lock _{m_mtx};
 
+    g_contexts.add(this);
+
     if (param != nullptr) {
         m_plugin_initialize_param = param;
         uevr::API::initialize(m_plugin_initialize_param);
@@ -65,6 +72,7 @@ ScriptContext::ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param)
 
 ScriptContext::~ScriptContext() {
     std::scoped_lock _{m_mtx};
+    ScriptContext::log("ScriptContext destructor called");
 
     // TODO: this probably does not support multiple states
     if (m_plugin_initialize_param != nullptr) {
@@ -75,11 +83,12 @@ ScriptContext::~ScriptContext() {
         m_callbacks_to_remove.clear();
     }
 
-    g_contexts.remove(shared_from_this());
+    g_contexts.remove(this);
 }
 
 void ScriptContext::log(const std::string& message) {
     std::cout << "[LuaVR] " << message << std::endl;
+    API::get()->log_info("[LuaVR] %s", message.c_str());
 }
 
 void ScriptContext::test_function() {
@@ -149,8 +158,6 @@ void ScriptContext::setup_callback_bindings() {
 }
 
 int ScriptContext::setup_bindings() {
-    g_contexts.add(shared_from_this());
-
     m_lua.registry()["uevr_context"] = this;
 
     m_lua.set_function("test_function", ScriptContext::test_function);
diff --git a/src/mods/LuaLoader.cpp b/src/mods/LuaLoader.cpp
index 89e8f744..5aad714a 100644
--- a/src/mods/LuaLoader.cpp
+++ b/src/mods/LuaLoader.cpp
@@ -226,7 +226,7 @@ void LuaLoader::reset_scripts() {
         m_last_script_error.clear();
     }
 
-    if (!m_states.empty()) {
+    if (m_main_state != nullptr) {
         /*auto& mods = g_framework->get_mods()->get_mods();
 
         for (auto& mod : mods) {
@@ -239,6 +239,8 @@ void LuaLoader::reset_scripts() {
     m_main_state.reset();
     m_states.clear();
 
+    spdlog::info("[LuaLoader] Destroyed all Lua states.");
+
     m_main_state = std::make_shared<ScriptState>(make_gc_data(), &g_plugin_initialize_param, true);
     m_states.insert(m_states.begin(), m_main_state);