From 3e3b39f68083af24de32a6f3723ff8d8e4cf952f Mon Sep 17 00:00:00 2001
From: DashingCat <65317616+DashingCat@users.noreply.github.com>
Date: Sun, 3 Nov 2024 19:59:33 +0100
Subject: [PATCH] Prevent leftovers MenuButtons when rebuilding menu

Close issues #430 and #766
---
 Celeste.Mod.mm/Mod/Core/CoreModuleSettings.cs |  6 +++---
 Celeste.Mod.mm/Mod/Everest/Everest.Loader.cs  |  2 +-
 .../Mod/UI/OuiDependencyDownloader.cs         |  2 +-
 Celeste.Mod.mm/Patches/OuiMainMenu.cs         | 21 ++++++++++++++++++-
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/Celeste.Mod.mm/Mod/Core/CoreModuleSettings.cs b/Celeste.Mod.mm/Mod/Core/CoreModuleSettings.cs
index 32b404077..ef53ff6d6 100644
--- a/Celeste.Mod.mm/Mod/Core/CoreModuleSettings.cs
+++ b/Celeste.Mod.mm/Mod/Core/CoreModuleSettings.cs
@@ -79,7 +79,7 @@ public bool DebugModeInEverest {
                         Engine.Commands.Enabled = false;
                 }
 
-                ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.RebuildMainAndTitle();
+                ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.NeedsRebuild();
             }
         }
 
@@ -236,7 +236,7 @@ public string MainMenuMode {
                 _MainMenuMode = value;
                 if (value != originalValue) {
                     // the main menu mode was changed; rebuild the main menu
-                    ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.RebuildMainAndTitle();
+                    ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.NeedsRebuild();
                 }
             }
         }
@@ -265,7 +265,7 @@ public bool WarnOnEverestYamlErrors {
                 _WarnOnEverestYamlErrors = value;
 
                 // rebuild the main menu to make sure we show/hide the yaml error notice.
-                ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.RebuildMainAndTitle();
+                ((patch_OuiMainMenu) (Engine.Scene as Overworld)?.GetUI<OuiMainMenu>())?.NeedsRebuild();
             }
         }
 
diff --git a/Celeste.Mod.mm/Mod/Everest/Everest.Loader.cs b/Celeste.Mod.mm/Mod/Everest/Everest.Loader.cs
index 72a1f7eff..a2018fe62 100644
--- a/Celeste.Mod.mm/Mod/Everest/Everest.Loader.cs
+++ b/Celeste.Mod.mm/Mod/Everest/Everest.Loader.cs
@@ -237,7 +237,7 @@ private static void LoadAutoUpdated(object source, FileSystemEventArgs e) {
                         LoadDir(e.FullPath);
                     else if (e.FullPath.EndsWith(".zip"))
                         LoadZip(e.FullPath);
-                    ((patch_OuiMainMenu) (AssetReloadHelper.ReturnToScene as Overworld)?.GetUI<OuiMainMenu>())?.RebuildMainAndTitle();
+                    ((patch_OuiMainMenu) (AssetReloadHelper.ReturnToScene as Overworld)?.GetUI<OuiMainMenu>())?.NeedsRebuild();
                 })));
             }
 
diff --git a/Celeste.Mod.mm/Mod/UI/OuiDependencyDownloader.cs b/Celeste.Mod.mm/Mod/UI/OuiDependencyDownloader.cs
index 284a8f856..e2e2fffcc 100755
--- a/Celeste.Mod.mm/Mod/UI/OuiDependencyDownloader.cs
+++ b/Celeste.Mod.mm/Mod/UI/OuiDependencyDownloader.cs
@@ -523,7 +523,7 @@ public override void Update() {
         public void Exit() {
             task = null;
             Lines.Clear();
-            MainThreadHelper.Schedule(() => ((patch_OuiMainMenu) Overworld.GetUI<OuiMainMenu>())?.RebuildMainAndTitle());
+            MainThreadHelper.Schedule(() => ((patch_OuiMainMenu) Overworld.GetUI<OuiMainMenu>())?.NeedsRebuild());
             Audio.Play(SFX.ui_main_button_back);
             Overworld.Goto<OuiModOptions>();
         }
diff --git a/Celeste.Mod.mm/Patches/OuiMainMenu.cs b/Celeste.Mod.mm/Patches/OuiMainMenu.cs
index 4116d8d95..dc0a25dc1 100644
--- a/Celeste.Mod.mm/Patches/OuiMainMenu.cs
+++ b/Celeste.Mod.mm/Patches/OuiMainMenu.cs
@@ -18,6 +18,8 @@ class patch_OuiMainMenu : OuiMainMenu {
         public List<MenuButton> Buttons => buttons;
         private MainMenuClimb climbButton;
 
+        private bool needsRebuild = false;
+
         public extern void orig_CreateButtons();
         public new void CreateButtons() {
             orig_CreateButtons();
@@ -27,7 +29,22 @@ class patch_OuiMainMenu : OuiMainMenu {
             UpdateLayout();
         }
 
-        public void RebuildMainAndTitle() {
+        public void NeedsRebuild() {
+            needsRebuild = true;
+        }
+
+        public extern void orig_Update();
+        public new void Update() {
+            // rebuild only on update to prevent multiple rebuilds per frame
+            if (needsRebuild) {
+                RebuildMainAndTitle();
+                needsRebuild = false;
+            }
+
+            orig_Update();
+        }
+
+        private void RebuildMainAndTitle() {
             Overworld oui = Overworld;
             oui.UIs.Remove(oui.GetUI<OuiTitleScreen>());
             Oui title = new OuiTitleScreen() {
@@ -45,6 +62,8 @@ public void RebuildMainAndTitle() {
                 break;
             }
 
+            // this cannot be called more than once per frame, otherwise the
+            // scene keeps orphaned menu buttons which messes with selection
             CreateButtons();
 
             if (selected is MainMenuClimb) {