Optimize WeakReferenceMessenger

CommunityToolkit.Mvvm/Attributes/MaybeNullWhenAttribute.cs

@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#if NETSTANDARD2_0
+
+namespace System.Diagnostics.CodeAnalysis;
+
+/// <summary>
+/// Specifies that when a method returns <see cref="ReturnValue"/>, the parameter may be null even if the corresponding type disallows it.
+/// </summary>
+/// <remarks>Internal copy from the BCL attribute.</remarks>
+[AttributeUsage(AttributeTargets.Parameter, Inherited = false)]
+internal sealed class MaybeNullWhenAttribute : Attribute
+{
+    /// <summary>
+    /// Initializes the attribute with the specified return value condition.
+    /// </summary>
+    /// <param name="returnValue">The return value condition. If the method returns this value, the associated parameter may be null.</param>
+    public MaybeNullWhenAttribute(bool returnValue)
+    {
+        ReturnValue = returnValue;
+    }
+
+    /// <summary>
+    /// Gets the return value condition.
+    /// </summary>
+    public bool ReturnValue { get; }
+}
+
+#endif

CommunityToolkit.Mvvm/Attributes/MemberNotNullAttribute.cs

@@ -0,0 +1,40 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#if !NET6_0_OR_GREATER
+
+namespace System.Diagnostics.CodeAnalysis;
+
+/// <summary>
+/// Specifies that the method or property will ensure that the listed field and property members have not-null values.
+/// </summary>
+/// <remarks>Internal copy from the BCL attribute.</remarks>
+[AttributeUsage(AttributeTargets.Method | AttributeTargets.Property, Inherited = false, AllowMultiple = true)]
+internal sealed class MemberNotNullAttribute : Attribute
+{
+    /// <summary>
+    /// Initializes the attribute with a field or property member.
+    /// </summary>
+    /// <param name="member">The field or property member that is promised to be not-null.</param>
+    public MemberNotNullAttribute(string member)
+    {
+        Members = new[] { member };
+    }
+
+    /// <summary>
+    /// Initializes the attribute with the list of field and property members.
+    /// </summary>
+    /// <param name="members">The list of field and property members that are promised to be not-null.</param>
+    public MemberNotNullAttribute(params string[] members)
+    {
+        Members = members;
+    }
+
+    /// <summary>
+    /// Gets field or property member names.
+    /// </summary>
+    public string[] Members { get; }
+}
+
+#endif

CommunityToolkit.Mvvm/CommunityToolkit.Mvvm.csproj

@@ -26,12 +26,13 @@
     <PackageReference Include="Microsoft.Bcl.AsyncInterfaces" Version="5.0.0" />
     <PackageReference Include="System.ComponentModel.Annotations" Version="5.0.0" />
     <PackageReference Include="System.Memory" Version="4.5.4" />
+    <PackageReference Include="System.Runtime.CompilerServices.Unsafe" Version="6.0.0" />
   </ItemGroup>
 
   <!-- .NET Standard 2.1 doesn't have the Unsafe type -->
   <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.1'">
     <PackageReference Include="System.ComponentModel.Annotations" Version="5.0.0" />
-    <PackageReference Include="System.Runtime.CompilerServices.Unsafe" Version="5.0.0" />
+    <PackageReference Include="System.Runtime.CompilerServices.Unsafe" Version="6.0.0" />
   </ItemGroup>
 
   <!-- Source generator project reference for packing -->

CommunityToolkit.Mvvm/Messaging/IMessengerExtensions.cs

@@ -226,7 +226,14 @@ from registrationMethod in registrationMethods
     public static void Register<TMessage>(this IMessenger messenger, IRecipient<TMessage> recipient)
         where TMessage : class
     {
-        messenger.Register<IRecipient<TMessage>, TMessage, Unit>(recipient, default, static (r, m) => r.Receive(m));
+        if (messenger is WeakReferenceMessenger weakReferenceMessenger)
+        {
+            weakReferenceMessenger.Register<TMessage, Unit>(recipient, default);
+        }
+        else
+        {
+            messenger.Register<IRecipient<TMessage>, TMessage, Unit>(recipient, default, static (r, m) => r.Receive(m));
+        }
     }
 
     /// <summary>
@@ -243,7 +250,14 @@ public static void Register<TMessage, TToken>(this IMessenger messenger, IRecipi
         where TMessage : class
         where TToken : IEquatable<TToken>
     {
-        messenger.Register<IRecipient<TMessage>, TMessage, TToken>(recipient, token, static (r, m) => r.Receive(m));
+        if (messenger is WeakReferenceMessenger weakReferenceMessenger)
+        {
+            weakReferenceMessenger.Register(recipient, token);
+        }
+        else
+        {
+            messenger.Register<IRecipient<TMessage>, TMessage, TToken>(recipient, token, static (r, m) => r.Receive(m));
+        }
     }
 
     /// <summary>

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"/>

CommunityToolkit.Mvvm/Messaging/Internals/MessageHandlerDispatcher.cs

@@ -0,0 +1,56 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Runtime.CompilerServices;
+
+namespace CommunityToolkit.Mvvm.Messaging.Internals;
+
+/// <summary>
+/// A dispatcher type that invokes a given <see cref="MessageHandler{TRecipient, TMessage}"/> callback.
+/// </summary>
+/// <remarks>
+/// This type is used to avoid type aliasing with <see cref="Unsafe.As{T}(object)"/> when the generic
+/// arguments are not known. Additionally, this is an abstract class and not an interface so that when
+/// <see cref="Invoke(object, object)"/> is called, virtual dispatch will be used instead of interface
+/// stub dispatch, which is much slower and with more indirections.
+/// </remarks>
+internal abstract class MessageHandlerDispatcher
+{
+    /// <summary>
+    /// Invokes the current callback on a target recipient, with a specified message.
+    /// </summary>
+    /// <param name="recipient">The target recipient for the message.</param>
+    /// <param name="message">The message being broadcast.</param>
+    public abstract void Invoke(object recipient, object message);
+
+    /// <summary>
+    /// A generic version of <see cref="MessageHandlerDispatcher"/>.
+    /// </summary>
+    /// <typeparam name="TRecipient">The type of recipient for the message.</typeparam>
+    /// <typeparam name="TMessage">The type of message to receive.</typeparam>
+    public sealed class For<TRecipient, TMessage> : MessageHandlerDispatcher
+        where TRecipient : class
+        where TMessage : class
+    {
+        /// <summary>
+        /// The underlying <see cref="MessageHandler{TRecipient, TMessage}"/> callback to invoke.
+        /// </summary>
+        private readonly MessageHandler<TRecipient, TMessage> handler;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="For{TRecipient, TMessage}"/> class.
+        /// </summary>
+        /// <param name="handler">The input <see cref="MessageHandler{TRecipient, TMessage}"/> instance.</param>
+        public For(MessageHandler<TRecipient, TMessage> handler)
+        {
+            this.handler = handler;
+        }
+
+        /// <inheritdoc/>
+        public override void Invoke(object recipient, object message)
+        {
+            this.handler(Unsafe.As<TRecipient>(recipient), Unsafe.As<TMessage>(message));
+        }
+    }
+}