From c1c088118bb03b4922479571896f88e118b7b413 Mon Sep 17 00:00:00 2001
From: maelstrom <maelstrom@maelstrom.dev>
Date: Fri, 13 Jun 2025 01:45:04 +0200
Subject: [PATCH] fix(datatypes): InstanceRef is now shared_ptr and uses the
 same userdata when possible

---
 core/src/datatypes/ref.cpp | 62 ++++++++++++++++++++++++++++++++++----
 core/src/datatypes/ref.h   |  5 ++-
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/core/src/datatypes/ref.cpp b/core/src/datatypes/ref.cpp
index 699243b..20fc26d 100644
--- a/core/src/datatypes/ref.cpp
+++ b/core/src/datatypes/ref.cpp
@@ -12,8 +12,8 @@
 
 TypeMeta::TypeMeta(const InstanceType* instType) : descriptor(&InstanceRef::TYPE), instType(instType) {}
 
-InstanceRef::InstanceRef() {};
-InstanceRef::InstanceRef(std::weak_ptr<Instance> instance) : ref(instance) {};
+InstanceRef::InstanceRef() : ref(nullptr) {};
+InstanceRef::InstanceRef(std::weak_ptr<Instance> instance) : ref(instance.expired() ? nullptr : instance.lock()) {};
 InstanceRef::~InstanceRef() = default;
 
 const TypeDesc InstanceRef::TYPE = {
@@ -27,13 +27,21 @@ const TypeDesc InstanceRef::TYPE = {
 };
 
 const std::string InstanceRef::ToString() const {
-    return ref.expired() ? "" : ref.lock()->name;
+    return ref == nullptr ? "NULL" : ref->name;
+}
+
+InstanceRef::operator std::shared_ptr<Instance>() {
+    return ref;
 }
 
 InstanceRef::operator std::weak_ptr<Instance>() {
     return ref;
 }
 
+bool InstanceRef::operator ==(InstanceRef other) const {
+    return this->ref == other.ref;
+}
+
 // Serialization
 
 void InstanceRef::Serialize(pugi::xml_node node) const {
@@ -50,16 +58,43 @@ static int inst_gc(lua_State*);
 static int inst_index(lua_State*);
 static int inst_newindex(lua_State*);
 static int inst_tostring(lua_State*);
+static int inst_eq(lua_State*);
 static const struct luaL_Reg metatable [] = {
     {"__gc", inst_gc},
     {"__index", inst_index},
     {"__newindex", inst_newindex},
     {"__tostring", inst_tostring},
+    {"__eq", inst_eq},
     {NULL, NULL} /* end of array */
 };
 
 void InstanceRef::PushLuaValue(lua_State* L) const {
-    if (ref.expired()) return lua_pushnil(L);
+    if (ref == nullptr) return lua_pushnil(L);
+
+    // Get or create InstanceRef table
+    lua_getfield(L, LUA_REGISTRYINDEX, "__instances");
+    if (lua_isnil(L, -1)) {
+        lua_pop(L, 1);
+        lua_newtable(L);
+        
+        // Set metatable
+        lua_newtable(L);
+        lua_pushstring(L, "kv");
+        lua_setfield(L, -2, "__mode");
+        lua_setmetatable(L, -2);
+
+        lua_pushvalue(L, -1);
+        lua_setfield(L, LUA_REGISTRYINDEX, "__instances");
+    }
+
+    // Check if value already exists, and if so, return that instead
+    lua_pushlightuserdata(L, ref.get());
+    lua_rawget(L, -2);
+    if (!lua_isnil(L, -1)) {
+        lua_remove(L, -2); // Remove __instances
+        return;
+    }
+    lua_pop(L, 1);
 
     int n = lua_gettop(L);
 
@@ -72,8 +107,13 @@ void InstanceRef::PushLuaValue(lua_State* L) const {
     // Create the instance's metatable
     luaL_newmetatable(L, "__mt_instance");
     luaL_register(L, NULL, metatable);
-
     lua_setmetatable(L, n+1);
+
+    // Add instance to __instances
+    lua_pushlightuserdata(L, ref.get());
+    lua_pushvalue(L, -2); // Push userdata
+    lua_rawset(L, -4); // Put into __instance
+    lua_remove(L, -2); // Remove __instance
 }
 
 result<Variant, LuaCastError> InstanceRef::FromLuaValue(lua_State* L, int idx) {
@@ -152,4 +192,14 @@ static int inst_tostring(lua_State* L) {
     lua_pushstring(L, inst->name.c_str());
 
     return 1;
-}
\ No newline at end of file
+}
+
+static int inst_eq(lua_State* L) {
+    auto userdata = (std::shared_ptr<Instance>**)lua_touserdata(L, 1);
+    std::shared_ptr<Instance> inst = **userdata;
+    auto userdata2 = (std::shared_ptr<Instance>**)luaL_checkudata(L, 2, "__mt_instance");
+    std::shared_ptr<Instance> inst2 = **userdata2;
+
+    lua_pushboolean(L, inst == inst2);
+    return 1;
+}
diff --git a/core/src/datatypes/ref.h b/core/src/datatypes/ref.h
index 6e22878..afca5a1 100644
--- a/core/src/datatypes/ref.h
+++ b/core/src/datatypes/ref.h
@@ -7,7 +7,7 @@
 class Instance;
 
 class InstanceRef {
-    std::weak_ptr<Instance> ref;
+    std::shared_ptr<Instance> ref;
 public:
     InstanceRef();
     InstanceRef(std::weak_ptr<Instance>);
@@ -15,6 +15,7 @@ public:
 
     static const TypeDesc TYPE;
 
+    operator std::shared_ptr<Instance>();
     operator std::weak_ptr<Instance>();
 
     virtual const std::string ToString() const;
@@ -22,4 +23,6 @@ public:
     virtual void PushLuaValue(lua_State*) const;
     static result<InstanceRef, DataParseError> Deserialize(pugi::xml_node node);
     static result<Variant, LuaCastError> FromLuaValue(lua_State*, int idx);
+
+    bool operator ==(InstanceRef) const;
 };
\ No newline at end of file