Optimize ILookup<TKey, TElement> implementations

Author: Sergio0694

CommunityToolkit.Mvvm/Collections/ObservableGroupedCollectionExtensions.cs

@@ -17,7 +17,7 @@ namespace CommunityToolkit.Mvvm.Collections;
 public static class ObservableGroupedCollectionExtensions
 {
     /// <summary>
-    /// Return the first group with <paramref name="key"/> key.
+    /// Returns the first group with <paramref name="key"/> key.
     /// </summary>
     /// <typeparam name="TKey">The type of the group key.</typeparam>
     /// <typeparam name="TElement">The type of the items in the collection.</typeparam>
@@ -49,13 +49,13 @@ static void ThrowArgumentExceptionForKeyNotFound()
     }
 
     /// <summary>
-    /// Return the first group with <paramref name="key"/> key or null if not found.
+    /// Returns the first group with <paramref name="key"/> key or <see langword="null"/> if not found.
     /// </summary>
     /// <typeparam name="TKey">The type of the group key.</typeparam>
     /// <typeparam name="TElement">The type of the items in the collection.</typeparam>
     /// <param name="source">The source <see cref="ObservableGroupedCollection{TKey, TElement}"/> instance.</param>
     /// <param name="key">The key of the group to query.</param>
-    /// <returns>The first group matching <paramref name="key"/> or null.</returns>
+    /// <returns>The first group matching <paramref name="key"/> or <see langword="null"/>.</returns>
     /// <exception cref="ArgumentNullException">Thrown if <paramref name="source"/> or <paramref name="key"/> are <see langword="null"/>.</exception>
     public static ObservableGroup<TKey, TElement>? FirstGroupByKeyOrDefault<TKey, TElement>(this ObservableGroupedCollection<TKey, TElement> source, TKey key)
         where TKey : notnull
@@ -65,7 +65,7 @@ static void ThrowArgumentExceptionForKeyNotFound()
 
         if (source.TryGetList(out List<ObservableGroup<TKey, TElement>>? list))
         {
-            foreach (ObservableGroup<TKey, TElement>? group in list)
+            foreach (ObservableGroup<TKey, TElement> group in list)
             {
                 if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
                 {
@@ -77,12 +77,12 @@ static void ThrowArgumentExceptionForKeyNotFound()
         }
 
         [MethodImpl(MethodImplOptions.NoInlining)]
-        static ObservableGroup<TKey, TElement>? FirstOrDefaultFallback(ObservableGroupedCollection<TKey, TElement> source, TKey key)
+        static ObservableGroup<TKey, TElement>? FirstGroupByKeyOrDefaultFallback(ObservableGroupedCollection<TKey, TElement> source, TKey key)
         {
             return Enumerable.FirstOrDefault<ObservableGroup<TKey, TElement>>(source, group => EqualityComparer<TKey>.Default.Equals(group.Key, key));
         }
 
-        return FirstOrDefaultFallback(source, key);
+        return FirstGroupByKeyOrDefaultFallback(source, key);
     }
 
     /// <summary>
@@ -720,7 +720,7 @@ public static void RemoveGroup<TKey, TValue>(this ObservableGroupedCollection<TK
         {
             int index = 0;
 
-            foreach (ObservableGroup<TKey, TValue>? group in list)
+            foreach (ObservableGroup<TKey, TValue> group in list)
             {
                 if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
                 {
@@ -739,7 +739,7 @@ static void RemoveGroupFallback(ObservableGroupedCollection<TKey, TValue> source
             {
                 int index = 0;
 
-                foreach (ObservableGroup<TKey, TValue>? group in source)
+                foreach (ObservableGroup<TKey, TValue> group in source)
                 {
                     if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
                     {
@@ -776,7 +776,7 @@ public static void RemoveItem<TKey, TValue>(this ObservableGroupedCollection<TKe
         {
             int index = 0;
 
-            foreach (ObservableGroup<TKey, TValue>? group in list)
+            foreach (ObservableGroup<TKey, TValue> group in list)
             {
                 if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
                 {
@@ -800,7 +800,7 @@ static void RemoveItemFallback(ObservableGroupedCollection<TKey, TValue> source,
             {
                 int index = 0;
 
-                foreach (ObservableGroup<TKey, TValue>? group in source)
+                foreach (ObservableGroup<TKey, TValue> group in source)
                 {
                     if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
                     {

CommunityToolkit.Mvvm/Collections/ObservableGroupedCollection{TKey,TElement}.cs

@@ -40,8 +40,14 @@ IEnumerable<TElement> ILookup<TKey, TElement>.this[TKey key]
     {
         get
         {
-            // TODO: optimize this
-            return Enumerable.FirstOrDefault<ObservableGroup<TKey, TElement>>(this, item => EqualityComparer<TKey>.Default.Equals(item.Key, key)) ?? Enumerable.Empty<TElement>();
+            IEnumerable<TElement>? result = null;
+
+            if (key is not null)
+            {
+                result = this.FirstGroupByKeyOrDefault(key);
+            }
+
+            return result ?? Enumerable.Empty<TElement>();
         }
     }
 
@@ -61,8 +67,7 @@ internal bool TryGetList([NotNullWhen(true)] out List<ObservableGroup<TKey, TEle
     /// <inheritdoc/>
     bool ILookup<TKey, TElement>.Contains(TKey key)
     {
-        // TODO: optimize this
-        return Enumerable.Any<ObservableGroup<TKey, TElement>>(this, item => EqualityComparer<TKey>.Default.Equals(item.Key, key));
+        return key is not null && this.FirstGroupByKey(key) is not null;
     }
 
     /// <inheritdoc/>

CommunityToolkit.Mvvm/Collections/ReadOnlyObservableGroupedCollection{TKey,TElement}.cs

@@ -8,6 +8,7 @@
 using System.Collections.Specialized;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Runtime.CompilerServices;
 
 namespace CommunityToolkit.Mvvm.Collections;
 
@@ -45,16 +46,21 @@ IEnumerable<TElement> ILookup<TKey, TElement>.this[TKey key]
     {
         get
         {
-            // TODO: optimize this
-            return Enumerable.FirstOrDefault<ReadOnlyObservableGroup<TKey, TElement>>(this, item => EqualityComparer<TKey>.Default.Equals(item.Key, key)) ?? Enumerable.Empty<TElement>();
+            IEnumerable<TElement>? result = null;
+
+            if (key is not null)
+            {
+                result = FirstGroupByKeyOrDefault(key);
+            }
+
+            return result ?? Enumerable.Empty<TElement>();
         }
     }
 
     /// <inheritdoc/>
     bool ILookup<TKey, TElement>.Contains(TKey key)
     {
-        // TODO: optimize this
-        return Enumerable.Any<ReadOnlyObservableGroup<TKey, TElement>>(this, item => EqualityComparer<TKey>.Default.Equals(item.Key, key));
+        return key is not null && FirstGroupByKeyOrDefault(key) is not null;
     }
 
     /// <inheritdoc/>
@@ -153,4 +159,33 @@ static void ThrowNotSupportedExceptionForRangeOperation()
                 break;
         }
     }
+
+    /// <summary>
+    /// Returns the first group with <paramref name="key"/> key or <see langword="null"/> if not found.
+    /// </summary>
+    /// <param name="key">The key of the group to query (assumed not to be <see langword="null"/>).</param>
+    /// <returns>The first group matching <paramref name="key"/>.</returns>
+    private IEnumerable<TElement>? FirstGroupByKeyOrDefault(TKey key)
+    {
+        if (Items is List<ReadOnlyObservableGroup<TKey, TElement>> list)
+        {
+            foreach (ReadOnlyObservableGroup<TKey, TElement> group in list)
+            {
+                if (EqualityComparer<TKey>.Default.Equals(group.Key, key))
+                {
+                    return group;
+                }
+            }
+
+            return null;
+        }
+
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static IEnumerable<TElement>? FirstGroupByKeyOrDefaultFallback(ReadOnlyObservableGroupedCollection<TKey, TElement> source, TKey key)
+        {
+            return Enumerable.FirstOrDefault<ReadOnlyObservableGroup<TKey, TElement>>(source, group => EqualityComparer<TKey>.Default.Equals(group.Key, key));
+        }
+
+        return FirstGroupByKeyOrDefaultFallback(this, key);
+    }
 }