From 2c313791d972bd2c3753e3761df7a60839433e24 Mon Sep 17 00:00:00 2001
From: Deukhoofd <Deukhoofd@gmail.com>
Date: Fri, 14 Jun 2019 17:35:05 +0200
Subject: [PATCH] Ensure memory is handled properly with userdata

---
 src/UserData/UserDataScriptType.hpp |  2 +-
 src/UserData/UserDataStorage.cpp    |  2 +-
 src/UserData/UserDataStorage.hpp    | 13 +++++++++----
 src/UserData/UserDataValue.hpp      |  9 +++++----
 tests/integration/UserData.cpp      |  4 +++-
 5 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/UserData/UserDataScriptType.hpp b/src/UserData/UserDataScriptType.hpp
index 7907e95..d236440 100644
--- a/src/UserData/UserDataScriptType.hpp
+++ b/src/UserData/UserDataScriptType.hpp
@@ -8,7 +8,7 @@
 #include "UserDataStorage.hpp"
 
 class UserDataScriptType : public ScriptType{
-    UserData* _userData;
+    shared_ptr<UserData> _userData;
 public:
     explicit UserDataScriptType(int id) : ScriptType(TypeClass::UserData){
         _userData = UserDataStorage::GetUserDataType(id);
diff --git a/src/UserData/UserDataStorage.cpp b/src/UserData/UserDataStorage.cpp
index 4f3b6df..b008b02 100644
--- a/src/UserData/UserDataStorage.cpp
+++ b/src/UserData/UserDataStorage.cpp
@@ -1,4 +1,4 @@
 
 #include "UserDataStorage.hpp"
 
-std::unordered_map<int, UserData*> UserDataStorage::_userData = {};
\ No newline at end of file
+UserDataStorage::_internalDataStorage UserDataStorage::_internal = UserDataStorage::_internalDataStorage();
diff --git a/src/UserData/UserDataStorage.hpp b/src/UserData/UserDataStorage.hpp
index 3e826b6..0fbc356 100644
--- a/src/UserData/UserDataStorage.hpp
+++ b/src/UserData/UserDataStorage.hpp
@@ -6,14 +6,19 @@
 #include "UserData.hpp"
 
 class UserDataStorage {
-    static std::unordered_map<int, UserData*> _userData;
+    class _internalDataStorage{
+    public:
+        std::unordered_map<int, shared_ptr<UserData>> _userData;
+    };
+    static _internalDataStorage _internal;
+
 public:
     static void RegisterType(int i, UserData* ud){
-        UserDataStorage::_userData.insert({i, ud});
+        UserDataStorage::_internal._userData.insert({i, shared_ptr<UserData>(ud)});
     }
 
-    static UserData* GetUserDataType(int i){
-        return UserDataStorage::_userData[i];
+    static shared_ptr<UserData> GetUserDataType(int i){
+        return UserDataStorage::_internal._userData[i];
     }
 };
 
diff --git a/src/UserData/UserDataValue.hpp b/src/UserData/UserDataValue.hpp
index df56f7f..43243a4 100644
--- a/src/UserData/UserDataValue.hpp
+++ b/src/UserData/UserDataValue.hpp
@@ -1,18 +1,19 @@
 
+
 #ifndef PORYGONLANG_USERDATAVALUE_HPP
 #define PORYGONLANG_USERDATAVALUE_HPP
 
-
+#include <utility>
 #include "../Evaluator/EvalValues/EvalValue.hpp"
 #include "UserData.hpp"
 #include "UserDataStorage.hpp"
 
 class UserDataValue : public EvalValue{
-    UserData* _userData;
+    shared_ptr<UserData> _userData;
     void* _obj;
 public:
-    UserDataValue(UserData* userData, void* obj){
-        _userData = userData;
+    UserDataValue(shared_ptr<UserData> userData, void* obj){
+        _userData = std::move(userData);
         _obj = obj;
     }
 
diff --git a/tests/integration/UserData.cpp b/tests/integration/UserData.cpp
index a69824f..58495aa 100644
--- a/tests/integration/UserData.cpp
+++ b/tests/integration/UserData.cpp
@@ -56,8 +56,10 @@ end
     auto obj = new UserDataTestObject();
     auto parameter = new UserDataValue(HashedString::ConstHash("testObject"), obj);
     script->CallFunction("testFunc", {parameter});
-    REQUIRE(obj->foo == 5000);
     delete script;
+    REQUIRE(obj->foo == 5000);
+    delete obj;
+    delete parameter;
 }