From 07de3cb17211eed31c338964d8c6c0daa841155a Mon Sep 17 00:00:00 2001
From: Garrett Cox <garrettjcox@gmail.com>
Date: Sun, 28 Jan 2024 22:02:30 +0000
Subject: [PATCH] Support for hook unregistration (#109)

---
 mm/2s2h/BenGui/BenMenuBar.cpp                 |  7 +++--
 mm/2s2h/Enhancements/Enhancements.cpp         | 18 ++++++++----
 mm/2s2h/Enhancements/Enhancements.h           |  1 +
 .../GameInteractor/GameInteractor.h           | 29 ++++++++++++++++---
 4 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/mm/2s2h/BenGui/BenMenuBar.cpp b/mm/2s2h/BenGui/BenMenuBar.cpp
index 54f2e1c67..e373981d0 100644
--- a/mm/2s2h/BenGui/BenMenuBar.cpp
+++ b/mm/2s2h/BenGui/BenMenuBar.cpp
@@ -7,6 +7,7 @@
 #include <unordered_map>
 #include <string>
 #include "z64.h"
+#include "2s2h/Enhancements/Enhancements.h"
 
 extern bool ShouldClearTextureCacheAtEndOfFrame;
 
@@ -286,9 +287,11 @@ void DrawDeveloperToolsMenu() {
         });
         UIWidgets::CVarCheckbox("No Clip", "gDeveloperTools.NoClip");
         UIWidgets::CVarCheckbox("Better Map Select", "gDeveloperTools.BetterMapSelect.Enabled");
-        UIWidgets::CVarCheckbox("Moon Jump on L", "gDeveloperTools.MoonJumpOnL", {
+        if (UIWidgets::CVarCheckbox("Moon Jump on L", "gDeveloperTools.MoonJumpOnL", {
             .tooltip = "Holding L makes you float into the air"
-        });
+        })) {
+            RegisterMoonJumpOnL();
+        }
         if (gPlayState != NULL) {
             ImGui::Separator();
             UIWidgets::Checkbox("Frame Advance", (bool*)&gPlayState->frameAdvCtx.enabled, {
diff --git a/mm/2s2h/Enhancements/Enhancements.cpp b/mm/2s2h/Enhancements/Enhancements.cpp
index 82aedcf1a..0000f37f1 100644
--- a/mm/2s2h/Enhancements/Enhancements.cpp
+++ b/mm/2s2h/Enhancements/Enhancements.cpp
@@ -10,18 +10,24 @@ extern "C" {
 extern PlayState* gPlayState;
 }
 
+static uint32_t moonJumpOnLGameStateUpdateHookId = 0;
 void RegisterMoonJumpOnL() {
-    GameInteractor::Instance->RegisterGameHook<GameInteractor::OnGameStateUpdate>([]() {
-        if (!gPlayState) return;
-        
-        if (CVarGetInteger("gDeveloperTools.MoonJumpOnL", 0)) {
+    if (moonJumpOnLGameStateUpdateHookId) {
+        GameInteractor::Instance->UnregisterGameHook<GameInteractor::OnGameStateUpdate>(moonJumpOnLGameStateUpdateHookId);
+        moonJumpOnLGameStateUpdateHookId = 0;
+    }
+
+    if (CVarGetInteger("gDeveloperTools.MoonJumpOnL", 0)) {
+        moonJumpOnLGameStateUpdateHookId = GameInteractor::Instance->RegisterGameHook<GameInteractor::OnGameStateUpdate>([]() {
+            if (!gPlayState) return;
+
             Player* player = GET_PLAYER(gPlayState);
 
             if (CHECK_BTN_ANY(gPlayState->state.input[0].cur.button, BTN_L)) {
                 player->actor.velocity.y = 6.34375f;
             }
-        }
-    });
+        });
+    }
 }
 
 void InitEnhancements() {
diff --git a/mm/2s2h/Enhancements/Enhancements.h b/mm/2s2h/Enhancements/Enhancements.h
index b716a6e5a..9565c3625 100644
--- a/mm/2s2h/Enhancements/Enhancements.h
+++ b/mm/2s2h/Enhancements/Enhancements.h
@@ -3,6 +3,7 @@ extern "C" {
 #endif
 
 void InitEnhancements();
+void RegisterMoonJumpOnL();
 
 #ifdef __cplusplus
 }
diff --git a/mm/2s2h/Enhancements/GameInteractor/GameInteractor.h b/mm/2s2h/Enhancements/GameInteractor/GameInteractor.h
index a9828765e..1947deb74 100644
--- a/mm/2s2h/Enhancements/GameInteractor/GameInteractor.h
+++ b/mm/2s2h/Enhancements/GameInteractor/GameInteractor.h
@@ -13,6 +13,8 @@ extern "C" {
 
 #include <vector>
 #include <functional>
+#include <unordered_map>
+#include <cstdint>
 
 #define DEFINE_HOOK(name, type)         \
     struct name {                       \
@@ -29,11 +31,30 @@ public:
     };
 
     // Game Hooks
-    template <typename H> struct RegisteredGameHooks { inline static std::vector<typename H::fn> functions; };
-    template <typename H> void RegisterGameHook(typename H::fn h) { RegisteredGameHooks<H>::functions.push_back(h); }
+    uint32_t nextHookId = 1;
+    template <typename H> struct RegisteredGameHooks { inline static std::unordered_map<uint32_t, typename H::fn> functions; };
+    template <typename H> struct HooksToUnregister { inline static std::vector<uint32_t> hooks; };
+    template <typename H> uint32_t RegisterGameHook(typename H::fn h) {
+        // Ensure hook id is unique and not 0, which is reserved for invalid hooks
+        if (this->nextHookId == 0 || this->nextHookId >= UINT32_MAX) this->nextHookId = 1;
+        while (RegisteredGameHooks<H>::functions.find(this->nextHookId) != RegisteredGameHooks<H>::functions.end()) {
+            this->nextHookId++;
+        }
+
+        RegisteredGameHooks<H>::functions[this->nextHookId] = h;
+        return this->nextHookId++;
+    }
+    template <typename H> void UnregisterGameHook(uint32_t id) {
+        HooksToUnregister<H>::hooks.push_back(id);
+    }
+
     template <typename H, typename... Args> void ExecuteHooks(Args&&... args) {
-        for (auto& fn : RegisteredGameHooks<H>::functions) {
-            fn(std::forward<Args>(args)...);
+        for (auto& hookId : HooksToUnregister<H>::hooks) {
+            RegisteredGameHooks<H>::functions.erase(hookId);
+        }
+        HooksToUnregister<H>::hooks.clear();
+        for (auto& hook : RegisteredGameHooks<H>::functions) {
+            hook.second(std::forward<Args>(args)...);
         }
     }