⚡ optimize IEnumerable.Shuffle to use ArrayPool

src/GameUtils/Extensions/CollectionExtensions.cs

@@ -52,6 +52,8 @@ public static T GetRandom<T>(this IReadOnlyList<T> list)
     /// </summary>
     public static IEnumerable<T> Shuffle<T>(this IEnumerable<T> source)
     {
+        ArgumentNullException.ThrowIfNull(source);
+
         if (source.TryGetNonEnumeratedCount(out int count))
         {
             if (count == 0)
@@ -81,18 +83,45 @@ public static IEnumerable<T> Shuffle<T>(this IEnumerable<T> source)
         }
         else
         {
-            var list = new List<T>();
-            foreach (var item in source)
+            int initialCapacity = 16;
+            T[] rentedArray = System.Buffers.ArrayPool<T>.Shared.Rent(initialCapacity);
+            int itemCapacity = rentedArray.Length;
+            int numItems = 0;
+            try
             {
-                list.Add(item);
+                foreach (var item in source)
+                {
+                    if (numItems == itemCapacity)
+                    {
+                        T[] newArray = System.Buffers.ArrayPool<T>.Shared.Rent(itemCapacity * 2);
+                        rentedArray.AsSpan(0, numItems).CopyTo(newArray);
+                        System.Buffers.ArrayPool<T>.Shared.Return(rentedArray, clearArray: System.Runtime.CompilerServices.RuntimeHelpers.IsReferenceOrContainsReferences<T>());
+                        rentedArray = newArray;
+                        itemCapacity = newArray.Length;
+                    }
+                    rentedArray[numItems++] = item;
+                }
+
+                if (numItems == 0)
+                {
+                    return [];
+                }
+
+                var span = rentedArray.AsSpan(0, numItems);
+                for (var i = span.Length - 1; i > 0; i--)
+                {
+                    var j = Random.Shared.Next(i + 1);
+                    (span[i], span[j]) = (span[j], span[i]);
+                }
+
+                var result = new T[numItems];
+                span.CopyTo(result);
+                return result;
             }
-            var span = System.Runtime.InteropServices.CollectionsMarshal.AsSpan(list);
-            for (var i = span.Length - 1; i > 0; i--)
+            finally
             {
-                var j = Random.Shared.Next(i + 1);
-                (span[i], span[j]) = (span[j], span[i]);
+                System.Buffers.ArrayPool<T>.Shared.Return(rentedArray, clearArray: System.Runtime.CompilerServices.RuntimeHelpers.IsReferenceOrContainsReferences<T>());
             }
-            return list;
         }
     }
 

tests/GameUtils.Benchmarks/CollectionExtensionsBenchmarks.cs

@@ -0,0 +1,51 @@
+using System.Collections.Generic;
+using System.Linq;
+using BenchmarkDotNet.Attributes;
+
+namespace GameUtils.Benchmarks
+{
+    [MemoryDiagnoser]
+    public class CollectionExtensionsBenchmarks
+    {
+        private int[] _data = null!;
+
+        [Params(10, 100, 1000)]
+        public int N;
+
+        [GlobalSetup]
+        public void Setup()
+        {
+            _data = Enumerable.Range(0, N).ToArray();
+        }
+
+        private static IEnumerable<T> YieldSequence<T>(IEnumerable<T> source)
+        {
+            foreach (var item in source)
+            {
+                yield return item;
+            }
+        }
+
+        [Benchmark]
+        public int Shuffle_UnknownCount()
+        {
+            int sum = 0;
+            foreach (var item in GameUtils.Extensions.CollectionExtensions.Shuffle(YieldSequence(_data)))
+            {
+                sum += item;
+            }
+            return sum;
+        }
+
+        [Benchmark]
+        public int Shuffle_KnownCount()
+        {
+            int sum = 0;
+            foreach (var item in GameUtils.Extensions.CollectionExtensions.Shuffle(_data))
+            {
+                sum += item;
+            }
+            return sum;
+        }
+    }
+}

tests/GameUtils.Tests/Extensions/CollectionExtensionsTests.cs

@@ -46,4 +46,43 @@ public void WeightedRandom_IList_ReturnsElement()
         var result = elements.WeightedRandom(x => x);
         Assert.IsTrue(elements.Contains(result));
     }
+
+    [TestMethod]
+    public void Shuffle_Array_ReturnsShuffledElements()
+    {
+        var source = new[] { 1, 2, 3, 4, 5 };
+        var result = GameUtils.Extensions.CollectionExtensions.Shuffle(source).ToList();
+        Assert.AreEqual(5, result.Count);
+        CollectionAssert.AreEquivalent(source, result);
+    }
+
+    [TestMethod]
+    public void Shuffle_IEnumerableNonEnumerated_ReturnsShuffledElements()
+    {
+        static IEnumerable<int> YieldItems()
+        {
+            for (int i = 1; i <= 10; i++) yield return i;
+        }
+        var result = GameUtils.Extensions.CollectionExtensions.Shuffle(YieldItems()).ToList();
+        Assert.AreEqual(10, result.Count);
+        CollectionAssert.AreEquivalent(Enumerable.Range(1, 10).ToList(), result);
+    }
+
+    [TestMethod]
+    public void Shuffle_EmptyIEnumerable_ReturnsEmpty()
+    {
+        static IEnumerable<int> YieldEmpty()
+        {
+            yield break;
+        }
+        var result = GameUtils.Extensions.CollectionExtensions.Shuffle(YieldEmpty()).ToList();
+        Assert.AreEqual(0, result.Count);
+    }
+
+    [TestMethod]
+    public void Shuffle_NullSource_ThrowsArgumentNullException()
+    {
+        IEnumerable<int> source = null!;
+        Assert.ThrowsExactly<ArgumentNullException>(() => GameUtils.Extensions.CollectionExtensions.Shuffle(source));
+    }
 }