Factor out pooling logic from Playfield

osu.Game/Rulesets/Objects/Pooling/HitObjectEntryManager.cs

@@ -0,0 +1,69 @@
+// Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+#nullable enable
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+
+namespace osu.Game.Rulesets.Objects.Pooling
+{
+    /// <summary>
+    /// Manages a mapping between <see cref="HitObject"/> and <see cref="HitObjectLifetimeEntry"/>
+    /// </summary>
+    internal class HitObjectEntryManager
+    {
+        /// <summary>
+        /// All entries, including entries of the nested hit objects.
+        /// </summary>
+        public IEnumerable<HitObjectLifetimeEntry> AllEntries => entryMap.Values;
+
+        public event Action<HitObjectLifetimeEntry, HitObject?>? OnEntryAdded;
+        public event Action<HitObjectLifetimeEntry, HitObject?>? OnEntryRemoved;
+
+        private readonly Func<HitObject, HitObjectLifetimeEntry> createLifetimeEntry;
+
+        private readonly Dictionary<HitObject, HitObjectLifetimeEntry> entryMap = new Dictionary<HitObject, HitObjectLifetimeEntry>();
+        private readonly Dictionary<HitObject, HitObject> parentMap = new Dictionary<HitObject, HitObject>();
+
+        public HitObjectEntryManager(Func<HitObject, HitObjectLifetimeEntry> createLifetimeEntry)
+        {
+            this.createLifetimeEntry = createLifetimeEntry;
+        }
+
+        public HitObjectLifetimeEntry Add(HitObject hitObject, HitObject? parentHitObject)
+        {
+            if (parentHitObject != null && !entryMap.TryGetValue(parentHitObject, out var parentEntry))
+                throw new InvalidOperationException($@"The parent {nameof(HitObject)} must be added to this {nameof(HitObjectEntryManager)} before nested {nameof(HitObject)} is added.");
+
+            if (entryMap.ContainsKey(hitObject))
+                throw new InvalidOperationException($@"The {nameof(HitObject)} is already added to this {nameof(HitObjectEntryManager)}.");
+
+            if (parentHitObject != null)
+                parentMap[hitObject] = parentHitObject;
+
+            var entry = createLifetimeEntry(hitObject);
+            entryMap[hitObject] = entry;
+
+            OnEntryAdded?.Invoke(entry, parentHitObject);
+            return entry;
+        }
+
+        public bool Remove(HitObject hitObject)
+        {
+            if (!entryMap.TryGetValue(hitObject, out var entry))
+                return false;
+
+            parentMap.Remove(hitObject, out var parentHitObject);
+
+            OnEntryRemoved?.Invoke(entry, parentHitObject);
+            return true;
+        }
+
+        public bool TryGet(HitObject hitObject, [MaybeNullWhen(false)] out HitObjectLifetimeEntry entry)
+        {
+            return entryMap.TryGetValue(hitObject, out entry);
+        }
+    }
+}

osu.Game/Rulesets/UI/Playfield.cs

@@ -21,6 +21,7 @@ using osu.Game.Rulesets.Objects;
 using osu.Game.Rulesets.Objects.Drawables;
 using osu.Game.Skinning;
 using osuTK;
+using osu.Game.Rulesets.Objects.Pooling;
 
 namespace osu.Game.Rulesets.UI
 {
@@ -94,6 +95,8 @@ namespace osu.Game.Rulesets.UI
         [Resolved(CanBeNull = true)]
         private IReadOnlyList<Mod> mods { get; set; }
 
+        private readonly HitObjectEntryManager entryManager;
+
         /// <summary>
         /// Creates a new <see cref="Playfield"/>.
         /// </summary>
@@ -108,6 +111,10 @@ namespace osu.Game.Rulesets.UI
                 h.HitObjectUsageBegan += o => HitObjectUsageBegan?.Invoke(o);
                 h.HitObjectUsageFinished += o => HitObjectUsageFinished?.Invoke(o);
             }));
+
+            entryManager = new HitObjectEntryManager(CreateLifetimeEntry);
+            entryManager.OnEntryAdded += onEntryAdded;
+            entryManager.OnEntryRemoved += onEntryRemoved;
         }
 
         [BackgroundDependencyLoader]
@@ -171,6 +178,7 @@ namespace osu.Game.Rulesets.UI
         /// <param name="hitObject">The added <see cref="HitObject"/>.</param>
         protected virtual void OnHitObjectAdded(HitObject hitObject)
         {
+            preloadSamples(hitObject);
         }
 
         /// <summary>
@@ -263,13 +271,7 @@ namespace osu.Game.Rulesets.UI
         /// <param name="hitObject"></param>
         public virtual void Add(HitObject hitObject)
         {
-            var entry = CreateLifetimeEntry(hitObject);
+            entryManager.Add(hitObject, null);
-            lifetimeEntryMap[entry.HitObject] = entry;
-
-            preloadSamples(hitObject);
-
-            HitObjectContainer.Add(entry);
-            OnHitObjectAdded(entry.HitObject);
         }
 
         private void preloadSamples(HitObject hitObject)
@@ -292,14 +294,23 @@ namespace osu.Game.Rulesets.UI
         /// <returns>Whether the <see cref="HitObject"/> was successfully removed.</returns>
         public virtual bool Remove(HitObject hitObject)
         {
-            if (lifetimeEntryMap.Remove(hitObject, out var entry))
+            return entryManager.Remove(hitObject) || nestedPlayfields.Any(p => p.Remove(hitObject));
-            {
+        }
-                HitObjectContainer.Remove(entry);
-                OnHitObjectRemoved(hitObject);
-                return true;
-            }
 
-            return nestedPlayfields.Any(p => p.Remove(hitObject));
+        private void onEntryAdded(HitObjectLifetimeEntry entry, [CanBeNull] HitObject parentHitObject)
+        {
+            if (parentHitObject != null) return;
+
+            HitObjectContainer.Add(entry);
+            OnHitObjectAdded(entry.HitObject);
+        }
+
+        private void onEntryRemoved(HitObjectLifetimeEntry entry, [CanBeNull] HitObject parentHitObject)
+        {
+            if (parentHitObject != null) return;
+
+            HitObjectContainer.Remove(entry);
+            OnHitObjectRemoved(entry.HitObject);
         }
 
         /// <summary>
@@ -366,8 +377,8 @@ namespace osu.Game.Rulesets.UI
                     }
                 }
 
-                if (!lifetimeEntryMap.TryGetValue(hitObject, out var entry))
+                if (!entryManager.TryGet(hitObject, out var entry))
-                    lifetimeEntryMap[hitObject] = entry = CreateLifetimeEntry(hitObject);
+                    entry = entryManager.Add(hitObject, parent?.HitObject);
 
                 dho.ParentHitObject = parent;
                 dho.Apply(entry);
@@ -442,8 +453,6 @@ namespace osu.Game.Rulesets.UI
         /// </remarks>
         internal event Action<HitObject> HitObjectUsageFinished;
 
-        private readonly Dictionary<HitObject, HitObjectLifetimeEntry> lifetimeEntryMap = new Dictionary<HitObject, HitObjectLifetimeEntry>();
-
         /// <summary>
         /// Sets whether to keep a given <see cref="HitObject"/> always alive within this or any nested <see cref="Playfield"/>.
         /// </summary>
@@ -451,7 +460,7 @@ namespace osu.Game.Rulesets.UI
         /// <param name="keepAlive">Whether to keep <paramref name="hitObject"/> always alive.</param>
         internal void SetKeepAlive(HitObject hitObject, bool keepAlive)
         {
-            if (lifetimeEntryMap.TryGetValue(hitObject, out var entry))
+            if (entryManager.TryGet(hitObject, out var entry))
             {
                 entry.KeepAlive = keepAlive;
                 return;
@@ -466,7 +475,7 @@ namespace osu.Game.Rulesets.UI
         /// </summary>
         internal void KeepAllAlive()
         {
-            foreach (var (_, entry) in lifetimeEntryMap)
+            foreach (var entry in entryManager.AllEntries)
                 entry.KeepAlive = true;
 
             foreach (var p in nestedPlayfields)