Merge pull request #616 from CommunityToolkit/dev/array-writer-segment

Add ArrayPoolBufferWriter<T>.DangerousGetArray() API

Author: Sergio0694

src/CommunityToolkit.HighPerformance/Buffers/ArrayPoolBufferWriter{T}.cs

@@ -244,6 +244,59 @@ public Span<T> GetSpan(int sizeHint = 0)
         return this.array.AsSpan(this.index);
     }
 
+    /// <summary>
+    /// Gets an <see cref="ArraySegment{T}"/> instance wrapping the underlying <typeparamref name="T"/> array in use.
+    /// </summary>
+    /// <returns>An <see cref="ArraySegment{T}"/> instance wrapping the underlying <typeparamref name="T"/> array in use.</returns>
+    /// <exception cref="ObjectDisposedException">Thrown when the buffer in use has already been disposed.</exception>
+    /// <remarks>
+    /// This method is meant to be used when working with APIs that only accept an array as input, and should be used with caution.
+    /// In particular, the returned array is rented from an array pool, and it is responsibility of the caller to ensure that it's
+    /// not used after the current <see cref="ArrayPoolBufferWriter{T}"/> instance is disposed. Doing so is considered undefined
+    /// behavior, as the same array might be in use within another <see cref="ArrayPoolBufferWriter{T}"/> instance.
+    /// </remarks>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public ArraySegment<T> DangerousGetArray()
+    {
+        T[]? array = this.array;
+
+        if (array is null)
+        {
+            ThrowObjectDisposedException();
+        }
+
+        return new(array!, 0, this.index);
+    }
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        T[]? array = this.array;
+
+        if (array is null)
+        {
+            return;
+        }
+
+        this.array = null;
+
+        this.pool.Return(array);
+    }
+
+    /// <inheritdoc/>
+    public override string ToString()
+    {
+        // See comments in MemoryOwner<T> about this
+        if (typeof(T) == typeof(char) &&
+            this.array is char[] chars)
+        {
+            return new(chars, 0, this.index);
+        }
+
+        // Same representation used in Span<T>
+        return $"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(T)}>[{this.index}]";
+    }
+
     /// <summary>
     /// Ensures that <see cref="array"/> has enough free space to contain a given number of new items.
     /// </summary>
@@ -296,35 +349,6 @@ private void ResizeBuffer(int sizeHint)
         this.pool.Resize(ref this.array, (int)minimumSize);
     }
 
-    /// <inheritdoc/>
-    public void Dispose()
-    {
-        T[]? array = this.array;
-
-        if (array is null)
-        {
-            return;
-        }
-
-        this.array = null;
-
-        this.pool.Return(array);
-    }
-
-    /// <inheritdoc/>
-    public override string ToString()
-    {
-        // See comments in MemoryOwner<T> about this
-        if (typeof(T) == typeof(char) &&
-            this.array is char[] chars)
-        {
-            return new(chars, 0, this.index);
-        }
-
-        // Same representation used in Span<T>
-        return $"CommunityToolkit.HighPerformance.Buffers.ArrayPoolBufferWriter<{typeof(T)}>[{this.index}]";
-    }
-
     /// <summary>
     /// Throws an <see cref="ArgumentOutOfRangeException"/> when the requested count is negative.
     /// </summary>

tests/CommunityToolkit.HighPerformance.UnitTests/Buffers/Test_ArrayPoolBufferWriter{T}.cs

@@ -3,9 +3,11 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Buffers;
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.InteropServices;
 using CommunityToolkit.HighPerformance;
 using CommunityToolkit.HighPerformance.Buffers;
 using CommunityToolkit.HighPerformance.UnitTests.Buffers.Internals;
@@ -218,4 +220,30 @@ public void Test_ArrayPoolBufferWriterOfT_AsStream()
         // Now check that the writer is actually disposed instead
         _ = Assert.ThrowsException<ObjectDisposedException>(() => writer.Capacity);
     }
+
+    [TestMethod]
+    public void Test_ArrayPoolBufferWriterOfT_AllocateAndGetArray()
+    {
+        ArrayPoolBufferWriter<int>? bufferWriter = new();
+
+        // Write some random data
+        bufferWriter.Write(Enumerable.Range(0, 127).ToArray());
+
+        // Get the array for the written segment
+        ArraySegment<int> segment = bufferWriter.DangerousGetArray();
+
+        Assert.IsNotNull(segment.Array);
+        Assert.IsTrue(segment.Array.Length >= bufferWriter.WrittenSpan.Length);
+        Assert.AreEqual(segment.Offset, 0);
+        Assert.AreEqual(segment.Count, bufferWriter.WrittenSpan.Length);
+
+        _ = MemoryMarshal.TryGetArray(bufferWriter.WrittenMemory, out ArraySegment<int> writtenSegment);
+
+        // The array is the same one as the one from the written span
+        Assert.AreSame(segment.Array, writtenSegment.Array);
+
+        bufferWriter.Dispose();
+
+        _ = Assert.ThrowsException<ObjectDisposedException>(() => bufferWriter.DangerousGetArray());
+    }
 }