Added MessageHandler<in TRecipient> "covariance"

Author: Sergio0694

Microsoft.Toolkit.Mvvm/Messaging/IMessenger.cs

@@ -13,10 +13,13 @@ namespace Microsoft.Toolkit.Mvvm.Messaging
     /// closures: if an instance method on a recipient needs to be invoked it is possible to just
     /// cast the recipient to the right type and then access the local method from that instance.
     /// </summary>
+    /// <typeparam name="TRecipient">The type of recipient for the message.</typeparam>
     /// <typeparam name="TMessage">The type of message to receive.</typeparam>
     /// <param name="recipient">The recipient that is receiving the message.</param>
     /// <param name="message">The message being received.</param>
-    public delegate void MessageHandler<in TMessage>(object recipient, TMessage message);
+    public delegate void MessageHandler<in TRecipient, in TMessage>(TRecipient recipient, TMessage message)
+        where TRecipient : class
+        where TMessage : class;
 
     /// <summary>
     /// An interface for a type providing the ability to exchange messages between different objects.
@@ -39,13 +42,15 @@ bool IsRegistered<TMessage, TToken>(object recipient, TToken token)
         /// <summary>
         /// Registers a recipient for a given type of message.
         /// </summary>
+        /// <typeparam name="TRecipient">The type of recipient for the message.</typeparam>
         /// <typeparam name="TMessage">The type of message to receive.</typeparam>
         /// <typeparam name="TToken">The type of token to use to pick the messages to receive.</typeparam>
         /// <param name="recipient">The recipient that will receive the messages.</param>
         /// <param name="token">A token used to determine the receiving channel to use.</param>
-        /// <param name="handler">The <see cref="MessageHandler{T}"/> to invoke when a message is received.</param>
+        /// <param name="handler">The <see cref="MessageHandler{TRecipient,TMessage}"/> to invoke when a message is received.</param>
         /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
-        void Register<TMessage, TToken>(object recipient, TToken token, MessageHandler<TMessage> handler)
+        void Register<TRecipient, TMessage, TToken>(TRecipient recipient, TToken token, MessageHandler<TRecipient, TMessage> handler)
+            where TRecipient : class
             where TMessage : class
             where TToken : IEquatable<TToken>;
 

Microsoft.Toolkit.Mvvm/Messaging/Messenger.cs

@@ -50,8 +50,8 @@ namespace Microsoft.Toolkit.Mvvm.Messaging
     /// <code>
     /// Messenger.Default.Register&lt;LoginCompletedMessage&gt;(this, Receive);
     /// </code>
-    /// The C# compiler will automatically convert that expression to a <see cref="MessageHandler{TMessage}"/> instance
-    /// compatible with the <see cref="MessengerExtensions.Register{T}(IMessenger,object,MessageHandler{T})"/> method.
+    /// The C# compiler will automatically convert that expression to a <see cref="MessageHandler{TRecipient,TMessage}"/> instance
+    /// compatible with <see cref="MessengerExtensions.Register{TRecipient,TMessage}(IMessenger,TRecipient,MessageHandler{TRecipient,TMessage})"/>.
     /// This will also work if multiple overloads of that method are available, each handling a different
     /// message type: the C# compiler will automatically pick the right one for the current message type.
     /// For info on the other available features, check the <see cref="IMessenger"/> interface.
@@ -66,18 +66,23 @@ public sealed class Messenger : IMessenger
         //                    |   ________(recipients registrations)___________\________/            /          __/
         //                    |  /           _______(channel registrations)_____\___________________/          /
         //                    | /           /                                    \                            /
-        // DictionarySlim<Recipient, DictionarySlim<TToken, MessageHandler<TMessage>>> mapping = Mapping<TMessage, TToken>
-        //                                            /               / \        /                   /
-        //                      ___(Type2.tToken)____/               /   \______/___________________/
+        // DictionarySlim<Recipient, DictionarySlim<TToken, MessageHandler<object, TMessage>>> mapping = Mapping<TMessage, TToken>
+        //                                            /               / \        /                         /
+        //                      ___(Type2.tToken)____/               /   \______/_________________________/
         //                     /________________(Type2.tMessage)____/          /
         //                    /       ________________________________________/
         //                   /       /
         // DictionarySlim<Type2, IMapping> typesMap;
         // --------------------------------------------------------------------------------------------------------
-        // Each combination of <TMessage, TToken> results in a concrete Mapping<TMessage, TToken> type, which holds
-        // the references from registered recipients to handlers. The handlers are stored in a <TToken, Action<TMessage>>
-        // dictionary, so that each recipient can have up to one registered handler for a given token, for each
-        // message type. Each mapping is stored in the types map, which associates each pair of concrete types to its
+        // Each combination of <TMessage, TToken> results in a concrete Mapping<TMessage, TToken> type, which holds the references
+        // from registered recipients to handlers. The handlers are stored in a <TToken, MessageHandler<object, TMessage>>
+        // dictionary, so that each recipient can have up to one registered handler for a given token, for each message type.
+        // Note that the registered handlers are unsafe-cast to a MessageHandler<object, TMessage> instance even if they were
+        // actually of type MessageHandler<TRecipient, TMessage>. This allows the messenger to track and invoke type-specific
+        // handlers without using reflection and without having to capture the input handler in a proxy delegate, causing one
+        // extra memory allocations and adding overhead. The type conversion is guaranteed to be respected due to how the
+        // messenger type itself works - as registered handlers are always invoked on their respective recipients.
+        // Each mapping is stored in the types map, which associates each pair of concrete types to its
         // mapping instance. Mapping instances are exposed as IMapping items, as each will be a closed type over
         // a different combination of TMessage and TToken generic type parameters. Each existing recipient is also stored in
         // the main recipients map, along with a set of all the existing dictionaries of handlers for that recipient (for all
@@ -137,7 +142,8 @@ public bool IsRegistered<TMessage, TToken>(object recipient, TToken token)
         }
 
         /// <inheritdoc/>
-        public void Register<TMessage, TToken>(object recipient, TToken token, MessageHandler<TMessage> handler)
+        public void Register<TRecipient, TMessage, TToken>(TRecipient recipient, TToken token, MessageHandler<TRecipient, TMessage> handler)
+            where TRecipient : class
             where TMessage : class
             where TToken : IEquatable<TToken>
         {
@@ -146,19 +152,20 @@ public void Register<TMessage, TToken>(object recipient, TToken token, MessageHa
                 // Get the <TMessage, TToken> registration list for this recipient
                 Mapping<TMessage, TToken> mapping = GetOrAddMapping<TMessage, TToken>();
                 var key = new Recipient(recipient);
-                ref DictionarySlim<TToken, MessageHandler<TMessage>>? map = ref mapping.GetOrAddValueRef(key);
+                ref DictionarySlim<TToken, MessageHandler<object, TMessage>>? map = ref mapping.GetOrAddValueRef(key);
 
-                map ??= new DictionarySlim<TToken, MessageHandler<TMessage>>();
+                map ??= new DictionarySlim<TToken, MessageHandler<object, TMessage>>();
 
                 // Add the new registration entry
-                ref MessageHandler<TMessage>? registeredHandler = ref map.GetOrAddValueRef(token);
+                ref MessageHandler<object, TMessage>? registeredHandler = ref map.GetOrAddValueRef(token);
 
                 if (!(registeredHandler is null))
                 {
                     ThrowInvalidOperationExceptionForDuplicateRegistration();
                 }
 
-                registeredHandler = handler;
+                // Treat the input delegate as if it was covariant (see comments below in the Send method)
+                registeredHandler = Unsafe.As<MessageHandler<object, TMessage>>(handler);
 
                 // Update the total counter for handlers for the current type parameters
                 mapping.TotalHandlersCount++;
@@ -356,7 +363,7 @@ public void Unregister<TMessage, TToken>(object recipient, TToken token)
 
                 var key = new Recipient(recipient);
 
-                if (!mapping!.TryGetValue(key, out DictionarySlim<TToken, MessageHandler<TMessage>>? dictionary))
+                if (!mapping!.TryGetValue(key, out DictionarySlim<TToken, MessageHandler<object, TMessage>>? dictionary))
                 {
                     return;
                 }
@@ -461,9 +468,14 @@ public unsafe TMessage Send<TMessage, TToken>(TMessage message, TToken token)
                 {
                     // We're doing an unsafe cast to skip the type checks again.
                     // See the comments in the UnregisterAll method for more info.
-                    MessageHandler<TMessage> handler = Unsafe.As<MessageHandler<TMessage>>(Unsafe.Add(ref handlersRef, (IntPtr)(void*)(uint)j));
-
-                    handler(Unsafe.Add(ref recipientsRef, (IntPtr)(void*)(uint)j), message);
+                    object handler = Unsafe.Add(ref handlersRef, (IntPtr)(void*)(uint)j);
+                    object recipient = Unsafe.Add(ref recipientsRef, (IntPtr)(void*)(uint)j);
+
+                    // Here we perform an unsafe cast to enable covariance for delegate types.
+                    // We know that the input recipient will always respect the type constraints
+                    // of each original input delegate, and doing so allows us to still invoke
+                    // them all from here without worrying about specific generic type arguments.
+                    Unsafe.As<MessageHandler<object, TMessage>>(handler)(recipient, message);
                 }
             }
             finally
@@ -559,7 +571,7 @@ private Mapping<TMessage, TToken> GetOrAddMapping<TMessage, TToken>()
         /// This type is defined for simplicity and as a workaround for the lack of support for using type aliases
         /// over open generic types in C# (using type aliases can only be used for concrete, closed types).
         /// </remarks>
-        private sealed class Mapping<TMessage, TToken> : DictionarySlim<Recipient, DictionarySlim<TToken, MessageHandler<TMessage>>>, IMapping
+        private sealed class Mapping<TMessage, TToken> : DictionarySlim<Recipient, DictionarySlim<TToken, MessageHandler<object, TMessage>>>, IMapping
             where TMessage : class
             where TToken : IEquatable<TToken>
         {

Microsoft.Toolkit.Mvvm/Messaging/MessengerExtensions.cs

@@ -174,7 +174,7 @@ static Action<IMessenger, object, TToken> GetRegistrationAction(Type type, Metho
         public static void Register<TMessage>(this IMessenger messenger, IRecipient<TMessage> recipient)
             where TMessage : class
         {
-            messenger.Register<TMessage, Unit>(recipient, default, (r, m) => Unsafe.As<IRecipient<TMessage>>(r).Receive(m));
+            messenger.Register<IRecipient<TMessage>, TMessage, Unit>(recipient, default, (r, m) => r.Receive(m));
         }
 
         /// <summary>
@@ -191,7 +191,7 @@ public static void Register<TMessage, TToken>(this IMessenger messenger, IRecipi
             where TMessage : class
             where TToken : IEquatable<TToken>
         {
-            messenger.Register<TMessage, TToken>(recipient, token, (r, m) => Unsafe.As<IRecipient<TMessage>>(r).Receive(m));
+            messenger.Register<IRecipient<TMessage>, TMessage, TToken>(recipient, token, (r, m) => r.Receive(m));
         }
 
         /// <summary>
@@ -200,15 +200,49 @@ public static void Register<TMessage, TToken>(this IMessenger messenger, IRecipi
         /// <typeparam name="TMessage">The type of message to receive.</typeparam>
         /// <param name="messenger">The <see cref="IMessenger"/> instance to use to register the recipient.</param>
         /// <param name="recipient">The recipient that will receive the messages.</param>
-        /// <param name="handler">The <see cref="MessageHandler{T}"/> to invoke when a message is received.</param>
+        /// <param name="handler">The <see cref="MessageHandler{TRecipient,TMessage}"/> to invoke when a message is received.</param>
         /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
         /// <remarks>This method will use the default channel to perform the requested registration.</remarks>
-        public static void Register<TMessage>(this IMessenger messenger, object recipient, MessageHandler<TMessage> handler)
+        public static void Register<TMessage>(this IMessenger messenger, object recipient, MessageHandler<object, TMessage> handler)
             where TMessage : class
         {
             messenger.Register(recipient, default(Unit), handler);
         }
 
+        /// <summary>
+        /// Registers a recipient for a given type of message.
+        /// </summary>
+        /// <typeparam name="TRecipient">The type of recipient for the message.</typeparam>
+        /// <typeparam name="TMessage">The type of message to receive.</typeparam>
+        /// <param name="messenger">The <see cref="IMessenger"/> instance to use to register the recipient.</param>
+        /// <param name="recipient">The recipient that will receive the messages.</param>
+        /// <param name="handler">The <see cref="MessageHandler{TRecipient,TMessage}"/> to invoke when a message is received.</param>
+        /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
+        /// <remarks>This method will use the default channel to perform the requested registration.</remarks>
+        public static void Register<TRecipient, TMessage>(this IMessenger messenger, TRecipient recipient, MessageHandler<TRecipient, TMessage> handler)
+            where TRecipient : class
+            where TMessage : class
+        {
+            messenger.Register(recipient, default(Unit), handler);
+        }
+
+        /// <summary>
+        /// Registers a recipient for a given type of message.
+        /// </summary>
+        /// <typeparam name="TMessage">The type of message to receive.</typeparam>
+        /// <typeparam name="TToken">The type of token to use to pick the messages to receive.</typeparam>
+        /// <param name="messenger">The <see cref="IMessenger"/> instance to use to register the recipient.</param>
+        /// <param name="recipient">The recipient that will receive the messages.</param>
+        /// <param name="token">A token used to determine the receiving channel to use.</param>
+        /// <param name="handler">The <see cref="MessageHandler{TRecipient,TMessage}"/> to invoke when a message is received.</param>
+        /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
+        public static void Register<TMessage, TToken>(this IMessenger messenger, object recipient, TToken token, MessageHandler<object, TMessage> handler)
+            where TMessage : class
+            where TToken : IEquatable<TToken>
+        {
+            messenger.Register(recipient, token, handler);
+        }
+
         /// <summary>
         /// Unregisters a recipient from messages of a given type.
         /// </summary>

UnitTests/UnitTests.Shared/Mvvm/Test_Messenger.cs

@@ -297,8 +297,25 @@ public void Test_Messenger_IRecipient_SomeMessages_WithToken()
             Assert.IsFalse(messenger.IsRegistered<MessageB>(recipient));
         }
 
+        [TestCategory("Mvvm")]
+        [TestMethod]
+        public void Test_Messenger_RegisterWithTypeParameter()
+        {
+            var messenger = new Messenger();
+            var recipient = new RecipientWithNoMessages { Number = 42 };
+
+            int number = 0;
+
+            messenger.Register<RecipientWithNoMessages, MessageA>(recipient, (r, m) => number = r.Number);
+
+            messenger.Send<MessageA>();
+
+            Assert.AreEqual(number, 42);
+        }
+
         public sealed class RecipientWithNoMessages
         {
+            public int Number { get; set; }
         }
 
         public sealed class RecipientWithSomeMessages