Merge pull request #701 from CommunityToolkit/dev/generator-array-pool

Improve ImmutableArrayBuilder<T> type in MVVM Toolkit generators

Author: Sergio0694

src/CommunityToolkit.Mvvm.SourceGenerators/CommunityToolkit.Mvvm.SourceGenerators.projitems

@@ -67,7 +67,6 @@
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\EquatableArray{T}.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\HashCode.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\ImmutableArrayBuilder{T}.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Helpers\ObjectPool{T}.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Input\Models\CanExecuteExpressionType.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Input\Models\CommandInfo.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Input\RelayCommandGenerator.cs" />

src/CommunityToolkit.Mvvm.SourceGenerators/Helpers/ImmutableArrayBuilder{T}.cs

@@ -6,7 +6,9 @@
 // more info in ThirdPartyNotices.txt in the root of the project.
 
 using System;
+using System.Buffers;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
 
 namespace CommunityToolkit.Mvvm.SourceGenerators.Helpers;
@@ -15,13 +17,8 @@ namespace CommunityToolkit.Mvvm.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>
@@ -33,7 +30,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>
@@ -46,7 +43,7 @@ private ImmutableArrayBuilder(Writer writer)
     }
 
     /// <inheritdoc cref="ImmutableArray{T}.Builder.Count"/>
-    public int Count
+    public readonly int Count
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         get => this.writer!.Count;
@@ -55,6 +52,7 @@ public int Count
     /// <summary>
     /// Gets the data written to the underlying buffer so far, as a <see cref="ReadOnlySpan{T}"/>.
     /// </summary>
+    [UnscopedRef]
     public readonly ReadOnlySpan<T> WrittenSpan
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -71,7 +69,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);
     }
@@ -96,30 +94,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"/>.
@@ -131,15 +124,7 @@ 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;
         }
 
@@ -154,38 +139,38 @@ public int Count
         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>
@@ -195,7 +180,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);
             }
@@ -209,13 +194,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;
+
+            ArrayPool<T?>.Shared.Return(oldArray, clearArray: typeof(T) != typeof(char));
         }
     }
 }