Merge pull request #20641 from pfgithub/patch-fix-save-collections

Preserve collections when saving a beatmap

osu.Game.Tests/Beatmaps/WorkingBeatmapManagerTest.cs

@@ -9,8 +9,10 @@ using osu.Framework.Extensions;
 using osu.Framework.Platform;
 using osu.Framework.Testing;
 using osu.Game.Beatmaps;
+using osu.Game.Collections;
 using osu.Game.Database;
 using osu.Game.Rulesets;
+using osu.Game.Rulesets.Osu;
 using osu.Game.Tests.Resources;
 using osu.Game.Tests.Visual;
 
@@ -96,5 +98,41 @@ namespace osu.Game.Tests.Beatmaps
             var second = beatmaps.GetWorkingBeatmap(beatmap, true);
             Assert.That(first, Is.Not.SameAs(second));
         });
+
+        [Test]
+        public void TestSavePreservesCollections() => AddStep("run test", () =>
+        {
+            var beatmap = Realm.Run(r => r.Find<BeatmapInfo>(importedSet.Beatmaps.First().ID).Detach());
+
+            var working = beatmaps.GetWorkingBeatmap(beatmap);
+
+            Assert.That(working.BeatmapInfo.BeatmapSet?.Files, Has.Count.GreaterThan(0));
+
+            string initialHash = working.BeatmapInfo.MD5Hash;
+
+            var preserveCollection = new BeatmapCollection("test contained");
+            preserveCollection.BeatmapMD5Hashes.Add(initialHash);
+
+            var noNewCollection = new BeatmapCollection("test not contained");
+
+            Realm.Write(r =>
+            {
+                r.Add(preserveCollection);
+                r.Add(noNewCollection);
+            });
+
+            Assert.That(preserveCollection.BeatmapMD5Hashes, Does.Contain(initialHash));
+            Assert.That(noNewCollection.BeatmapMD5Hashes, Does.Not.Contain(initialHash));
+
+            beatmaps.Save(working.BeatmapInfo, working.GetPlayableBeatmap(new OsuRuleset().RulesetInfo));
+
+            string finalHash = working.BeatmapInfo.MD5Hash;
+
+            Assert.That(finalHash, Is.Not.SameAs(initialHash));
+
+            Assert.That(preserveCollection.BeatmapMD5Hashes, Does.Not.Contain(initialHash));
+            Assert.That(preserveCollection.BeatmapMD5Hashes, Does.Contain(finalHash));
+            Assert.That(noNewCollection.BeatmapMD5Hashes, Does.Not.Contain(finalHash));
+        });
     }
 }

osu.Game/Beatmaps/BeatmapImporter.cs

@@ -141,18 +141,9 @@ namespace osu.Game.Beatmaps
             // Handle collections using permissive difficulty name to track difficulties.
             foreach (var originalBeatmap in original.Beatmaps)
             {
-                var updatedBeatmap = updated.Beatmaps.FirstOrDefault(b => b.DifficultyName == originalBeatmap.DifficultyName);
-
-                if (updatedBeatmap == null)
-                    continue;
-
-                var collections = realm.All<BeatmapCollection>().AsEnumerable().Where(c => c.BeatmapMD5Hashes.Contains(originalBeatmap.MD5Hash));
-
-                foreach (var c in collections)
-                {
-                    c.BeatmapMD5Hashes.Remove(originalBeatmap.MD5Hash);
-                    c.BeatmapMD5Hashes.Add(updatedBeatmap.MD5Hash);
-                }
+                updated.Beatmaps
+                       .FirstOrDefault(b => b.DifficultyName == originalBeatmap.DifficultyName)?
+                       .TransferCollectionReferences(realm, originalBeatmap.MD5Hash);
             }
         }
 

osu.Game/Beatmaps/BeatmapInfo.cs

@@ -8,6 +8,7 @@ using JetBrains.Annotations;
 using Newtonsoft.Json;
 using osu.Framework.Testing;
 using osu.Game.Beatmaps.ControlPoints;
+using osu.Game.Collections;
 using osu.Game.Database;
 using osu.Game.Models;
 using osu.Game.Online.API.Requests.Responses;
@@ -213,6 +214,23 @@ namespace osu.Game.Beatmaps
             return fileHashX == fileHashY;
         }
 
+        /// <summary>
+        /// When updating a beatmap, its hashes will change. Collections currently track beatmaps by hash, so they need to be updated.
+        /// This method will handle updating
+        /// </summary>
+        /// <param name="realm">A realm instance in an active write transaction.</param>
+        /// <param name="previousMD5Hash">The previous MD5 hash of the beatmap before update.</param>
+        public void TransferCollectionReferences(Realm realm, string previousMD5Hash)
+        {
+            var collections = realm.All<BeatmapCollection>().AsEnumerable().Where(c => c.BeatmapMD5Hashes.Contains(previousMD5Hash));
+
+            foreach (var c in collections)
+            {
+                c.BeatmapMD5Hashes.Remove(previousMD5Hash);
+                c.BeatmapMD5Hashes.Add(MD5Hash);
+            }
+        }
+
         IBeatmapMetadataInfo IBeatmapInfo.Metadata => Metadata;
         IBeatmapSetInfo? IBeatmapInfo.BeatmapSet => BeatmapSet;
         IRulesetInfo IBeatmapInfo.Ruleset => Ruleset;

osu.Game/Beatmaps/BeatmapManager.cs

@@ -311,6 +311,8 @@ namespace osu.Game.Beatmaps
                 if (existingFileInfo != null)
                     DeleteFile(setInfo, existingFileInfo);
 
+                string oldMd5Hash = beatmapInfo.MD5Hash;
+
                 beatmapInfo.MD5Hash = stream.ComputeMD5Hash();
                 beatmapInfo.Hash = stream.ComputeSHA2Hash();
 
@@ -327,6 +329,8 @@ namespace osu.Game.Beatmaps
 
                     setInfo.CopyChangesToRealm(liveBeatmapSet);
 
+                    beatmapInfo.TransferCollectionReferences(r, oldMd5Hash);
+
                     ProcessBeatmap?.Invoke((liveBeatmapSet, false));
                 });
             }