// 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.Buffers;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Threading;
using CommunityToolkit.Mvvm.Messaging.Internals;
namespace CommunityToolkit.Mvvm.Messaging;
///
/// A class providing a reference implementation for the interface.
///
///
/// This implementation uses strong references to track the registered
/// recipients, so it is necessary to manually unregister them when they're no longer needed.
///
public sealed class StrongReferenceMessenger : IMessenger
{
// This messenger uses the following logic to link stored instances together:
// --------------------------------------------------------------------------------------------------------
// Dictionary2> recipientsMap;
// | \________________[*]IDictionary2>
// | \_______________[*]IDictionary2 /
// | \_________/_________/___ /
// |\ _(recipients registrations)_/ / \ /
// | \__________________ / _____(channel registrations)_____/______\____/
// | \ / / __________________________/ \
// | / / / \
// | Dictionary2 mapping = Mapping________________\
// | __________________/ / | / \
// |/ / | / \
// Dictionary2> mapping = Mapping____________\
// / / / /
// ___(Type2.TToken)____/ / / /
// /________________(Type2.TMessage)_______/_______/__/
// / ________________________________/
// / /
// Dictionary2 typesMap;
// --------------------------------------------------------------------------------------------------------
// Each combination of results in a concrete Mapping type (if TToken is Unit) or Mapping type,
// which holds the references from registered recipients to handlers. Mapping is used when the default channel is being
// requested, as in that case there will only ever be up to a handler per recipient, per message type. In that case,
// each recipient will only track the message dispatcher (stored as an object?, see notes below), instead of a dictionary
// mapping each TToken value to the corresponding dispatcher for that recipient. When a custom channel is used, the
// dispatchers are stored in a dictionary, so that each recipient can have up to one registered handler
// for a given token, for each message type. Note that the registered dispatchers are only stored as object references, as
// they can either be null or a MessageHandlerDispatcher.For instance.
//
// The first case happens if the handler was registered through an IRecipient instance, while the second one is
// used to wrap input MessageHandler instances. The MessageHandlerDispatcher.For
// instances will just be cast to MessageHandlerDispatcher when invoking it. This allows users to retain type information on
// each registered recipient, instead of having to manually cast each recipient to the right type within the handler
// (additionally, using double dispatch here avoids the need to alias delegate types). 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 (or just of TMessage, for the default channel). 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 message types and token types, if any).
//
// A recipient is stored in the main map as long as it has at least one registered handler in any of the existing mappings for every
// message/token type combination. The shared map is used to access the set of all registered handlers for a given recipient, without
// having to know in advance the type of message or token being used for the registration, and without having to use reflection. This
// is the same approach used in the types map, as we expose saved items as IMapping values too.
//
// Note that each mapping stored in the associated set for each recipient also indirectly implements either IDictionary2
// or IDictionary2, with any token type currently in use by that recipient (or none, if using the default channel). This allows
// to retrieve the type-closed mappings of registered handlers with a given token type, for any message type, for every receiver, again
// without having to use reflection. This shared map is used to unregister messages from a given recipients either unconditionally, by
// message type, by token, or for a specific pair of message type and token value.
///
/// The collection of currently registered recipients, with a link to their linked message receivers.
///
///
/// This collection is used to allow reflection-free access to all the existing
/// registered recipients from and other methods in this type,
/// so that all the existing handlers can be removed without having to dynamically create
/// the generic types for the containers of the various dictionaries mapping the handlers.
///
private readonly Dictionary2> recipientsMap = new();
///
/// The and instance for types combination.
///
///
/// The values are just of type as we don't know the type parameters in advance.
/// Each method relies on to get the type-safe instance of the
/// or class for each pair of generic arguments in use.
///
private readonly Dictionary2 typesMap = new();
///
/// Gets the default instance.
///
public static StrongReferenceMessenger Default { get; } = new();
///
public bool IsRegistered(object recipient, TToken token)
where TMessage : class
where TToken : IEquatable
{
ArgumentNullException.ThrowIfNull(recipient);
ArgumentNullException.For.ThrowIfNull(token);
lock (this.recipientsMap)
{
if (typeof(TToken) == typeof(Unit))
{
if (!TryGetMapping(out Mapping? mapping))
{
return false;
}
Recipient key = new(recipient);
return mapping.ContainsKey(key);
}
else
{
if (!TryGetMapping(out Mapping? mapping))
{
return false;
}
Recipient key = new(recipient);
return
mapping.TryGetValue(key, out Dictionary2? handlers) &&
handlers.ContainsKey(token);
}
}
}
///
public void Register(TRecipient recipient, TToken token, MessageHandler handler)
where TRecipient : class
where TMessage : class
where TToken : IEquatable
{
ArgumentNullException.ThrowIfNull(recipient);
ArgumentNullException.For.ThrowIfNull(token);
ArgumentNullException.ThrowIfNull(handler);
Register(recipient, token, new MessageHandlerDispatcher.For(handler));
}
///
internal void Register(IRecipient recipient, TToken token)
where TMessage : class
where TToken : IEquatable
{
Register(recipient, token, null);
}
///
/// Registers a recipient for a given type of message.
///
/// The type of message to receive.
/// The type of token to use to pick the messages to receive.
/// The recipient that will receive the messages.
/// A token used to determine the receiving channel to use.
/// The input instance to register, or null.
/// Thrown when trying to register the same message twice.
private void Register(object recipient, TToken token, MessageHandlerDispatcher? dispatcher)
where TMessage : class
where TToken : IEquatable
{
lock (this.recipientsMap)
{
Recipient key = new(recipient);
IMapping mapping;
// Fast path for unit tokens
if (typeof(TToken) == typeof(Unit))
{
// Get the registration list for this recipient
Mapping underlyingMapping = GetOrAddMapping();
ref object? registeredHandler = ref underlyingMapping.GetOrAddValueRef(key);
if (registeredHandler is not null)
{
ThrowInvalidOperationExceptionForDuplicateRegistration();
}
// Store the input handler
registeredHandler = dispatcher;
mapping = underlyingMapping;
}
else
{
// Get the registration list for this recipient
Mapping underlyingMapping = GetOrAddMapping();
ref Dictionary2? map = ref underlyingMapping.GetOrAddValueRef(key);
map ??= new Dictionary2();
// Add the new registration entry
ref object? registeredHandler = ref map.GetOrAddValueRef(token);
if (registeredHandler is not null)
{
ThrowInvalidOperationExceptionForDuplicateRegistration();
}
registeredHandler = dispatcher;
mapping = underlyingMapping;
}
// Make sure this registration map is tracked for the current recipient
ref HashSet? set = ref this.recipientsMap.GetOrAddValueRef(key);
set ??= new HashSet();
_ = set.Add(mapping);
}
}
///
public void UnregisterAll(object recipient)
{
ArgumentNullException.ThrowIfNull(recipient);
lock (this.recipientsMap)
{
// If the recipient has no registered messages at all, ignore
Recipient key = new(recipient);
if (!this.recipientsMap.TryGetValue(key, out HashSet? set))
{
return;
}
// Removes all the lists of registered handlers for the recipient
foreach (IMapping mapping in set)
{
if (mapping.TryRemove(key) &&
mapping.Count == 0)
{
// Maps here are really of type Mapping<,> and with unknown type arguments.
// If after removing the current recipient a given map becomes empty, it means
// that there are no registered recipients at all for a given pair of message
// and token types. In that case, we also remove the map from the types map.
// The reason for keeping a key in each mapping is that removing items from a
// dictionary (a hashed collection) only costs O(1) in the best case, while
// if we had tried to iterate the whole dictionary every time we would have
// paid an O(n) minimum cost for each single remove operation.
_ = this.typesMap.TryRemove(mapping.TypeArguments);
}
}
// Remove the associated set in the recipients map
_ = this.recipientsMap.TryRemove(key);
}
}
///
public void UnregisterAll(object recipient, TToken token)
where TToken : IEquatable
{
ArgumentNullException.ThrowIfNull(recipient);
ArgumentNullException.For.ThrowIfNull(token);
// This method is never called with the unit type, so this path is not implemented. This
// exception should not ever be thrown, it's here just to double check for regressions in
// case a bug was introduced that caused this path to somehow be invoked with the Unit type.
// This type is internal, so consumers of the library would never be able to pass it here,
// and there are (and shouldn't be) any APIs publicly exposed from the library that would
// cause this path to be taken either. When using the default channel, only UnregisterAll(object)
// is supported, which would just unregister all recipients regardless of the selected channel.
if (typeof(TToken) == typeof(Unit))
{
throw new NotImplementedException();
}
bool lockTaken = false;
object[]? maps = null;
int i = 0;
// We use an explicit try/finally block here instead of the lock syntax so that we can use a single
// one both to release the lock and to clear the rented buffer and return it to the pool. The reason
// why we're declaring the buffer here and clearing and returning it in this outer finally block is
// that doing so doesn't require the lock to be kept, and releasing it before performing this last
// step reduces the total time spent while the lock is acquired, which in turn reduces the lock
// contention in multi-threaded scenarios where this method is invoked concurrently.
try
{
Monitor.Enter(this.recipientsMap, ref lockTaken);
// Get the shared set of mappings for the recipient, if present
Recipient key = new(recipient);
if (!this.recipientsMap.TryGetValue(key, out HashSet? set))
{
return;
}
// Copy the candidate mappings for the target recipient to a local array, as we can't modify the
// contents of the set while iterating it. The rented buffer is oversized and will also include
// mappings for handlers of messages that are registered through a different token. Note that
// we're using just an object array to minimize the number of total rented buffers, that would
// just remain in the shared pool unused, other than when they are rented here. Instead, we're
// using a type that would possibly also be used by the users of the library, which increases
// the opportunities to reuse existing buffers for both. When we need to reference an item
// stored in the buffer with the type we know it will have, we use Unsafe.As to avoid the
// expensive type check in the cast, since we already know the assignment will be valid.
maps = ArrayPool