From d6e243c4a762a989373e5e1379c13aea72298185 Mon Sep 17 00:00:00 2001
From: Isaac0-dev <62234577+Isaac0-dev@users.noreply.github.com>
Date: Tue, 4 Mar 2025 15:09:09 +1000
Subject: [PATCH] add obj_get_model_id_extended, remember regular model ids

---
 autogen/lua_definitions/functions.lua |  7 ++++
 data/dynos_mgr_models.cpp             | 11 +++--
 docs/lua/functions-6.md               | 23 +++++++++++
 docs/lua/functions.md                 |  1 +
 include/model_ids.h                   |  2 +
 src/pc/lua/smlua_functions_autogen.c  | 18 ++++++++
 src/pc/lua/utils/smlua_model_utils.c  | 59 ++++++++++++++++++++++++++-
 src/pc/lua/utils/smlua_model_utils.h  |  3 ++
 src/pc/lua/utils/smlua_obj_utils.c    | 11 +++--
 src/pc/lua/utils/smlua_obj_utils.h    |  2 +
 10 files changed, 130 insertions(+), 7 deletions(-)

diff --git a/autogen/lua_definitions/functions.lua b/autogen/lua_definitions/functions.lua
index 08651ea69..ef0fbccbe 100644
--- a/autogen/lua_definitions/functions.lua
+++ b/autogen/lua_definitions/functions.lua
@@ -10076,6 +10076,13 @@ function obj_get_first_with_behavior_id_and_field_s32(behaviorId, fieldIndex, va
     -- ...
 end
 
+--- @param o Object
+--- @return ModelExtendedId
+--- Returns an object's extended model id
+function obj_get_model_id_extended(o)
+    -- ...
+end
+
 --- @param o Object
 --- @param behaviorId BehaviorId
 --- @return Object
diff --git a/data/dynos_mgr_models.cpp b/data/dynos_mgr_models.cpp
index 9ceec81cd..90386fac2 100644
--- a/data/dynos_mgr_models.cpp
+++ b/data/dynos_mgr_models.cpp
@@ -6,10 +6,9 @@ extern "C" {
 #include "engine/geo_layout.h"
 #include "engine/graph_node.h"
 #include "model_ids.h"
+#include "pc/lua/utils/smlua_model_utils.h"
 }
 
-#define VANILLA_ID_END 255
-
 enum ModelLoadType {
     MLT_GEO,
     MLT_DL,
@@ -58,7 +57,7 @@ void DynOS_Model_Dump() {
     }
 }
 
-struct GraphNode* DynOS_Model_LoadCommon(u32* aId, enum ModelPool aModelPool, void* aAsset, u8 aLayer, struct GraphNode* aGraphNode, bool aDeDuplicate, enum ModelLoadType mlt) {
+static struct GraphNode* DynOS_Model_LoadCommonInternal(u32* aId, enum ModelPool aModelPool, void* aAsset, u8 aLayer, struct GraphNode* aGraphNode, bool aDeDuplicate, enum ModelLoadType mlt) {
     // sanity check pool
     if (aModelPool >= MODEL_POOL_MAX) { return NULL; }
 
@@ -124,6 +123,12 @@ struct GraphNode* DynOS_Model_LoadCommon(u32* aId, enum ModelPool aModelPool, vo
     return node;
 }
 
+static struct GraphNode* DynOS_Model_LoadCommon(u32* aId, enum ModelPool aModelPool, void* aAsset, u8 aLayer, struct GraphNode* aGraphNode, bool aDeDuplicate, enum ModelLoadType mlt) {
+    struct GraphNode* node = DynOS_Model_LoadCommonInternal(aId, aModelPool, aAsset, aLayer, aGraphNode, aDeDuplicate, mlt);
+    smlua_model_util_register_model_id(*aId, aAsset);
+    return node;
+}
+
 struct GraphNode* DynOS_Model_LoadGeo(u32* aId, enum ModelPool aModelPool, void* aAsset, bool aDeDuplicate) {
     return DynOS_Model_LoadCommon(aId, aModelPool, aAsset, 0, NULL, aDeDuplicate, MLT_GEO);
 }
diff --git a/docs/lua/functions-6.md b/docs/lua/functions-6.md
index ec51fa095..13e0bf7f3 100644
--- a/docs/lua/functions-6.md
+++ b/docs/lua/functions-6.md
@@ -2215,6 +2215,29 @@ Gets the first object loaded with `behaviorId` and object signed 32-bit integer
 
 <br />
 
+## [obj_get_model_id_extended](#obj_get_model_id_extended)
+
+### Description
+Returns an object's extended model id
+
+### Lua Example
+`local enumValue = obj_get_model_id_extended(o)`
+
+### Parameters
+| Field | Type |
+| ----- | ---- |
+| o | [Object](structs.md#Object) |
+
+### Returns
+[enum ModelExtendedId](constants.md#enum-ModelExtendedId)
+
+### C Prototype
+`enum ModelExtendedId obj_get_model_id_extended(struct Object *o);`
+
+[:arrow_up_small:](#)
+
+<br />
+
 ## [obj_get_nearest_object_with_behavior_id](#obj_get_nearest_object_with_behavior_id)
 
 ### Description
diff --git a/docs/lua/functions.md b/docs/lua/functions.md
index fe1ab79c3..fd8d8b7b0 100644
--- a/docs/lua/functions.md
+++ b/docs/lua/functions.md
@@ -1847,6 +1847,7 @@
    - [obj_get_first_with_behavior_id](functions-6.md#obj_get_first_with_behavior_id)
    - [obj_get_first_with_behavior_id_and_field_f32](functions-6.md#obj_get_first_with_behavior_id_and_field_f32)
    - [obj_get_first_with_behavior_id_and_field_s32](functions-6.md#obj_get_first_with_behavior_id_and_field_s32)
+   - [obj_get_model_id_extended](functions-6.md#obj_get_model_id_extended)
    - [obj_get_nearest_object_with_behavior_id](functions-6.md#obj_get_nearest_object_with_behavior_id)
    - [obj_get_next](functions-6.md#obj_get_next)
    - [obj_get_next_with_same_behavior_id](functions-6.md#obj_get_next_with_same_behavior_id)
diff --git a/include/model_ids.h b/include/model_ids.h
index e7e900e2f..8cd85e7fc 100644
--- a/include/model_ids.h
+++ b/include/model_ids.h
@@ -19,6 +19,8 @@
 #define COIN_FORMATION_FLAG_ARROW     (1 << 2)
 #define COIN_FORMATION_FLAG_FLYING    (1 << 4)
 
+#define VANILLA_ID_END 255
+
 #define MODEL_NONE                        0x00
 
 /* Global models that are loaded for every level */
diff --git a/src/pc/lua/smlua_functions_autogen.c b/src/pc/lua/smlua_functions_autogen.c
index a64cf1ac2..f3b781eb3 100644
--- a/src/pc/lua/smlua_functions_autogen.c
+++ b/src/pc/lua/smlua_functions_autogen.c
@@ -30547,6 +30547,23 @@ int smlua_func_obj_get_first_with_behavior_id_and_field_s32(lua_State* L) {
     return 1;
 }
 
+int smlua_func_obj_get_model_id_extended(lua_State* L) {
+    if (L == NULL) { return 0; }
+
+    int top = lua_gettop(L);
+    if (top != 1) {
+        LOG_LUA_LINE("Improper param count for '%s': Expected %u, Received %u", "obj_get_model_id_extended", 1, top);
+        return 0;
+    }
+
+    struct Object* o = (struct Object*)smlua_to_cobject(L, 1, LOT_OBJECT);
+    if (!gSmLuaConvertSuccess) { LOG_LUA("Failed to convert parameter %u for function '%s'", 1, "obj_get_model_id_extended"); return 0; }
+
+    lua_pushinteger(L, obj_get_model_id_extended(o));
+
+    return 1;
+}
+
 int smlua_func_obj_get_nearest_object_with_behavior_id(lua_State* L) {
     if (L == NULL) { return 0; }
 
@@ -33857,6 +33874,7 @@ void smlua_bind_functions_autogen(void) {
     smlua_bind_function(L, "obj_get_first_with_behavior_id", smlua_func_obj_get_first_with_behavior_id);
     smlua_bind_function(L, "obj_get_first_with_behavior_id_and_field_f32", smlua_func_obj_get_first_with_behavior_id_and_field_f32);
     smlua_bind_function(L, "obj_get_first_with_behavior_id_and_field_s32", smlua_func_obj_get_first_with_behavior_id_and_field_s32);
+    smlua_bind_function(L, "obj_get_model_id_extended", smlua_func_obj_get_model_id_extended);
     smlua_bind_function(L, "obj_get_nearest_object_with_behavior_id", smlua_func_obj_get_nearest_object_with_behavior_id);
     smlua_bind_function(L, "obj_get_next", smlua_func_obj_get_next);
     smlua_bind_function(L, "obj_get_next_with_same_behavior_id", smlua_func_obj_get_next_with_same_behavior_id);
diff --git a/src/pc/lua/utils/smlua_model_utils.c b/src/pc/lua/utils/smlua_model_utils.c
index c94375d70..759f4d227 100644
--- a/src/pc/lua/utils/smlua_model_utils.c
+++ b/src/pc/lua/utils/smlua_model_utils.c
@@ -486,8 +486,10 @@ void smlua_model_util_store_in_slot(u32 slot, const char* name) {
     dynos_model_overwrite_slot(slot, loadedId);
 }
 
+// Takes an extended model id, loads the model, and returns the regular model id
 u16 smlua_model_util_load(enum ModelExtendedId extId) {
-    if ((u32)extId >= (u32)E_MODEL_MAX + (u32)sCustomModelsCount) { extId = E_MODEL_ERROR_MODEL; }
+    if (extId == E_MODEL_NONE) { return MODEL_NONE; }
+    if (!extId || (u32)extId >= (u32)E_MODEL_MAX + (u32)sCustomModelsCount) { extId = E_MODEL_ERROR_MODEL; }
 
     struct ModelUtilsInfo* info = (extId < E_MODEL_MAX)
         ? &sModels[extId]
@@ -502,6 +504,61 @@ u16 smlua_model_util_load(enum ModelExtendedId extId) {
     return (u16)id;
 }
 
+// Links the regular model id created by DynOS to our models list
+void smlua_model_util_register_model_id(u32 id, const void *asset) {
+    if (id < VANILLA_ID_END) {
+        for (u32 i = 0; i < E_MODEL_MAX; i++) {
+            struct ModelUtilsInfo* m = &sModels[i];
+            if (m->asset == asset) {
+                m->loadedId = id;
+                return;
+            }
+        }
+    } else {
+        for (u32 i = 0; i < sCustomModelsCount; i++) {
+            struct ModelUtilsInfo* m = &sCustomModels[i];
+            if (m->asset == asset) {
+                m->loadedId = id;
+                return;
+            }
+        }
+    }
+}
+
+// Translates an extended model id to a regular model id
+u16 smlua_model_util_ext_id_to_id(enum ModelExtendedId extId) {
+    if (extId == E_MODEL_NONE) { return MODEL_NONE; }
+    if ((u32)extId >= (u32)E_MODEL_MAX + (u32)sCustomModelsCount) { return MODEL_ERROR_MODEL; }
+
+    struct ModelUtilsInfo* info = (extId < E_MODEL_MAX)
+        ? &sModels[extId]
+        : &sCustomModels[extId - E_MODEL_MAX];
+    return info->loadedId != UNLOADED_ID ? info->loadedId : MODEL_ERROR_MODEL;
+}
+
+// Translates a regular model id to an extended model id
+enum ModelExtendedId smlua_model_util_id_to_ext_id(u16 id) {
+    if (!id) { return E_MODEL_NONE; }
+
+    // Check built-in models
+    for (u32 i = 0; i < E_MODEL_MAX; i++) {
+        struct ModelUtilsInfo* m = &sModels[i];
+        if (m->loadedId == id) {
+            return m->extId;
+        }
+    }
+
+    // Check custom models
+    for (u32 i = 0; i < sCustomModelsCount; i++) {
+        struct ModelUtilsInfo* m = &sCustomModels[i];
+        if (m->loadedId == id) {
+            return m->extId;
+        }
+    }
+
+    return E_MODEL_ERROR_MODEL;
+}
+
 enum ModelExtendedId smlua_model_util_get_id(const char* name) {
     // Find geolayout
     const void* asset = dynos_geolayout_get(name);
diff --git a/src/pc/lua/utils/smlua_model_utils.h b/src/pc/lua/utils/smlua_model_utils.h
index 13e06f68a..3e3a955d2 100644
--- a/src/pc/lua/utils/smlua_model_utils.h
+++ b/src/pc/lua/utils/smlua_model_utils.h
@@ -399,6 +399,9 @@ void smlua_model_util_initialize(void);
 void smlua_model_util_clear(void);
 void smlua_model_util_store_in_slot(u32 slot, const char* name);
 u16 smlua_model_util_load(enum ModelExtendedId extId);
+void smlua_model_util_register_model_id(u32 id, const void *asset);
+u16 smlua_model_util_ext_id_to_id(enum ModelExtendedId extId);
+enum ModelExtendedId smlua_model_util_id_to_ext_id(u16 id);
 /* |description|Gets the extended model ID for the `name` of a `GeoLayout`|descriptionEnd| */
 enum ModelExtendedId smlua_model_util_get_id(const char* name);
 
diff --git a/src/pc/lua/utils/smlua_obj_utils.c b/src/pc/lua/utils/smlua_obj_utils.c
index 3f9e63db3..ca651591e 100644
--- a/src/pc/lua/utils/smlua_obj_utils.c
+++ b/src/pc/lua/utils/smlua_obj_utils.c
@@ -92,9 +92,14 @@ s32 obj_has_behavior_id(struct Object *o, enum BehaviorId behaviorId) {
 
 s32 obj_has_model_extended(struct Object *o, enum ModelExtendedId modelId) {
     if (!o) { return 0; }
-    u16 slot = smlua_model_util_load(modelId);
-    struct GraphNode *model = dynos_model_get_geo(slot);
-    return o->header.gfx.sharedChild == model;
+    if (!o->header.gfx.sharedChild && modelId == E_MODEL_NONE) { return 1; }
+    return dynos_model_get_id_from_graph_node(o->header.gfx.sharedChild) == smlua_model_util_ext_id_to_id(modelId);
+}
+
+enum ModelExtendedId obj_get_model_id_extended(struct Object *o) {
+    if (!o) { return E_MODEL_NONE; }
+    if (!o->header.gfx.sharedChild) { return E_MODEL_NONE; }
+    return smlua_model_util_id_to_ext_id(dynos_model_get_id_from_graph_node(o->header.gfx.sharedChild));
 }
 
 void obj_set_model_extended(struct Object *o, enum ModelExtendedId modelId) {
diff --git a/src/pc/lua/utils/smlua_obj_utils.h b/src/pc/lua/utils/smlua_obj_utils.h
index a36ed4d9b..c833ffe2b 100644
--- a/src/pc/lua/utils/smlua_obj_utils.h
+++ b/src/pc/lua/utils/smlua_obj_utils.h
@@ -20,6 +20,8 @@ struct Object* spawn_non_sync_object(enum BehaviorId behaviorId, enum ModelExten
 s32 obj_has_behavior_id(struct Object *o, enum BehaviorId behaviorId);
 /* |description|Checks if an object's model is equal to `modelId`|descriptionEnd| */
 s32 obj_has_model_extended(struct Object *o, enum ModelExtendedId modelId);
+/* |description|Returns an object's extended model id|descriptionEnd| */
+enum ModelExtendedId obj_get_model_id_extended(struct Object *o);
 /* |description|Sets an object's model to `modelId`|descriptionEnd| */
 void obj_set_model_extended(struct Object *o, enum ModelExtendedId modelId);