From 8707a1b857deaf29c608012f9ac25fee305d9698 Mon Sep 17 00:00:00 2001
From: Cvolton <cvolton@cvolton.eu>
Date: Tue, 14 Jan 2025 19:17:06 +0100
Subject: [PATCH] fix SceneManager use after frees

---
 loader/src/hooks/persist.cpp         | 24 ++++++++++++++++++++++++
 loader/src/ui/nodes/SceneManager.cpp |  4 ++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/loader/src/hooks/persist.cpp b/loader/src/hooks/persist.cpp
index 0c0ea4681..9d278f6cc 100644
--- a/loader/src/hooks/persist.cpp
+++ b/loader/src/hooks/persist.cpp
@@ -1,4 +1,5 @@
 #include <Geode/ui/SceneManager.hpp>
+#include <Geode/modify/CCDirector.hpp>
 
 using namespace geode::prelude;
 
@@ -30,4 +31,27 @@ struct SceneSwitch : Modify<SceneSwitch, AchievementNotifier> {
 
 #endif
 
+struct SceneSwitch2 : Modify<SceneSwitch2, CCDirector> {
+    GEODE_FORWARD_COMPAT_DISABLE_HOOKS("persist disabled")
+    // CCDirector does not call willSwitchToScene in these 2 instances,
+    // so we have to do it ourselves to make everything behave as expected
+    void popScene() {
+        CCDirector::popScene();
+    #ifdef GEODE_IS_WINDOWS
+        AppDelegate::get()->willSwitchToScene(m_pNextScene);
+    #else
+        AchievementNotifier::get()->willSwitchToScene(m_pNextScene);
+    #endif
+    }
+
+    void popToSceneStackLevel(int level) {
+        CCDirector::popToSceneStackLevel(level);
+    #ifdef GEODE_IS_WINDOWS
+        AppDelegate::get()->willSwitchToScene(m_pNextScene);
+    #else
+        AchievementNotifier::get()->willSwitchToScene(m_pNextScene);
+    #endif
+    }
+};
+
 }
\ No newline at end of file
diff --git a/loader/src/ui/nodes/SceneManager.cpp b/loader/src/ui/nodes/SceneManager.cpp
index 8a9982da4..eeaf14500 100644
--- a/loader/src/ui/nodes/SceneManager.cpp
+++ b/loader/src/ui/nodes/SceneManager.cpp
@@ -18,11 +18,11 @@ void SceneManager::keepAcrossScenes(CCNode* node) {
     if (ranges::contains(m_persistedNodes, node)) {
         return;
     }
+    m_persistedNodes.push_back(node);
     if (m_lastScene) {
         node->removeFromParentAndCleanup(false);
         m_lastScene->addChild(node);
     }
-    m_persistedNodes.push_back(node);
 }
 
 void SceneManager::forget(CCNode* node) {
@@ -37,7 +37,7 @@ void SceneManager::willSwitchToScene(CCScene* scene) {
     for (auto& node : m_persistedNodes) {
         // no cleanup in order to keep actions running
         node->removeFromParentAndCleanup(false);
-        scene->addChild(node);
+        if(scene) scene->addChild(node);
     }
     m_lastScene = scene;
 }