From 15bd652a572fbf7735ba53e62945473d03fa6add Mon Sep 17 00:00:00 2001
From: Sergio Pedri <sergio0694@live.com>
Date: Sat, 27 Nov 2021 13:48:02 +0100
Subject: [PATCH] Skip array covariance checks in ArrayPoolBufferWriter<T>

---
 .../Internals/ArrayPoolBufferWriter{T}.cs     | 40 +++++++++++++------
 .../Messaging/WeakReferenceMessenger.cs       |  6 +--
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/CommunityToolkit.Mvvm/Messaging/Internals/ArrayPoolBufferWriter{T}.cs b/CommunityToolkit.Mvvm/Messaging/Internals/ArrayPoolBufferWriter{T}.cs
index 3c8acf76a..5a8672990 100644
--- a/CommunityToolkit.Mvvm/Messaging/Internals/ArrayPoolBufferWriter{T}.cs
+++ b/CommunityToolkit.Mvvm/Messaging/Internals/ArrayPoolBufferWriter{T}.cs
@@ -4,7 +4,6 @@
 
 using System;
 using System.Buffers;
-using System.Diagnostics.Contracts;
 using System.Runtime.CompilerServices;
 
 namespace CommunityToolkit.Mvvm.Messaging.Internals;
@@ -30,6 +29,12 @@ internal ref struct ArrayPoolBufferWriter<T>
     /// </summary>
     private T[] array;
 
+    /// <summary>
+    /// The span mapping to <see cref="array"/>.
+    /// </summary>
+    /// <remarks>All writes are done through this to avoid covariance checks.</remarks>
+    private Span<T> span;
+
     /// <summary>
     /// The starting offset within <see cref="array"/>.
     /// </summary>
@@ -38,12 +43,11 @@ internal ref struct ArrayPoolBufferWriter<T>
     /// <summary>
     /// Creates a new instance of the <see cref="ArrayPoolBufferWriter{T}"/> struct.
     /// </summary>
-    /// <returns>A new <see cref="ArrayPoolBufferWriter{T}"/> instance.</returns>
-    [Pure]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static ArrayPoolBufferWriter<T> Create()
+    public ArrayPoolBufferWriter()
     {
-        return new() { array = ArrayPool<T>.Shared.Rent(DefaultInitialBufferSize) };
+        this.span = this.array = ArrayPool<T>.Shared.Rent(DefaultInitialBufferSize);
+        this.index = 0;
     }
 
     /// <summary>
@@ -52,7 +56,7 @@ public static ArrayPoolBufferWriter<T> Create()
     public ReadOnlySpan<T> Span
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        get => this.array.AsSpan(0, this.index);
+        get => this.span.Slice(0, this.index);
     }
 
     /// <summary>
@@ -62,12 +66,19 @@ public ReadOnlySpan<T> Span
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public void Add(T item)
     {
-        if (this.index == this.array.Length)
+        Span<T> span = this.span;
+        int index = this.index;
+
+        if ((uint)index < (uint)span.Length)
         {
-            ResizeBuffer();
-        }
+            span[index] = item;
 
-        this.array[this.index++] = item;
+            this.index = index + 1;
+        }
+        else
+        {
+            ResizeBufferAndAdd(item);
+        }
     }
 
     /// <summary>
@@ -81,10 +92,11 @@ public void Reset()
     }
 
     /// <summary>
-    /// Resizes <see cref="array"/> when there is no space left for new items.
+    /// Resizes <see cref="array"/> when there is no space left for new items, then adds one
     /// </summary>
+    /// <param name="item">The item to add.</param>
     [MethodImpl(MethodImplOptions.NoInlining)]
-    private void ResizeBuffer()
+    private void ResizeBufferAndAdd(T item)
     {
         T[] rent = ArrayPool<T>.Shared.Rent(this.index << 2);
 
@@ -93,7 +105,9 @@ private void ResizeBuffer()
 
         ArrayPool<T>.Shared.Return(this.array);
 
-        this.array = rent;
+        this.span = this.array = rent;
+
+        this.span[this.index++] = item;
     }
 
     /// <inheritdoc cref="IDisposable.Dispose"/>
diff --git a/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs b/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
index 1825aa430..ce65dd14c 100644
--- a/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
+++ b/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
@@ -276,7 +276,7 @@ public TMessage Send<TMessage, TToken>(TMessage message, TToken token)
                 return message;
             }
 
-            bufferWriter = ArrayPoolBufferWriter<object>.Create();
+            bufferWriter = new ArrayPoolBufferWriter<object>();
 
             // We need a local, temporary copy of all the pending recipients and handlers to
             // invoke, to avoid issues with handlers unregistering from messages while we're
@@ -436,8 +436,8 @@ private void CleanupWithNonBlockingLock()
     /// </summary>
     private void CleanupWithoutLock()
     {
-        using ArrayPoolBufferWriter<Type2> type2s = ArrayPoolBufferWriter<Type2>.Create();
-        using ArrayPoolBufferWriter<object> emptyRecipients = ArrayPoolBufferWriter<object>.Create();
+        using ArrayPoolBufferWriter<Type2> type2s = new();
+        using ArrayPoolBufferWriter<object> emptyRecipients = new();
 
         Dictionary2<Type2, RecipientsTable>.Enumerator type2Enumerator = this.recipientsMap.GetEnumerator();