From bc21e30b09ef42fa17f8297a58bab4873181a100 Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Thu, 30 Jan 2020 17:35:54 +0900
Subject: [PATCH] Allow specifying a valid list of types to performFromMenu

---
 osu.Game/OsuGame.cs | 45 ++++++++++++++++++++++++++-------------------
 1 file changed, 26 insertions(+), 19 deletions(-)

diff --git a/osu.Game/OsuGame.cs b/osu.Game/OsuGame.cs
index ff46c6d849..f54e6c02af 100644
--- a/osu.Game/OsuGame.cs
+++ b/osu.Game/OsuGame.cs
@@ -327,10 +327,10 @@ namespace osu.Game
                 return;
             }
 
-            performFromMainMenu(() =>
+            performFromScreen(screen =>
             {
                 // we might already be at song select, so a check is required before performing the load to solo.
-                if (menuScreen.IsCurrentScreen())
+                if (screen is MainMenu)
                     menuScreen.LoadToSolo();
 
                 // we might even already be at the song
@@ -344,7 +344,7 @@ namespace osu.Game
 
                 Ruleset.Value = first.Ruleset;
                 Beatmap.Value = BeatmapManager.GetWorkingBeatmap(first);
-            }, $"load {beatmap}", bypassScreenAllowChecks: true, targetScreen: typeof(PlaySongSelect));
+            }, $"load {beatmap}", bypassScreenAllowChecks: true, validScreens: new[] { typeof(PlaySongSelect) });
         }
 
         /// <summary>
@@ -381,11 +381,11 @@ namespace osu.Game
                 return;
             }
 
-            performFromMainMenu(() =>
+            performFromScreen(screen =>
             {
                 Beatmap.Value = BeatmapManager.GetWorkingBeatmap(databasedBeatmap);
 
-                menuScreen.Push(new ReplayPlayerLoader(databasedScore));
+                screen.Push(new ReplayPlayerLoader(databasedScore));
             }, $"watch {databasedScoreInfo}", bypassScreenAllowChecks: true);
         }
 
@@ -441,17 +441,20 @@ namespace osu.Game
         private ScheduledDelegate performFromMainMenuTask;
 
         /// <summary>
-        /// Perform an action only after returning to the main menu.
+        /// Perform an action only after returning to a specific screen specification.
         /// Eagerly tries to exit the current screen until it succeeds.
         /// </summary>
         /// <param name="action">The action to perform once we are in the correct state.</param>
         /// <param name="taskName">The task name to display in a notification (if we can't immediately reach the main menu state).</param>
-        /// <param name="targetScreen">An optional target screen type. If this screen is already current we can immediately perform the action without returning to the menu.</param>
+        /// <param name="validScreens">An optional collection of valid screen types. If any of these screens are already current we can immediately perform the action immediately, else the first valid parent will be made current before performing the action. <see cref="MainMenu"/> is used if not specified.</param>
         /// <param name="bypassScreenAllowChecks">Whether checking <see cref="IOsuScreen.AllowExternalScreenChange"/> should be bypassed.</param>
-        private void performFromMainMenu(Action action, string taskName, Type targetScreen = null, bool bypassScreenAllowChecks = false)
+        private void performFromScreen(Action<IScreen> action, string taskName, IEnumerable<Type> validScreens = null, bool bypassScreenAllowChecks = false)
         {
             performFromMainMenuTask?.Cancel();
 
+            validScreens ??= Enumerable.Empty<Type>();
+            validScreens = validScreens.Append(typeof(MainMenu));
+
             // if the current screen does not allow screen changing, give the user an option to try again later.
             if (!bypassScreenAllowChecks && (ScreenStack.CurrentScreen as IOsuScreen)?.AllowExternalScreenChange == false)
             {
@@ -460,7 +463,7 @@ namespace osu.Game
                     Text = $"Click here to {taskName}",
                     Activated = () =>
                     {
-                        performFromMainMenu(action, taskName, targetScreen, true);
+                        performFromScreen(action, taskName, validScreens, true);
                         return true;
                     }
                 });
@@ -471,23 +474,27 @@ namespace osu.Game
             CloseAllOverlays(false);
 
             // we may already be at the target screen type.
-            if (targetScreen != null && ScreenStack.CurrentScreen?.GetType() == targetScreen)
+            if (validScreens.Contains(ScreenStack.CurrentScreen?.GetType()) && !Beatmap.Disabled)
             {
-                action();
+                action(ScreenStack.CurrentScreen);
                 return;
             }
 
-            // all conditions have been met to continue with the action.
-            if (menuScreen?.IsCurrentScreen() == true && !Beatmap.Disabled)
+            // find closest valid target
+            IScreen screen = ScreenStack.CurrentScreen;
+
+            while (screen != null)
             {
-                action();
-                return;
+                if (validScreens.Contains(screen.GetType()))
+                {
+                    screen.MakeCurrent();
+                    break;
+                }
+
+                screen = screen.GetParentScreen();
             }
 
-            // menuScreen may not be initialised yet (null check required).
-            menuScreen?.MakeCurrent();
-
-            performFromMainMenuTask = Schedule(() => performFromMainMenu(action, taskName));
+            performFromMainMenuTask = Schedule(() => performFromScreen(action, taskName, validScreens));
         }
 
         /// <summary>