Add collection transfer logic to beatmap import-as-update flow

osu.Game.Tests/Database/BeatmapImporterUpdateTests.cs

@@ -9,6 +9,7 @@
 using NUnit.Framework;
 using osu.Framework.Allocation;
 using osu.Game.Beatmaps;
+using osu.Game.Collections;
 using osu.Game.Database;
 using osu.Game.Models;
 using osu.Game.Overlays.Notifications;
@@ -432,6 +433,126 @@ public void TestMetadataTransferred()
             });
         }
 
+        /// <summary>
+        /// If all difficulties in the original beatmap set are in a collection, presume the user also wants new difficulties added.
+        /// </summary>
+        [TestCase(false)]
+        [TestCase(true)]
+        public void TestCollectionTransferNewBeatmap(bool allOriginalBeatmapsInCollection)
+        {
+            RunTestWithRealmAsync(async (realm, storage) =>
+            {
+                var importer = new BeatmapImporter(storage, realm);
+                using var rulesets = new RealmRulesetStore(realm, storage);
+
+                using var __ = getBeatmapArchive(out string pathOriginal);
+                using var _ = getBeatmapArchiveWithModifications(out string pathMissingOneBeatmap, directory =>
+                {
+                    // remove one difficulty before first import
+                    directory.GetFiles("*.osu").First().Delete();
+                });
+
+                var importBeforeUpdate = await importer.Import(new ImportTask(pathMissingOneBeatmap));
+
+                Assert.That(importBeforeUpdate, Is.Not.Null);
+                Debug.Assert(importBeforeUpdate != null);
+
+                int beatmapsToAddToCollection = 0;
+
+                importBeforeUpdate.PerformWrite(s =>
+                {
+                    var beatmapCollection = s.Realm.Add(new BeatmapCollection("test collection"));
+                    beatmapsToAddToCollection = s.Beatmaps.Count - (allOriginalBeatmapsInCollection ? 0 : 1);
+
+                    for (int i = 0; i < beatmapsToAddToCollection; i++)
+                        beatmapCollection.BeatmapMD5Hashes.Add(s.Beatmaps[i].MD5Hash);
+                });
+
+                // Second import matches first but contains one extra .osu file.
+                var importAfterUpdate = await importer.ImportAsUpdate(new ProgressNotification(), new ImportTask(pathOriginal), importBeforeUpdate.Value);
+
+                Assert.That(importAfterUpdate, Is.Not.Null);
+                Debug.Assert(importAfterUpdate != null);
+
+                importAfterUpdate.PerformRead(updated =>
+                {
+                    updated.Realm.Refresh();
+
+                    string[] hashes = updated.Realm.All<BeatmapCollection>().Single().BeatmapMD5Hashes.ToArray();
+
+                    if (allOriginalBeatmapsInCollection)
+                    {
+                        Assert.That(updated.Beatmaps.Count, Is.EqualTo(beatmapsToAddToCollection + 1));
+                        Assert.That(hashes, Has.Length.EqualTo(updated.Beatmaps.Count));
+                    }
+                    else
+                    {
+                        // Collection contains one less than the original beatmap, and two less after update (new difficulty included).
+                        Assert.That(updated.Beatmaps.Count, Is.EqualTo(beatmapsToAddToCollection + 2));
+                        Assert.That(hashes, Has.Length.EqualTo(beatmapsToAddToCollection));
+                    }
+                });
+            });
+        }
+
+        /// <summary>
+        /// If a difficulty in the original beatmap set is modified, the updated version should remain in any collections it was in.
+        /// </summary>
+        [Test]
+        public void TestCollectionTransferModifiedBeatmap()
+        {
+            RunTestWithRealmAsync(async (realm, storage) =>
+            {
+                var importer = new BeatmapImporter(storage, realm);
+                using var rulesets = new RealmRulesetStore(realm, storage);
+
+                using var __ = getBeatmapArchive(out string pathOriginal);
+                using var _ = getBeatmapArchiveWithModifications(out string pathModified, directory =>
+                {
+                    // Modify one .osu file with different content.
+                    var firstOsuFile = directory.GetFiles("*[Hard]*.osu").First();
+
+                    string existingContent = File.ReadAllText(firstOsuFile.FullName);
+
+                    File.WriteAllText(firstOsuFile.FullName, existingContent + "\n# I am new content");
+                });
+
+                var importBeforeUpdate = await importer.Import(new ImportTask(pathOriginal));
+
+                Assert.That(importBeforeUpdate, Is.Not.Null);
+                Debug.Assert(importBeforeUpdate != null);
+
+                string originalHash = string.Empty;
+
+                importBeforeUpdate.PerformWrite(s =>
+                {
+                    var beatmapCollection = s.Realm.Add(new BeatmapCollection("test collection"));
+                    originalHash = s.Beatmaps.Single(b => b.DifficultyName == "Hard").MD5Hash;
+
+                    beatmapCollection.BeatmapMD5Hashes.Add(originalHash);
+                });
+
+                // Second import matches first but contains a modified .osu file.
+                var importAfterUpdate = await importer.ImportAsUpdate(new ProgressNotification(), new ImportTask(pathModified), importBeforeUpdate.Value);
+
+                Assert.That(importAfterUpdate, Is.Not.Null);
+                Debug.Assert(importAfterUpdate != null);
+
+                importAfterUpdate.PerformRead(updated =>
+                {
+                    updated.Realm.Refresh();
+
+                    string[] hashes = updated.Realm.All<BeatmapCollection>().Single().BeatmapMD5Hashes.ToArray();
+                    string updatedHash = updated.Beatmaps.Single(b => b.DifficultyName == "Hard").MD5Hash;
+
+                    Assert.That(hashes, Has.Length.EqualTo(1));
+                    Assert.That(hashes.First(), Is.EqualTo(updatedHash));
+
+                    Assert.That(updatedHash, Is.Not.EqualTo(originalHash));
+                });
+            });
+        }
+
         private static void checkCount<T>(RealmAccess realm, int expected, Expression<Func<T, bool>>? condition = null) where T : RealmObject
         {
             var query = realm.Realm.All<T>();

osu.Game/Beatmaps/BeatmapImporter.cs

@@ -14,6 +14,7 @@
 using osu.Framework.Platform;
 using osu.Framework.Testing;
 using osu.Game.Beatmaps.Formats;
+using osu.Game.Collections;
 using osu.Game.Database;
 using osu.Game.Extensions;
 using osu.Game.IO;
@@ -71,6 +72,8 @@ public BeatmapImporter(Storage storage, RealmAccess realm)
                 // Transfer local values which should be persisted across a beatmap update.
                 updated.DateAdded = original.DateAdded;
 
+                transferCollectionReferences(realm, original, updated);
+
                 foreach (var beatmap in original.Beatmaps.ToArray())
                 {
                     var updatedBeatmap = updated.Beatmaps.FirstOrDefault(b => b.Hash == beatmap.Hash);
@@ -112,6 +115,40 @@ public BeatmapImporter(Storage storage, RealmAccess realm)
             return first;
         }
 
+        private static void transferCollectionReferences(Realm realm, BeatmapSetInfo original, BeatmapSetInfo updated)
+        {
+            // First check if every beatmap in the original set is in any collections.
+            // In this case, we will assume they also want any newly added difficulties added to the collection.
+            foreach (var c in realm.All<BeatmapCollection>())
+            {
+                if (original.Beatmaps.Select(b => b.MD5Hash).All(c.BeatmapMD5Hashes.Contains))
+                {
+                    foreach (var b in original.Beatmaps)
+                        c.BeatmapMD5Hashes.Remove(b.MD5Hash);
+
+                    foreach (var b in updated.Beatmaps)
+                        c.BeatmapMD5Hashes.Add(b.MD5Hash);
+                }
+            }
+
+            // 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);
+                }
+            }
+        }
+
         protected override bool ShouldDeleteArchive(string path) => Path.GetExtension(path).ToLowerInvariant() == ".osz";
 
         protected override void Populate(BeatmapSetInfo beatmapSet, ArchiveReader? archive, Realm realm, CancellationToken cancellationToken = default)

osu.Game/Database/Live.cs

@@ -2,6 +2,7 @@
 // See the LICENCE file in the repository root for full licence text.
 
 using System;
+using JetBrains.Annotations;
 
 namespace osu.Game.Database
 {
@@ -18,19 +19,19 @@ public abstract class Live<T> : IEquatable<Live<T>>
         /// Perform a read operation on this live object.
         /// </summary>
         /// <param name="perform">The action to perform.</param>
-        public abstract void PerformRead(Action<T> perform);
+        public abstract void PerformRead([InstantHandle] Action<T> perform);
 
         /// <summary>
         /// Perform a read operation on this live object.
         /// </summary>
         /// <param name="perform">The action to perform.</param>
-        public abstract TReturn PerformRead<TReturn>(Func<T, TReturn> perform);
+        public abstract TReturn PerformRead<TReturn>([InstantHandle] Func<T, TReturn> perform);
 
         /// <summary>
         /// Perform a write operation on this live object.
         /// </summary>
         /// <param name="perform">The action to perform.</param>
-        public abstract void PerformWrite(Action<T> perform);
+        public abstract void PerformWrite([InstantHandle] Action<T> perform);
 
         /// <summary>
         /// Whether this instance is tracking data which is managed by the database backing.