From e1a005703447948da875abb2d3f4b949cbe37c1d Mon Sep 17 00:00:00 2001
From: praydog <praydog@praydog.com>
Date: Fri, 18 Oct 2024 00:53:52 -0700
Subject: [PATCH] Lua: Cache MotionControllerState

---
 include/uevr/API.hpp                        |  4 +++
 lua-api/lib/include/ScriptPrerequisites.hpp |  4 +--
 lua-api/lib/src/ScriptContext.cpp           | 36 +++++++++++++++++----
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/include/uevr/API.hpp b/include/uevr/API.hpp
index eb6e2ecf..3d69c41f 100644
--- a/include/uevr/API.hpp
+++ b/include/uevr/API.hpp
@@ -1565,6 +1565,10 @@ class API {
         }
 
         struct MotionControllerState {
+            static consteval std::string_view internal_name() {
+                return "MotionControllerState";
+            }
+
             inline UEVR_UObjectHookMotionControllerStateHandle to_handle() { return (UEVR_UObjectHookMotionControllerStateHandle)this; }
             inline UEVR_UObjectHookMotionControllerStateHandle to_handle() const { return (UEVR_UObjectHookMotionControllerStateHandle)this; }
 
diff --git a/lua-api/lib/include/ScriptPrerequisites.hpp b/lua-api/lib/include/ScriptPrerequisites.hpp
index 027889d1..432eed29 100644
--- a/lua-api/lib/include/ScriptPrerequisites.hpp
+++ b/lua-api/lib/include/ScriptPrerequisites.hpp
@@ -9,13 +9,13 @@
 
 namespace detail {
 template<typename T>
-concept UObjectBased = std::is_base_of_v<uevr::API::UObject, T>;
+concept CacheablePointer = std::is_base_of_v<uevr::API::UObject, T> || std::is_base_of_v<uevr::API::UObjectHook::MotionControllerState, T>;
 
 constexpr uintptr_t FAKE_OBJECT_ADDR = 12345;
 }
 
 // Object pooling for UObject based pointers using a specialization for sol_lua_push
-template<detail::UObjectBased T>
+template<detail::CacheablePointer T>
 int sol_lua_push(sol::types<T*>, lua_State* l, T* obj) {
     static const std::string TABLE_NAME = std::string{"_sol_lua_push_objects"} + "_" + T::internal_name().data();
 
diff --git a/lua-api/lib/src/ScriptContext.cpp b/lua-api/lib/src/ScriptContext.cpp
index 89ed310e..28c8c3c0 100644
--- a/lua-api/lib/src/ScriptContext.cpp
+++ b/lua-api/lib/src/ScriptContext.cpp
@@ -198,6 +198,10 @@ int ScriptContext::setup_bindings() {
         _sol_lua_push_objects_Property = setmetatable({}, { __mode = "v" })
         _sol_lua_push_objects_Field = setmetatable({}, { __mode = "v" })
         _sol_lua_push_objects_Enum = setmetatable({}, { __mode = "v" })
+
+        -- Not a real UObject but is cacheable
+        _sol_lua_push_objects_MotionControllerState = setmetatable({}, { __mode = "v" })
+
         _sol_lua_push_usertypes = {}
         _sol_lua_push_ref_counts = {}
         _sol_lua_push_ephemeral_counts = {}
@@ -205,13 +209,13 @@ int ScriptContext::setup_bindings() {
 
 
     // templated lambda
-    auto create_uobject_ptr_gc = [&]<detail::UObjectBased T>(T* obj) {
-        m_lua["__UEVRUObjectPtrInternalCreate"] = [this]() -> sol::object {
+    auto create_uobject_ptr_gc = [&]<detail::CacheablePointer T>(T* obj) {
+        m_lua["__UEVRCachePtrInternalCreate"] = [this]() -> sol::object {
             return sol::make_object(m_lua, (T*)detail::FAKE_OBJECT_ADDR);
         };
 
         m_lua.do_string(R"(
-            local fake_obj = __UEVRUObjectPtrInternalCreate()
+            local fake_obj = __UEVRCachePtrInternalCreate()
             local mt = getmetatable(fake_obj)
 
             fake_obj = nil
@@ -222,7 +226,7 @@ int ScriptContext::setup_bindings() {
             end
         )");
 
-        m_lua["__UEVRUObjectPtrInternalCreate"] = sol::make_object(m_lua, sol::nil);
+        m_lua["__UEVRCachePtrInternalCreate"] = sol::make_object(m_lua, sol::nil);
     };
 
     lua::datatypes::bind_xinput(m_lua);
@@ -860,6 +864,8 @@ int ScriptContext::setup_bindings() {
         "set_permanent", &uevr::API::UObjectHook::MotionControllerState::set_permanent
     );
 
+    create_uobject_ptr_gc((API::UObjectHook::MotionControllerState*)nullptr);
+
     m_lua.new_usertype<uevr::API::UObjectHook>("UEVR_UObjectHook",
         "activate", &uevr::API::UObjectHook::activate,
         "exists", &uevr::API::UObjectHook::exists,
@@ -897,8 +903,26 @@ int ScriptContext::setup_bindings() {
 
             return sol::make_object(s, tbl);
         },
-        "get_or_add_motion_controller_state", &uevr::API::UObjectHook::get_or_add_motion_controller_state,
-        "get_motion_controller_state", &uevr::API::UObjectHook::get_motion_controller_state,
+        //"get_or_add_motion_controller_state", &uevr::API::UObjectHook::get_or_add_motion_controller_state,
+        "get_or_add_motion_controller_state", [](sol::this_state s, API::UObject* obj) -> sol::object {
+            auto state = API::UObjectHook::get_or_add_motion_controller_state(obj);
+
+            if (state == nullptr) {
+                return sol::make_object(s, sol::lua_nil);
+            }
+
+            return sol::make_object(s, state);
+        },
+        //"get_motion_controller_state", &uevr::API::UObjectHook::get_motion_controller_state,
+        "get_motion_controller_state", [](sol::this_state s, uevr::API::UObjectHook* self, API::UObject* obj) -> sol::object {
+            auto state = API::UObjectHook::get_motion_controller_state(obj);
+
+            if (state == nullptr) {
+                return sol::make_object(s, sol::lua_nil);
+            }
+
+            return sol::make_object(s, state);
+        },
         "remove_motion_controller_state", &uevr::API::UObjectHook::remove_motion_controller_state
     );