Remove Unsafe.As<T>(object) delegate type aliasing

CommunityToolkit.Mvvm/Messaging/Internals/MessageHandlerDispatcher.cs

@@ -0,0 +1,76 @@
+// 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;
+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));
+        }
+    }
+
+    /// <summary>
+    /// A marker type implementing <see cref="MessageHandlerDispatcher"/> for <see cref="IRecipient{TMessage}"/> types.
+    /// </summary>
+    public sealed class IRecipient : MessageHandlerDispatcher
+    {
+        /// <summary>
+        /// Gets the shared marker instance.
+        /// </summary>
+        public static IRecipient Instance { get; } = new();
+
+        /// <inheritdoc/>
+        public override void Invoke(object recipient, object message)
+        {
+            // This method always throws, as callers are intended to do explicit guarded
+            // devirtualization on this type and then just skip the indirection entirely.
+            throw new NotSupportedException();
+        }
+    }
+}

CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs

@@ -33,12 +33,12 @@ public sealed class WeakReferenceMessenger : IMessenger
 {
     // This messenger uses the following logic to link stored instances together:
     // --------------------------------------------------------------------------------------------------------
-    //                         Dictionary2<TToken, MessageHandler<TRecipient, TMessage>> mapping
-    //                                        /                                   /         /
-    //                   ___(Type2.TToken)___/                                   /         /
-    //                  /_________________(Type2.TMessage)______________________/         /
-    //                 /                                    _____________________________/
-    //                /                                    /                   \_______MessageHandler<TRecipient, TMessage>
+    //                          Dictionary2<TToken, MessageHandlerDispatcher> mapping
+    //                                        /                   /             /
+    //                   ___(Type2.TToken)___/                   /             /         ___(if Type2.TToken is Unit)
+    //                  /_________(Type2.TMessage)______________/             /         /
+    //                 /                                    _________________/___MessageHandler<TRecipient, TMessage>
+    //                /                                    /
     // Dictionary2<Type2, ConditionalWeakTable<object, object>> recipientsMap;
     // --------------------------------------------------------------------------------------------------------
     // Just like in the strong reference variant, each pair of message and token types is used as a key in the
@@ -131,7 +131,7 @@ public void Register<TRecipient, TMessage, TToken>(TRecipient recipient, TToken
             // Fast path for unit tokens
             if (typeof(TToken) == typeof(Unit))
             {
-                if (!mapping.TryAdd(recipient, handler))
+                if (!mapping.TryAdd(recipient, new MessageHandlerDispatcher.For<TRecipient, TMessage>(handler)))
                 {
                     ThrowInvalidOperationExceptionForDuplicateRegistration();
                 }
@@ -150,7 +150,7 @@ public void Register<TRecipient, TMessage, TToken>(TRecipient recipient, TToken
                 }
 
                 // Store the input handler
-                registeredHandler = handler;
+                registeredHandler = new MessageHandlerDispatcher.For<TRecipient, TMessage>(handler);
             }
         }
     }
@@ -279,7 +279,34 @@ public TMessage Send<TMessage, TToken>(TMessage message, TToken token)
 
             for (int j = 0; j < i; j++)
             {
-                Unsafe.As<MessageHandler<object, TMessage>>(pairs[2 * j])(pairs[(2 * j) + 1], message);
+                object recipient = pairs[(2 * j) + 1];
+                object handler = pairs[2 * j];
+
+                // This doesn't use reflection: a GetType() call being immediately compared to
+                // a specific type just results in a direct comparison of the method table pointer
+                // with a constant address corresponding to the method table address for that type.
+                // That is, for instance on x64 and assuming handler is in rcx, this will produce:
+                // =============================
+                // L0000: mov rax, 0x7ffcbc87cc98
+                // L000a: cmp [rcx], rax
+                // =============================
+                // Which is extremely fast. The reason for this conditional check in the first place
+                // is that we're doing manual guarded devirtualization: if the handler is the marker
+                // type and not an actual handler then we know that the recipient implements
+                // IRecipient<TMessage>, so we can just cast to it and invoke it directly. This avoids
+                // having to store the proxy callback when registering, and also skips an indirection
+                // (invoking the delegate that then invokes the actual method). Additional note: this
+                // pattern ensures that both casts below do not actually alias incompatible reference
+                // types (as in, they would both succeed if they were safe casts), which lets the code
+                // not rely on undefined behavior to run correctly (ie. we're not aliasing delegates).
+                if (handler.GetType() == typeof(MessageHandlerDispatcher.IRecipient))
+                {
+                    Unsafe.As<IRecipient<TMessage>>(recipient).Receive(message);
+                }
+                else
+                {
+                    Unsafe.As<MessageHandlerDispatcher>(handler).Invoke(recipient, message);
+                }
             }
         }
         finally