Merge pull request #199 from ovska/feat/highperformance/arraypool-ens…

…urecapacity

Add EnsureCapacity to ArrayPoolExtensions

CommunityToolkit.HighPerformance/Extensions/ArrayPoolExtensions.cs

@@ -53,4 +53,45 @@ public static void Resize<T>(this ArrayPool<T> pool, [NotNull] ref T[]? array, i
 
         array = newArray;
     }
+
+    /// <summary>
+    /// Ensures that when the method returns <paramref name="array"/> is not null and is at least <paramref name="capacity"/> in length.
+    /// Contents of <paramref name="array"/> are not copied if a new array is rented.
+    /// </summary>
+    /// <typeparam name="T">The type of items into the target array given as input.</typeparam>
+    /// <param name="pool">The target <see cref="ArrayPool{T}"/> instance used to rent and/or return the array.</param>
+    /// <param name="array">The rented <typeparamref name="T"/> array to ensure capacity for, or <see langword="null"/> to rent a new array.</param>
+    /// <param name="capacity">The minimum length of <paramref name="array"/> when the method returns.</param>
+    /// <param name="clearArray">Indicates whether the contents of the array should be cleared if returned to the pool.</param>
+    /// <exception cref="ArgumentOutOfRangeException">Thrown when <paramref name="capacity"/> is less than 0.</exception>
+    /// <remarks>When this method returns, the caller must not use any references to the old array anymore.</remarks>
+    public static void EnsureCapacity<T>(this ArrayPool<T> pool, [NotNull] ref T[]? array, int capacity, bool clearArray = false)
+    {
+        if (capacity < 0)
+        {
+            ThrowArgumentOutOfRangeExceptionForNegativeArrayCapacity();
+        }
+
+        if (array is null)
+        {
+            array = pool.Rent(capacity);
+        }
+        else if (array.Length < capacity)
+        {
+            // Ensure rent succeeds before returning the original array to the pool
+            T[] newArray = pool.Rent(capacity);
+
+            pool.Return(array, clearArray);
+
+            array = newArray;
+        }
+    }
+
+    /// <summary>
+    /// Throws an <see cref="ArgumentOutOfRangeException"/> when the "capacity" parameter is negative.
+    /// </summary>
+    private static void ThrowArgumentOutOfRangeExceptionForNegativeArrayCapacity()
+    {
+        throw new ArgumentOutOfRangeException("capacity", "The array capacity must be a positive number.");
+    }
 }

tests/CommunityToolkit.HighPerformance.UnitTests/Extensions/Test_ArrayPoolExtensions.cs

@@ -15,15 +15,15 @@ public class Test_ArrayPoolExtensions
 {
     [TestMethod]
     [ExpectedException(typeof(ArgumentOutOfRangeException))]
-    public void Test_ArrayExtensions_InvalidSize()
+    public void Test_ArrayPoolExtensions_Resize_InvalidSize()
     {
         int[]? array = null;
 
         ArrayPool<int>.Shared.Resize(ref array, -1);
     }
 
     [TestMethod]
-    public void Test_ArrayExtensions_NewArray()
+    public void Test_ArrayPoolExtensions_Resize_NewArray()
     {
         int[]? array = null;
 
@@ -34,7 +34,7 @@ public void Test_ArrayExtensions_NewArray()
     }
 
     [TestMethod]
-    public void Test_ArrayExtensions_SameSize()
+    public void Test_ArrayPoolExtensions_Resize_SameSize()
     {
         int[] array = ArrayPool<int>.Shared.Rent(10);
         int[] backup = array;
@@ -45,7 +45,7 @@ public void Test_ArrayExtensions_SameSize()
     }
 
     [TestMethod]
-    public void Test_ArrayExtensions_Expand()
+    public void Test_ArrayPoolExtensions_Resize_Expand()
     {
         int[] array = ArrayPool<int>.Shared.Rent(16);
         int[] backup = array;
@@ -60,7 +60,7 @@ public void Test_ArrayExtensions_Expand()
     }
 
     [TestMethod]
-    public void Test_ArrayExtensions_Shrink()
+    public void Test_ArrayPoolExtensions_Resize_Shrink()
     {
         int[] array = ArrayPool<int>.Shared.Rent(32);
         int[] backup = array;
@@ -75,7 +75,7 @@ public void Test_ArrayExtensions_Shrink()
     }
 
     [TestMethod]
-    public void Test_ArrayExtensions_Clear()
+    public void Test_ArrayPoolExtensions_Resize_Clear()
     {
         int[] array = ArrayPool<int>.Shared.Rent(16);
         int[] backup = array;
@@ -87,4 +87,73 @@ public void Test_ArrayExtensions_Clear()
         Assert.AreNotSame(array, backup);
         Assert.IsTrue(backup.AsSpan(0, 16).ToArray().All(i => i == 0));
     }
+
+    [TestMethod]
+    [ExpectedException(typeof(ArgumentOutOfRangeException))]
+    public void Test_ArrayPoolExtensions_EnsureCapacity_InvalidCapacity()
+    {
+        int[]? array = null;
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, -1);
+    }
+
+    [TestMethod]
+    public void Test_ArrayPoolExtensions_EnsureCapacity_IdenticalCapacity()
+    {
+        int[]? array = ArrayPool<int>.Shared.Rent(10);
+        int[]? backup = array;
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 10);
+        Assert.AreSame(backup, array);
+        Assert.IsTrue(array.Length >= 10);
+    }
+
+    [TestMethod]
+    public void Test_ArrayPoolExtensions_EnsureCapacity_NewArray()
+    {
+        int[]? array = null;
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 7);
+
+        Assert.IsNotNull(array);
+        Assert.IsTrue(array.Length >= 7);
+        int[]? backup = array;
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 64);
+
+        Assert.AreNotSame(backup, array);
+        Assert.IsTrue(array.Length >= 64);
+    }
+
+    [TestMethod]
+    public void Test_ArrayPoolExtensions_EnsureCapacity_SufficientCapacity()
+    {
+        int[]? array = ArrayPool<int>.Shared.Rent(16);
+        int[]? backup = array;
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 8);
+        Assert.AreSame(backup, array);
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 16);
+        Assert.AreSame(backup, array);
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 0);
+        Assert.AreSame(backup, array);
+    }
+
+    [TestMethod]
+    public void Test_ArrayPoolExtensions_EnsureCapacity_ClearArray()
+    {
+        int[]? array = ArrayPool<int>.Shared.Rent(16);
+        int[]? backup = array;
+
+        array.AsSpan().Fill(7);
+        Assert.IsTrue(backup.All(i => i == 7));
+
+        ArrayPool<int>.Shared.EnsureCapacity(ref array, 256, true);
+
+        Assert.AreNotSame(backup, array);
+        Assert.IsTrue(backup.All(i => i == default));
+        Assert.IsTrue(array.Length >= 256);
+    }
 }