Optimize ConditionalWeakTable2<,>.Enumerator

Also done some refactoring to streamline code across targets

CommunityToolkit.Mvvm/Messaging/Internals/System/Collections.Generic/Dictionary2.cs

@@ -321,19 +321,19 @@ public bool MoveNext()
         /// <summary>
         /// Gets the current key.
         /// </summary>
-        public TKey Key
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TKey GetKey()
         {
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => this.entries[this.index - 1].Key;
+            return this.entries[this.index - 1].Key;
         }
 
         /// <summary>
         /// Gets the current value.
         /// </summary>
-        public TValue Value
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TValue GetValue()
         {
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => this.entries[this.index - 1].Value!;
+            return this.entries[this.index - 1].Value!;
         }
     }
 

CommunityToolkit.Mvvm/Messaging/Internals/System/Runtime.CompilerServices/ConditionalWeakTable2{TKey,TValue}.Proxy.cs

@@ -0,0 +1,108 @@
+// 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_1
+
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Diagnostics.Contracts;
+
+namespace System.Runtime.CompilerServices;
+
+/// <summary>
+/// A wrapper for <see cref="ConditionalWeakTable{TKey,TValue}"/> with a custom enumerator.
+/// </summary>
+/// <typeparam name="TKey">Tke key of items to store in the table.</typeparam>
+/// <typeparam name="TValue">The values to store in the table.</typeparam>
+internal sealed class ConditionalWeakTable2<TKey, TValue>
+    where TKey : class
+    where TValue : class?
+{
+    /// <summary>
+    /// The underlying <see cref="ConditionalWeakTable{TKey,TValue}"/> instance.
+    /// </summary>
+    private readonly ConditionalWeakTable<TKey, TValue> table = new();
+
+    /// <inheritdoc cref="ConditionalWeakTable{TKey,TValue}.TryGetValue"/>
+    public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
+    {
+        return this.table.TryGetValue(key, out value);
+    }
+
+    /// <inheritdoc cref="ConditionalWeakTableExtensions.TryAdd{TKey, TValue}(ConditionalWeakTable{TKey, TValue}, TKey, TValue)"/>
+    public bool TryAdd(TKey key, TValue value)
+    {
+        return this.table.TryAdd(key, value);
+    }
+
+    /// <inheritdoc cref="ConditionalWeakTable{TKey,TValue}.GetValue"/>
+    public TValue GetValue(TKey key, ConditionalWeakTable<TKey, TValue>.CreateValueCallback createValueCallback)
+    {
+        return this.table.GetValue(key, createValueCallback);
+    }
+
+    /// <inheritdoc cref="ConditionalWeakTable{TKey,TValue}.Remove"/>
+    public bool Remove(TKey key)
+    {
+        return this.table.Remove(key);
+    }
+
+    /// <inheritdoc cref="IEnumerable{T}.GetEnumerator"/>
+    [Pure]
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public Enumerator GetEnumerator() => new(this);
+
+    /// <summary>
+    /// A custom enumerator that traverses items in a <see cref="ConditionalWeakTable2{TKey, TValue}"/> instance.
+    /// </summary>
+    public ref struct Enumerator
+    {
+        /// <summary>
+        /// The wrapped <see cref="IEnumerator{T}"/> instance for the enumerator.
+        /// </summary>
+        private readonly IEnumerator<KeyValuePair<TKey, TValue>> enumerator;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="Enumerator"/> struct.
+        /// </summary>
+        /// <param name="owner">The owner <see cref="ConditionalWeakTable2{TKey, TValue}"/> instance for the enumerator.</param>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public Enumerator(ConditionalWeakTable2<TKey, TValue> owner)
+        {
+            this.enumerator = ((IEnumerable<KeyValuePair<TKey, TValue>>)owner.table).GetEnumerator();
+        }
+
+        /// <inheritdoc cref="IDisposable.Dispose"/>
+        public void Dispose()
+        {
+            this.enumerator.Dispose();
+        }
+
+        /// <inheritdoc cref="Collections.IEnumerator.MoveNext"/>
+        public bool MoveNext()
+        {
+            return this.enumerator.MoveNext();
+        }
+
+        /// <summary>
+        /// Gets the current key.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TKey GetKey()
+        {
+            return this.enumerator.Current.Key;
+        }
+
+        /// <summary>
+        /// Gets the current value.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TValue GetValue()
+        {
+            return this.enumerator.Current.Value;
+        }
+    }
+}
+
+#endif

CommunityToolkit.Mvvm/Messaging/Internals/System/Runtime.CompilerServices/ConditionalWeakTable2{TKey,TValue}.ZeroAlloc.cs

@@ -157,9 +157,14 @@ public ref struct Enumerator
         private int currentIndex;
 
         /// <summary>
-        /// The current entry set by MoveNext and returned from <see cref="Current"/>.
+        /// The current key, if available.
         /// </summary>
-        private KeyValuePair<TKey, TValue> current;
+        private TKey? key;
+
+        /// <summary>
+        /// The current value, if available.
+        /// </summary>
+        private TValue? value;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="Enumerator"/> class.
@@ -170,22 +175,22 @@ public Enumerator(ConditionalWeakTable2<TKey, TValue> table)
             // Store a reference to the parent table and increase its active enumerator count
             this.table = table;
 
-            Container c = table.container;
+            Container container = table.container;
 
-            if (c is null || c.FirstFreeEntry == 0)
+            if (container is null || container.FirstFreeEntry == 0)
             {
                 // The max index is the same as the current to prevent enumeration
                 this.maxIndexInclusive = -1;
-                this.currentIndex = -1;
-                this.current = default;
             }
             else
             {
                 // Store the max index to be enumerated
-                this.maxIndexInclusive = table.container.FirstFreeEntry - 1;
-                this.currentIndex = -1;
-                this.current = default;
+                this.maxIndexInclusive = container.FirstFreeEntry - 1;
             }
+
+            this.currentIndex = -1;
+            this.key = null;
+            this.value = null;
         }
 
         /// <inheritdoc cref="IDisposable.Dispose"/>
@@ -197,7 +202,8 @@ public void Dispose()
             this.table = null!;
 
             // Ensure we don't keep the last current alive unnecessarily
-            this.current = default;
+            this.key = null;
+            this.value = null;
         }
 
         /// <inheritdoc cref="IEnumerator.MoveNext"/>
@@ -209,30 +215,44 @@ public bool MoveNext()
             // container at the time) has already been finalized, this will be null.
             Container c = this.table.container;
 
-            if (c != null)
+            int currentIndex = this.currentIndex;
+            int maxIndexInclusive = this.maxIndexInclusive;
+
+            // We have the container. Find the next entry to return, if there is one. We need to loop as we
+            // may try to get an entry that's already been removed or collected, in which case we try again.
+            while (currentIndex < maxIndexInclusive)
             {
-                // We have the container. Find the next entry to return, if there is one. We need to loop as we
-                // may try to get an entry that's already been removed or collected, in which case we try again.
-                while (this.currentIndex < this.maxIndexInclusive)
-                {
-                    this.currentIndex++;
+                currentIndex++;
 
-                    if (c.TryGetEntry(this.currentIndex, out TKey? key, out TValue? value))
-                    {
-                        this.current = new KeyValuePair<TKey, TValue>(key, value);
+                if (c.TryGetEntry(currentIndex, out this.key, out this.value))
+                {
+                    this.currentIndex = currentIndex;
 
-                        return true;
-                    }
+                    return true;
                 }
             }
 
+            this.currentIndex = currentIndex;
+
             return false;
         }
 
-        public KeyValuePair<TKey, TValue> Current
+        /// <summary>
+        /// Gets the current key.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TKey GetKey()
+        {
+            return this.key!;
+        }
+
+        /// <summary>
+        /// Gets the current value.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TValue GetValue()
         {
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => this.current;
+            return this.value!;
         }
     }
 

CommunityToolkit.Mvvm/Messaging/Internals/System/Runtime.CompilerServices/ConditionalWeakTable2{TKey,TValue}.cs

@@ -37,11 +37,7 @@ public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
         return this.table.TryGetValue(key, out value);
     }
 
-    /// <summary>
-    /// Tries to add a new pair to the table.
-    /// </summary>
-    /// <param name="key">The key to add.</param>
-    /// <param name="value">The value to associate with key.</param>
+    /// <inheritdoc cref="ConditionalWeakTableExtensions.TryAdd{TKey, TValue}(ConditionalWeakTable{TKey, TValue}, TKey, TValue)"/>
     public bool TryAdd(TKey key, TValue value)
     {
         if (!this.table.TryAdd(key, value))
@@ -116,9 +112,14 @@ public ref struct Enumerator
         private LinkedListNode<WeakReference<TKey>>? node;
 
         /// <summary>
-        /// The current <see cref="KeyValuePair{TKey, TValue}"/> to return.
+        /// The current key, if available.
+        /// </summary>
+        private TKey? key;
+
+        /// <summary>
+        /// The current value, if available.
         /// </summary>
-        private KeyValuePair<TKey, TValue> current;
+        private TValue? value;
 
         /// <summary>
         /// Indicates whether or not <see cref="MoveNext"/> has been called at least once.
@@ -134,10 +135,16 @@ public Enumerator(ConditionalWeakTable2<TKey, TValue> owner)
         {
             this.owner = owner;
             this.node = null;
-            this.current = default;
+            this.key = null;
+            this.value = null;
             this.isFirstMoveNextPending = true;
         }
 
+        /// <inheritdoc cref="IDisposable.Dispose"/>
+        public void Dispose()
+        {
+        }
+
         /// <inheritdoc cref="Collections.IEnumerator.MoveNext"/>
         public bool MoveNext()
         {
@@ -163,7 +170,8 @@ public bool MoveNext()
                     this.owner.table.TryGetValue(target!, out TValue? value))
                 {
                     this.node = node;
-                    this.current = new KeyValuePair<TKey, TValue>(target, value);
+                    this.key = target;
+                    this.value = value;
 
                     return true;
                 }
@@ -179,11 +187,22 @@ public bool MoveNext()
             return false;
         }
 
-        /// <inheritdoc cref="Collections.IEnumerator.MoveNext"/>
-        public readonly KeyValuePair<TKey, TValue> Current
+        /// <summary>
+        /// Gets the current key.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TKey GetKey()
+        {
+            return this.key!;
+        }
+
+        /// <summary>
+        /// Gets the current value.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public TValue GetValue()
         {
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => this.current;
+            return this.value!;
         }
     }
 }

CommunityToolkit.Mvvm/Messaging/Internals/System/Runtime.CompilerServices/ConditionalWeakTableExtensions.cs

@@ -23,7 +23,7 @@ public static bool TryAdd<TKey, TValue>(this ConditionalWeakTable<TKey, TValue>
         where TKey : class
         where TValue : class?
     {
-        // There is no way to do this on .NET Standard 2.0 without exception handling
+        // There is no way to do this on .NET Standard 2.0 or 2.1 without exception handling
         try
         {
             table.Add(key, value);

CommunityToolkit.Mvvm/Messaging/StrongReferenceMessenger.cs

@@ -392,15 +392,15 @@ public TMessage Send<TMessage, TToken>(TMessage message, TToken token)
             while (mappingEnumerator.MoveNext())
             {
                 // Pick the target handler, if the token is a match for the recipient
-                if (mappingEnumerator.Value.TryGetValue(token, out object? handler))
+                if (mappingEnumerator.GetValue().TryGetValue(token, out object? handler))
                 {
                     // This span access should always guaranteed to be valid due to the size of the
                     // array being set according to the current total number of registered handlers,
                     // which will always be greater or equal than the ones matching the previous test.
                     // We're still using a checked span accesses here though to make sure an out of
                     // bounds write can never happen even if an error was present in the logic above.
                     pairs[2 * i] = handler;
-                    pairs[(2 * i) + 1] = mappingEnumerator.Key.Target;
+                    pairs[(2 * i) + 1] = mappingEnumerator.GetKey().Target;
                     i++;
                 }
             }