From 3beca64cc514f4667340b004e6ee0553ca7cd92c Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Tue, 9 Jul 2024 21:53:57 +0900
Subject: [PATCH] Attempt to stay on correct difficulty

---
 osu.Game/Screens/Edit/ExternalEditScreen.cs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/osu.Game/Screens/Edit/ExternalEditScreen.cs b/osu.Game/Screens/Edit/ExternalEditScreen.cs
index fd438eacb3..ae5fad3ec0 100644
--- a/osu.Game/Screens/Edit/ExternalEditScreen.cs
+++ b/osu.Game/Screens/Edit/ExternalEditScreen.cs
@@ -169,6 +169,8 @@ namespace osu.Game.Screens.Edit
 
         private void finish()
         {
+            string originalDifficulty = editor.Beatmap.Value.Beatmap.BeatmapInfo.DifficultyName;
+
             showSpinner("Cleaning up...");
 
             EditOperation!.Finish().ContinueWith(t =>
@@ -178,12 +180,18 @@ namespace osu.Game.Screens.Edit
                     // Setting to null will allow exit to succeed.
                     operation = null;
 
-                    var beatmap = t.GetResultSafely();
+                    Live<BeatmapSetInfo>? beatmap = t.GetResultSafely();
 
                     if (beatmap == null)
                         this.Exit();
                     else
-                        editor.SwitchToDifficulty(beatmap.Value.Detach().Beatmaps.First());
+                    {
+                        var closestMatchingBeatmap =
+                            beatmap.Value.Beatmaps.FirstOrDefault(b => b.DifficultyName == originalDifficulty)
+                            ?? beatmap.Value.Beatmaps.First();
+
+                        editor.SwitchToDifficulty(closestMatchingBeatmap);
+                    }
                 });
             });
         }