refactor: update coffee.internal

Runtime/Internal/Extensions/ComponentExtensions.cs

@@ -18,10 +18,10 @@ namespace Coffee.UIParticleInternal
         public static T[] GetComponentsInChildren<T>(this Component self, int depth)
             where T : Component
         {
-            var results = ListPool<T>.Rent();
+            var results = InternalListPool<T>.Rent();
             self.GetComponentsInChildren_Internal(results, depth);
             var array = results.ToArray();
-            ListPool<T>.Return(ref results);
+            InternalListPool<T>.Return(ref results);
             return array;
         }
 

Runtime/Internal/Utilities/FastAction.cs

@@ -10,8 +10,9 @@ namespace Coffee.UIParticleInternal
     /// </summary>
     internal class FastActionBase<T>
     {
-        private static readonly ObjectPool<LinkedListNode<T>> s_NodePool =
+        private static readonly InternalObjectPool<LinkedListNode<T>> s_NodePool =
-            new ObjectPool<LinkedListNode<T>>(() => new LinkedListNode<T>(default), _ => true, x => x.Value = default);
+            new InternalObjectPool<LinkedListNode<T>>(() => new LinkedListNode<T>(default), _ => true,
+                x => x.Value = default);
 
         private readonly LinkedList<T> _delegates = new LinkedList<T>();
 

Runtime/Internal/Utilities/Misc.cs

@@ -1,6 +1,13 @@
+using System;
 using System.Diagnostics;
 using UnityEditor;
 using UnityEngine;
+using Object = UnityEngine.Object;
+#if UNITY_EDITOR && UNITY_2021_2_OR_NEWER
+using UnityEditor.SceneManagement;
+#elif UNITY_EDITOR
+using UnityEditor.Experimental.SceneManagement;
+#endif
 
 namespace Coffee.UIParticleInternal
 {
@@ -53,5 +60,15 @@ namespace Coffee.UIParticleInternal
             EditorUtility.SetDirty(obj);
 #endif
         }
+
+#if UNITY_EDITOR
+        public static T[] GetAllComponentsInPrefabStage<T>() where T : Component
+        {
+            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+            if (prefabStage == null) return Array.Empty<T>();
+
+            return prefabStage.prefabContentsRoot.GetComponentsInChildren<T>(true);
+        }
+#endif
     }
 }

Runtime/Internal/Utilities/ObjectPool.cs

@@ -6,15 +6,58 @@ namespace Coffee.UIParticleInternal
     /// <summary>
     /// Object pool.
     /// </summary>
-    internal class ObjectPool<T>
+    internal class InternalObjectPool<T> where T : class
     {
+#if UNITY_2021_1_OR_NEWER
+        private readonly Predicate<T> _onValid; // Delegate for checking if instances are valid
+        private readonly UnityEngine.Pool.ObjectPool<T> _pool;
+
+        public InternalObjectPool(Func<T> onCreate, Predicate<T> onValid, Action<T> onReturn)
+        {
+            _pool = new UnityEngine.Pool.ObjectPool<T>(onCreate, null, onReturn);
+            _onValid = onValid;
+        }
+
+        /// <summary>
+        /// Rent an instance from the pool.
+        /// When you no longer need it, return it with <see cref="Return" />.
+        /// </summary>
+        public T Rent()
+        {
+            while (0 < _pool.CountInactive)
+            {
+                var instance = _pool.Get();
+                if (_onValid(instance))
+                {
+                    return instance;
+                }
+            }
+
+            // If there are no instances in the pool, create a new one.
+            Logging.Log(this, $"A new instance is created (pooled: {_pool.CountInactive}, created: {_pool.CountAll}).");
+            return _pool.Get();
+        }
+
+        /// <summary>
+        /// Return an instance to the pool and assign null.
+        /// Be sure to return the instance obtained with <see cref="Rent" /> with this method.
+        /// </summary>
+        public void Return(ref T instance)
+        {
+            if (instance == null) return; // Ignore if already pooled or null.
+
+            _pool.Release(instance);
+            Logging.Log(this, $"An instance is released (pooled: {_pool.CountInactive}, created: {_pool.CountAll}).");
+            instance = default; // Set the reference to null.
+        }
+#else
         private readonly Func<T> _onCreate; // Delegate for creating instances
         private readonly Action<T> _onReturn; // Delegate for returning instances to the pool
         private readonly Predicate<T> _onValid; // Delegate for checking if instances are valid
         private readonly Stack<T> _pool = new Stack<T>(32); // Object pool
         private int _count; // Total count of created instances
 
-        public ObjectPool(Func<T> onCreate, Predicate<T> onValid, Action<T> onReturn)
+        public InternalObjectPool(Func<T> onCreate, Predicate<T> onValid, Action<T> onReturn)
         {
             _onCreate = onCreate;
             _onValid = onValid;
@@ -54,15 +97,40 @@ namespace Coffee.UIParticleInternal
             Logging.Log(this, $"An instance is released (pooled: {_pool.Count}, created: {_count}).");
             instance = default; // Set the reference to null.
         }
+#endif
     }
 
     /// <summary>
     /// Object pool for <see cref="List{T}" />.
     /// </summary>
-    internal static class ListPool<T>
+    internal static class InternalListPool<T>
     {
-        private static readonly ObjectPool<List<T>> s_ListPool =
+#if UNITY_2021_1_OR_NEWER
-            new ObjectPool<List<T>>(() => new List<T>(), _ => true, x => x.Clear());
+        /// <summary>
+        /// Rent an instance from the pool.
+        /// When you no longer need it, return it with <see cref="Return" />.
+        /// </summary>
+        public static List<T> Rent()
+        {
+            return UnityEngine.Pool.ListPool<T>.Get();
+        }
+
+        /// <summary>
+        /// Return an instance to the pool and assign null.
+        /// Be sure to return the instance obtained with <see cref="Rent" /> with this method.
+        /// </summary>
+        public static void Return(ref List<T> toRelease)
+        {
+            if (toRelease != null)
+            {
+                UnityEngine.Pool.ListPool<T>.Release(toRelease);
+            }
+
+            toRelease = null;
+        }
+#else
+        private static readonly InternalObjectPool<List<T>> s_ListPool =
+            new InternalObjectPool<List<T>>(() => new List<T>(), _ => true, x => x.Clear());
 
         /// <summary>
         /// Rent an instance from the pool.
@@ -81,5 +149,6 @@ namespace Coffee.UIParticleInternal
         {
             s_ListPool.Return(ref toRelease);
         }
+#endif
     }
 }

Runtime/UIParticleRenderer.cs

@@ -421,7 +421,7 @@ namespace Coffee.UIExtensions
                     workerMesh.LinearToGamma();
                 }
 
-                var components = ListPool<Component>.Rent();
+                var components = InternalListPool<Component>.Rent();
                 GetComponents(typeof(IMeshModifier), components);
                 for (var i = 0; i < components.Count; i++)
                 {
@@ -430,7 +430,7 @@ namespace Coffee.UIExtensions
 #pragma warning restore CS0618 // Type or member is obsolete
                 }
 
-                ListPool<Component>.Return(ref components);
+                InternalListPool<Component>.Return(ref components);
             }
 
             Profiler.EndSample();
@@ -442,7 +442,7 @@ namespace Coffee.UIExtensions
 
             // Get grouped renderers.
             Profiler.BeginSample("[UIParticleRenderer] Set Mesh");
-            var renderers = ListPool<UIParticleRenderer>.Rent();
+            var renderers = InternalListPool<UIParticleRenderer>.Rent();
             if (_parent.useMeshSharing)
             {
                 UIParticleUpdater.GetGroupedRenderers(_parent.groupId, _index, renderers);
@@ -459,7 +459,7 @@ namespace Coffee.UIExtensions
                 r.canvasRenderer.SetMaterial(materialForRendering, 0);
             }
 
-            ListPool<UIParticleRenderer>.Return(ref renderers);
+            InternalListPool<UIParticleRenderer>.Return(ref renderers);
 
             if (_parent.canRender)
             {