Switch ImmutableArrayBuilder<T> to array pool

src/PolySharp.SourceGenerators/Helpers/ImmutableArrayBuilder{T}.cs

@@ -3,7 +3,9 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Buffers;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 
 namespace PolySharp.SourceGenerators.Helpers;
@@ -12,13 +14,8 @@ namespace PolySharp.SourceGenerators.Helpers;
 /// A helper type to build sequences of values with pooled buffers.
 /// </summary>
 /// <typeparam name="T">The type of items to create sequences for.</typeparam>
-internal struct ImmutableArrayBuilder<T> : IDisposable
+internal ref struct ImmutableArrayBuilder<T>
 {
-    /// <summary>
-    /// The shared <see cref="ObjectPool{T}"/> instance to share <see cref="Writer"/> objects.
-    /// </summary>
-    private static readonly ObjectPool<Writer> SharedObjectPool = new(static () => new Writer());
-
     /// <summary>
     /// The rented <see cref="Writer"/> instance to use.
     /// </summary>
@@ -30,7 +27,7 @@ internal struct ImmutableArrayBuilder<T> : IDisposable
     /// <returns>A <see cref="ImmutableArrayBuilder{T}"/> instance to write data to.</returns>
     public static ImmutableArrayBuilder<T> Rent()
     {
-        return new(SharedObjectPool.Allocate());
+        return new(new Writer());
     }
 
     /// <summary>
@@ -42,22 +39,21 @@ private ImmutableArrayBuilder(Writer writer)
         this.writer = writer;
     }
 
-    /// <summary>
-    /// Gets the data written to the underlying buffer so far, as a <see cref="ReadOnlySpan{T}"/>.
-    /// </summary>
-    public readonly ReadOnlySpan<T> WrittenSpan
+    /// <inheritdoc cref="ImmutableArray{T}.Builder.Count"/>
+    public readonly int Count
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        get => this.writer!.WrittenSpan;
+        get => this.writer!.Count;
     }
 
     /// <summary>
-    /// Gets the number of elements currently written in the current instance.
+    /// Gets the data written to the underlying buffer so far, as a <see cref="ReadOnlySpan{T}"/>.
     /// </summary>
-    public readonly int Count
+    [UnscopedRef]
+    public readonly ReadOnlySpan<T> WrittenSpan
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        get => this.writer!.WrittenSpan.Length;
+        get => this.writer!.WrittenSpan;
     }
 
     /// <inheritdoc cref="ImmutableArray{T}.Builder.Add(T)"/>
@@ -70,7 +66,7 @@ public readonly void Add(T item)
     /// Adds the specified items to the end of the array.
     /// </summary>
     /// <param name="items">The items to add at the end of the array.</param>
-    public readonly void AddRange(ReadOnlySpan<T> items)
+    public readonly void AddRange(scoped ReadOnlySpan<T> items)
     {
         this.writer!.AddRange(items);
     }
@@ -95,30 +91,25 @@ public override readonly string ToString()
         return this.writer!.WrittenSpan.ToString();
     }
 
-    /// <inheritdoc/>
+    /// <inheritdoc cref="IDisposable.Dispose"/>
     public void Dispose()
     {
         Writer? writer = this.writer;
 
         this.writer = null;
 
-        if (writer is not null)
-        {
-            writer.Clear();
-
-            SharedObjectPool.Free(writer);
-        }
+        writer?.Dispose();
     }
 
     /// <summary>
     /// A class handling the actual buffer writing.
     /// </summary>
-    private sealed class Writer
+    private sealed class Writer : IDisposable
     {
         /// <summary>
         /// The underlying <typeparamref name="T"/> array.
         /// </summary>
-        private T[] array;
+        private T?[]? array;
 
         /// <summary>
         /// The starting offset within <see cref="array"/>.
@@ -130,54 +121,53 @@ private sealed class Writer
         /// </summary>
         public Writer()
         {
-            if (typeof(T) == typeof(char))
-            {
-                this.array = new T[1024];
-            }
-            else
-            {
-                this.array = new T[8];
-            }
-
+            this.array = ArrayPool<T?>.Shared.Rent(typeof(T) == typeof(char) ? 1024 : 8);
             this.index = 0;
         }
 
+        /// <inheritdoc cref="ImmutableArrayBuilder{T}.Count"/>
+        public int Count
+        {
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            get => this.index;
+        }
+
         /// <inheritdoc cref="ImmutableArrayBuilder{T}.WrittenSpan"/>
         public ReadOnlySpan<T> WrittenSpan
         {
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => new(this.array, 0, this.index);
+            get => new(this.array!, 0, this.index);
         }
 
         /// <inheritdoc cref="ImmutableArrayBuilder{T}.Add"/>
         public void Add(T value)
         {
             EnsureCapacity(1);
 
-            this.array[this.index++] = value;
+            this.array![this.index++] = value;
         }
 
         /// <inheritdoc cref="ImmutableArrayBuilder{T}.AddRange"/>
         public void AddRange(ReadOnlySpan<T> items)
         {
             EnsureCapacity(items.Length);
 
-            items.CopyTo(this.array.AsSpan(this.index));
+            items.CopyTo(this.array.AsSpan(this.index)!);
 
             this.index += items.Length;
         }
 
-        /// <summary>
-        /// Clears the items in the current writer.
-        /// </summary>
-        public void Clear()
+        /// <inheritdoc/>
+        public void Dispose()
         {
-            if (typeof(T) != typeof(char))
+            T?[]? array = this.array;
+
+            this.array = null;
+
+            if (array is not null)
             {
-                this.array.AsSpan(0, this.index).Clear();
+                ArrayPool<T?>.Shared.Return(array, clearArray: typeof(T) != typeof(char));
             }
-
-            this.index = 0;
         }
 
         /// <summary>
@@ -187,7 +177,7 @@ public void Clear()
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void EnsureCapacity(int requestedSize)
         {
-            if (requestedSize > this.array.Length - this.index)
+            if (requestedSize > this.array!.Length - this.index)
             {
                 ResizeBuffer(requestedSize);
             }
@@ -201,27 +191,15 @@ private void EnsureCapacity(int requestedSize)
         private void ResizeBuffer(int sizeHint)
         {
             int minimumSize = this.index + sizeHint;
-            int requestedSize = Math.Max(this.array.Length * 2, minimumSize);
 
-            T[] newArray = new T[requestedSize];
+            T?[] oldArray = this.array!;
+            T?[] newArray = ArrayPool<T?>.Shared.Rent(minimumSize);
 
-            Array.Copy(this.array, newArray, this.index);
+            Array.Copy(oldArray, newArray, this.index);
 
             this.array = newArray;
-        }
-    }
-}
 
-/// <summary>
-/// Private helpers for the <see cref="ImmutableArrayBuilder{T}"/> type.
-/// </summary>
-file static class ImmutableArrayBuilder
-{
-    /// <summary>
-    /// Throws an <see cref="ArgumentOutOfRangeException"/> for <c>"index"</c>.
-    /// </summary>
-    public static void ThrowArgumentOutOfRangeExceptionForIndex()
-    {
-        throw new ArgumentOutOfRangeException("index");
+            ArrayPool<T?>.Shared.Return(oldArray, clearArray: typeof(T) != typeof(char));
+        }
     }
 }