From 668544f62daa52394f2328b48e92ac56cb1afeb0 Mon Sep 17 00:00:00 2001
From: MysterD <myster@d>
Date: Mon, 30 May 2022 00:37:14 -0700
Subject: [PATCH] Clear custom actors on disconnect properly

---
 src/pc/lua/smlua.c                   |  2 ++
 src/pc/lua/utils/smlua_model_utils.c | 23 +++++++++++++++++++++++
 src/pc/lua/utils/smlua_model_utils.h |  1 +
 src/pc/network/network.c             |  1 -
 4 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/pc/lua/smlua.c b/src/pc/lua/smlua.c
index f91e0186..216dff3d 100644
--- a/src/pc/lua/smlua.c
+++ b/src/pc/lua/smlua.c
@@ -5,6 +5,7 @@
 #include "pc/crash_handler.h"
 #include "pc/lua/utils/smlua_text_utils.h"
 #include "pc/lua/utils/smlua_audio_utils.h"
+#include "pc/lua/utils/smlua_model_utils.h"
 #include "pc/djui/djui.h"
 
 lua_State* gLuaState = NULL;
@@ -186,6 +187,7 @@ void smlua_shutdown(void) {
     smlua_cobject_allowlist_shutdown();
     smlua_cpointer_allowlist_shutdown();
     smlua_clear_hooks();
+    smlua_model_util_reset();
     lua_State* L = gLuaState;
     if (L != NULL) {
         lua_close(L);
diff --git a/src/pc/lua/utils/smlua_model_utils.c b/src/pc/lua/utils/smlua_model_utils.c
index be44ef1b..45e8336a 100644
--- a/src/pc/lua/utils/smlua_model_utils.c
+++ b/src/pc/lua/utils/smlua_model_utils.c
@@ -62,6 +62,7 @@ struct ModelUtilsInfo {
     bool permanent;
     bool isDisplayList;
     const void* asset;
+    u8 shouldFreeAsset;
 };
 
 #define UNLOADED_ID 0xFF
@@ -509,6 +510,20 @@ void smlua_model_util_remember(u8 loadedId, UNUSED u8 layer, const void* asset,
     //LOG_INFO("Remember model: %u -> %u", found->extId, loadedId);
 }
 
+void smlua_model_util_reset(void) {
+    smlua_model_util_clear();
+    for (u32 i = 0; i < sCustomModelsCount; i++) {
+        struct ModelUtilsInfo* m = &sCustomModels[i];
+        m->loadedId = UNLOADED_ID;
+        if (m->asset && m->shouldFreeAsset) {
+            free((void*)m->asset);
+            m->asset = NULL;
+        }
+        m->shouldFreeAsset = false;
+    }
+    sCustomModelsCount = 0;
+}
+
 void smlua_model_util_clear(void) {
     // TODO: we need to restore replaced permanent models
     for (int i = 0; i < MAX_CACHED_ASSETS; i++) {
@@ -516,6 +531,11 @@ void smlua_model_util_clear(void) {
         if (m == NULL || m->permanent) { continue; }
         //LOG_INFO("Forget: %u -> %u", m->extId, m->loadedId);
         m->loadedId = UNLOADED_ID;
+        if (m->asset && m->shouldFreeAsset) {
+            free((void*)m->asset);
+            m->asset = NULL;
+        }
+        m->shouldFreeAsset = false;
         sCachedAssets[i] = NULL;
         sCachedAssetTaken[i] = false;
     }
@@ -566,6 +586,7 @@ u8 smlua_model_util_load_with_pool_and_cache_id(enum ModelExtendedId extId, stru
         pool = alloc_only_pool_init(main_pool_available() - sizeof(struct AllocOnlyPool), MEMORY_POOL_LEFT);
         resizePool = true;
     }
+    info->shouldFreeAsset = false;
 
     if (pool != NULL) {
         if (info->isDisplayList) {
@@ -586,6 +607,7 @@ u8 smlua_model_util_load_with_pool_and_cache_id(enum ModelExtendedId extId, stru
         if (info->isDisplayList) {
             const GeoLayout displayListToGeoLayout[] = { GEO_NODE_START(), GEO_DISPLAY_LIST(info->layer, info->asset), GEO_END() };
             info->asset = memcpy(calloc(1, sizeof(displayListToGeoLayout)), displayListToGeoLayout, sizeof(displayListToGeoLayout));
+            info->shouldFreeAsset = true;
             info->isDisplayList = false;
         }
         gLoadedGraphNodes[pickLoadedId] = dynos_geolayout_to_graphnode(info->asset, true);
@@ -634,6 +656,7 @@ u32 smlua_model_util_get_id(const char* name) {
     u16 customIndex = sCustomModelsCount++;
     struct ModelUtilsInfo* info = &sCustomModels[customIndex];
     info->asset = layout;
+    info->shouldFreeAsset = false;
     info->loadedId = UNLOADED_ID;
     info->extId = E_MODEL_MAX + customIndex;
     info->isDisplayList = false;
diff --git a/src/pc/lua/utils/smlua_model_utils.h b/src/pc/lua/utils/smlua_model_utils.h
index aff6a9e7..25187f4f 100644
--- a/src/pc/lua/utils/smlua_model_utils.h
+++ b/src/pc/lua/utils/smlua_model_utils.h
@@ -391,6 +391,7 @@ enum ModelExtendedId {
 };
 
 void smlua_model_util_remember(u8 loadedId, u8 layer, const void* asset, u8 isDisplayList);
+void smlua_model_util_reset(void);
 void smlua_model_util_clear(void);
 u8 smlua_model_util_load_with_pool_and_cache_id(enum ModelExtendedId extId, struct AllocOnlyPool* pool, u8 loadedId);
 u8 smlua_model_util_load_with_pool(enum ModelExtendedId extId, struct AllocOnlyPool* pool);
diff --git a/src/pc/network/network.c b/src/pc/network/network.c
index b4b8f4f4..44839b6e 100644
--- a/src/pc/network/network.c
+++ b/src/pc/network/network.c
@@ -497,7 +497,6 @@ void network_shutdown(bool sendLeaving, bool exiting) {
     mods_clear(&gActiveMods);
     mods_clear(&gRemoteMods);
     smlua_shutdown();
-    smlua_model_util_clear();
     extern s16 gChangeLevel;
     gChangeLevel = LEVEL_CASTLE_GROUNDS;
     network_player_init();