From 96606f208defb66e2ea752db004c68dec2a71f8d Mon Sep 17 00:00:00 2001
From: Sergio Pedri <sergio0694@live.com>
Date: Fri, 26 Nov 2021 14:24:01 +0100
Subject: [PATCH] Enable fast paths for IRecipient<TMessage> recipients

---
 .../Messaging/IMessengerExtensions.cs         | 18 ++++++++-
 .../Messaging/WeakReferenceMessenger.cs       | 40 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/CommunityToolkit.Mvvm/Messaging/IMessengerExtensions.cs b/CommunityToolkit.Mvvm/Messaging/IMessengerExtensions.cs
index cce5b277..67aa89ae 100644
--- a/CommunityToolkit.Mvvm/Messaging/IMessengerExtensions.cs
+++ b/CommunityToolkit.Mvvm/Messaging/IMessengerExtensions.cs
@@ -226,7 +226,14 @@ public static class IMessengerExtensions
     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 class IMessengerExtensions
         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>
diff --git a/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs b/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
index e9f57d1c..3688b66a 100644
--- a/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
+++ b/CommunityToolkit.Mvvm/Messaging/WeakReferenceMessenger.cs
@@ -118,6 +118,42 @@ public sealed class WeakReferenceMessenger : IMessenger
         where TRecipient : class
         where TMessage : class
         where TToken : IEquatable<TToken>
+    {
+        Register<TMessage, TToken>(recipient, token, new MessageHandlerDispatcher.For<TRecipient, TMessage>(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="recipient">The recipient that will receive the messages.</param>
+    /// <param name="token">A token used to determine the receiving channel to use.</param>
+    /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
+    /// <remarks>
+    /// This method is a variation of <see cref="Register{TRecipient, TMessage, TToken}(TRecipient, TToken, MessageHandler{TRecipient, TMessage})"/>
+    /// that is specialized for recipients implementing <see cref="IRecipient{TMessage}"/>. See more comments at the top of this type, as well as
+    /// within <see cref="Send{TMessage, TToken}(TMessage, TToken)"/> and in the <see cref="MessageHandlerDispatcher"/> types.
+    /// </remarks>
+    internal void Register<TMessage, TToken>(IRecipient<TMessage> recipient, TToken token)
+        where TMessage : class
+        where TToken : IEquatable<TToken>
+    {
+        Register<TMessage, TToken>(recipient, token, MessageHandlerDispatcher.IRecipient.Instance);
+    }
+
+    /// <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="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="dispatcher">The input <see cref="MessageHandlerDispatcher"/> instance to register.</param>
+    /// <exception cref="InvalidOperationException">Thrown when trying to register the same message twice.</exception>
+    private void Register<TMessage, TToken>(object recipient, TToken token, MessageHandlerDispatcher dispatcher)
+        where TMessage : class
+        where TToken : IEquatable<TToken>
     {
         lock (this.recipientsMap)
         {
@@ -131,7 +167,7 @@ public sealed class WeakReferenceMessenger : IMessenger
             // Fast path for unit tokens
             if (typeof(TToken) == typeof(Unit))
             {
-                if (!mapping.TryAdd(recipient, new MessageHandlerDispatcher.For<TRecipient, TMessage>(handler)))
+                if (!mapping.TryAdd(recipient, dispatcher))
                 {
                     ThrowInvalidOperationExceptionForDuplicateRegistration();
                 }
@@ -150,7 +186,7 @@ public sealed class WeakReferenceMessenger : IMessenger
                 }
 
                 // Store the input handler
-                registeredHandler = new MessageHandlerDispatcher.For<TRecipient, TMessage>(handler);
+                registeredHandler = dispatcher;
             }
         }
     }