Avoid generic virtual dispatch for frozen collections alternate lookup (#108732)

* Avoid generic virtual dispatch for frozen collections alternate lookup

* Various stylistic changes addressing PR feedback

* Tidy up some more xmldocs

---------

Co-authored-by: Stephen Toub <[email protected]>

Author: andrewjsaid

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/DefaultFrozenDictionary.AlternateLookup.cs

@@ -10,7 +10,33 @@ namespace System.Collections.Frozen
     internal sealed partial class DefaultFrozenDictionary<TKey, TValue>
     {
         /// <inheritdoc/>
-        private protected override ref readonly TValue GetValueRefOrNullRefCore<TAlternateKey>(TAlternateKey key)
+        private protected override AlternateLookupDelegate<TAlternateKey> GetAlternateLookupDelegate<TAlternateKey>()
+            => AlternateLookupDelegateHolder<TAlternateKey>.Instance;
+
+        private static class AlternateLookupDelegateHolder<TAlternateKey>
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
+        {
+            /// <summary>
+            /// Invokes <see cref="GetValueRefOrNullRefCoreAlternate{TAlternateKey}(TAlternateKey)"/>
+            /// on instances known to be of type <see cref="DefaultFrozenDictionary{TKey, TValue}"/>.
+            /// </summary>
+            public static readonly AlternateLookupDelegate<TAlternateKey> Instance = (dictionary, key)
+                => ref ((DefaultFrozenDictionary<TKey, TValue>)dictionary).GetValueRefOrNullRefCoreAlternate(key);
+        }
+
+        /// <inheritdoc cref="GetValueRefOrNullRefCore(TKey)" />
+        private ref readonly TValue GetValueRefOrNullRefCoreAlternate<TAlternateKey>(TAlternateKey key)
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
         {
             IAlternateEqualityComparer<TAlternateKey, TKey> comparer = GetAlternateEqualityComparer<TAlternateKey>();
 

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/DefaultFrozenSet.AlternateLookup.cs

@@ -7,8 +7,32 @@ namespace System.Collections.Frozen
 {
     internal sealed partial class DefaultFrozenSet<T>
     {
-        /// <inheritdoc />
-        private protected override int FindItemIndex<TAlternate>(TAlternate item)
+        /// <inheritdoc/>
+        private protected override AlternateLookupDelegate<TAlternate> GetAlternateLookupDelegate<TAlternate>()
+            => AlternateLookupDelegateHolder<TAlternate>.Instance;
+
+        private static class AlternateLookupDelegateHolder<TAlternate>
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            where TAlternate : allows ref struct
+#pragma warning restore SA1001
+#endif
+        {
+            /// <summary>
+            /// Invokes <see cref="FindItemIndexAlternate{TAlternate}(TAlternate)"/>
+            /// on instances known to be of type <see cref="DefaultFrozenSet{TValue}"/>.
+            /// </summary>
+            public static readonly AlternateLookupDelegate<TAlternate> Instance = (set, item)
+                => ((DefaultFrozenSet<T>)set).FindItemIndexAlternate(item);
+        }
+
+        /// <inheritdoc cref="FindItemIndex(T)" />
+        private int FindItemIndexAlternate<TAlternate>(TAlternate item)
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            where TAlternate : allows ref struct
+#pragma warning restore SA1001
+#endif
         {
             IAlternateEqualityComparer<TAlternate, T> comparer = GetAlternateEqualityComparer<TAlternate>();
 

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/FrozenDictionary.AlternateLookup.cs

@@ -53,7 +53,7 @@ public bool TryGetAlternateLookup<TAlternateKey>(out AlternateLookup<TAlternateK
             if (Comparer is IAlternateEqualityComparer<TAlternateKey, TKey> &&
                 (typeof(TKey) != typeof(string) || typeof(TAlternateKey) == typeof(ReadOnlySpan<char>)))
             {
-                lookup = new AlternateLookup<TAlternateKey>(this);
+                lookup = new AlternateLookup<TAlternateKey>(this, GetAlternateLookupDelegate<TAlternateKey>());
                 return true;
             }
 
@@ -76,12 +76,16 @@ private protected IAlternateEqualityComparer<TAlternateKey, TKey> GetAlternateEq
         /// <typeparam name="TAlternateKey">The alternate type of a key for performing lookups.</typeparam>
         public readonly struct AlternateLookup<TAlternateKey> where TAlternateKey : notnull, allows ref struct
         {
+            private readonly AlternateLookupDelegate<TAlternateKey> _alternateLookupDelegate;
+
             /// <summary>Initialize the instance. The dictionary must have already been verified to have a compatible comparer.</summary>
-            internal AlternateLookup(FrozenDictionary<TKey, TValue> dictionary)
+            internal AlternateLookup(FrozenDictionary<TKey, TValue> dictionary, AlternateLookupDelegate<TAlternateKey> alternateLookupDelegate)
             {
                 Debug.Assert(dictionary is not null);
                 Debug.Assert(dictionary.Comparer is IAlternateEqualityComparer<TAlternateKey, TKey>);
+                Debug.Assert(alternateLookupDelegate is not null);
                 Dictionary = dictionary;
+                _alternateLookupDelegate = alternateLookupDelegate;
             }
 
             /// <summary>Gets the <see cref="FrozenDictionary{TKey, TValue}"/> against which this instance performs operations.</summary>
@@ -99,7 +103,7 @@ public TValue this[TAlternateKey key]
             {
                 get
                 {
-                    ref readonly TValue valueRef = ref Dictionary.GetValueRefOrNullRefCore(key);
+                    ref readonly TValue valueRef = ref _alternateLookupDelegate(Dictionary, key);
                     if (Unsafe.IsNullRef(in valueRef))
                     {
                         ThrowHelper.ThrowKeyNotFoundException();
@@ -114,7 +118,7 @@ public TValue this[TAlternateKey key]
             /// <returns><see langword="true"/> if the key is in the dictionary; otherwise, <see langword="false"/>.</returns>
             /// <exception cref="ArgumentNullException"><paramref name="key"/> is <see langword="null"/>.</exception>
             public bool ContainsKey(TAlternateKey key) =>
-                !Unsafe.IsNullRef(in Dictionary.GetValueRefOrNullRefCore(key));
+                !Unsafe.IsNullRef(in _alternateLookupDelegate(Dictionary, key));
 
             /// <summary>Gets the value associated with the specified alternate key.</summary>
             /// <param name="key">The alternate key of the value to get.</param>
@@ -126,7 +130,7 @@ public bool ContainsKey(TAlternateKey key) =>
             /// <exception cref="ArgumentNullException"><paramref name="key"/> is <see langword="null"/>.</exception>
             public bool TryGetValue(TAlternateKey key, [MaybeNullWhen(false)] out TValue value)
             {
-                ref readonly TValue valueRef = ref Dictionary.GetValueRefOrNullRefCore(key);
+                ref readonly TValue valueRef = ref _alternateLookupDelegate(Dictionary, key);
 
                 if (!Unsafe.IsNullRef(in valueRef))
                 {

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/FrozenDictionary.cs

@@ -449,20 +449,20 @@ public ref readonly TValue GetValueRefOrNullRef(TKey key)
         /// <inheritdoc cref="GetValueRefOrNullRef" />
         private protected abstract ref readonly TValue GetValueRefOrNullRefCore(TKey key);
 
-        /// <inheritdoc cref="GetValueRefOrNullRef" />
+        /// <summary>
+        /// Retrieves a delegate which calls a method equivalent to <see cref="GetValueRefOrNullRef(TKey)"/>
+        /// for the <typeparamref name="TAlternateKey"/>.
+        /// </summary>
         /// <remarks>
         /// This is virtual rather than abstract because only some implementations need to support this, e.g. implementations that
         /// are only ever used with the default comparer won't ever hit code paths that use this, at least not
         /// until/if we make `EqualityComparer{string}.Default` implement `IAlternateEqualityComparer{ReadOnlySpan{char}, string}`.
         ///
-        /// This unfortunately needs to be a generic virtual method, but the only other known option involves having a dedicated
-        /// class instance such that the generic can be baked into that, where the methods on it are still virtual but don't have
-        /// extra generic methods. But for most implementations, either a) that class would need to be allocated as part of
-        /// TryGetAlternateLookup, which would be more expensive for use cases where someone needs a lookup for just a few operations,
-        /// or b) a dictionary of those instances would need to be maintained, which just replaces the runtime's dictionary for a GVM
-        /// with a custom one here.
+        /// Generic Virtual method invocation is slower than delegate invocation and could negate
+        /// much of the benefit of using Alternate Keys. By retrieving the delegate up-front when
+        /// the lookup is created, we only pay for generic virtual method invocation once.
         /// </remarks>
-        private protected virtual ref readonly TValue GetValueRefOrNullRefCore<TAlternateKey>(TAlternateKey key)
+        private protected virtual AlternateLookupDelegate<TAlternateKey> GetAlternateLookupDelegate<TAlternateKey>()
             where TAlternateKey : notnull
 #if NET9_0_OR_GREATER
 #pragma warning disable SA1001 // Commas should be spaced correctly
@@ -473,7 +473,34 @@ private protected virtual ref readonly TValue GetValueRefOrNullRefCore<TAlternat
             , allows ref struct
 #pragma warning restore SA1001
 #endif
-            => ref Unsafe.NullRef<TValue>();
+            => AlternateLookupDelegateHolder<TAlternateKey>.ReturnsNullRef;
+
+        /// <summary>
+        /// Invokes a method equivalent to <see cref="GetValueRefOrNullRef(TKey)"/>
+        /// for the <typeparamref name="TAlternateKey"/>.
+        /// </summary>
+        internal delegate ref readonly TValue AlternateLookupDelegate<TAlternateKey>(FrozenDictionary<TKey, TValue> dictionary, TAlternateKey key)
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
+            ;
+
+        /// <summary>
+        /// Holds an implementation of <see cref="AlternateLookupDelegate{TAlternateKey}"/> which always returns a null ref.
+        /// </summary>
+        private static class AlternateLookupDelegateHolder<TAlternateKey>
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
+        {
+            public static readonly AlternateLookupDelegate<TAlternateKey> ReturnsNullRef = (_, _) => ref Unsafe.NullRef<TValue>();
+        }
 
         /// <summary>Gets a reference to the value associated with the specified key.</summary>
         /// <param name="key">The key of the value to get.</param>

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/FrozenSet.AlternateLookup.cs

@@ -53,7 +53,7 @@ public bool TryGetAlternateLookup<TAlternate>(out AlternateLookup<TAlternate> lo
             if (Comparer is IAlternateEqualityComparer<TAlternate, T> &&
                 (typeof(T) != typeof(string) || typeof(TAlternate) == typeof(ReadOnlySpan<char>)))
             {
-                lookup = new AlternateLookup<TAlternate>(this);
+                lookup = new AlternateLookup<TAlternate>(this, GetAlternateLookupDelegate<TAlternate>());
                 return true;
             }
 
@@ -63,10 +63,10 @@ public bool TryGetAlternateLookup<TAlternate>(out AlternateLookup<TAlternate> lo
 
         /// <summary>Gets the <see cref="Comparer"/> as an <see cref="IAlternateEqualityComparer{TAlternate, T}"/>.</summary>
         /// <remarks>This must only be used when it's already been proven that the comparer implements the target interface.</remarks>
-        private protected IAlternateEqualityComparer<TAlternateKey, T> GetAlternateEqualityComparer<TAlternateKey>() where TAlternateKey : allows ref struct
+        private protected IAlternateEqualityComparer<TAlternate, T> GetAlternateEqualityComparer<TAlternate>() where TAlternate : allows ref struct
         {
-            Debug.Assert(Comparer is IAlternateEqualityComparer<TAlternateKey, T>, "Must have already been verified");
-            return Unsafe.As<IAlternateEqualityComparer<TAlternateKey, T>>(Comparer);
+            Debug.Assert(Comparer is IAlternateEqualityComparer<TAlternate, T>, "Must have already been verified");
+            return Unsafe.As<IAlternateEqualityComparer<TAlternate, T>>(Comparer);
         }
 
         /// <summary>
@@ -76,12 +76,16 @@ private protected IAlternateEqualityComparer<TAlternateKey, T> GetAlternateEqual
         /// <typeparam name="TAlternate">The alternate type of a key for performing lookups.</typeparam>
         public readonly struct AlternateLookup<TAlternate> where TAlternate : allows ref struct
         {
+            private readonly AlternateLookupDelegate<TAlternate> _alternateLookupDelegate;
+
             /// <summary>Initialize the instance. The set must have already been verified to have a compatible comparer.</summary>
-            internal AlternateLookup(FrozenSet<T> set)
+            internal AlternateLookup(FrozenSet<T> set, AlternateLookupDelegate<TAlternate> alternateLookupDelegate)
             {
                 Debug.Assert(set is not null);
                 Debug.Assert(set.Comparer is IAlternateEqualityComparer<TAlternate, T>);
+                Debug.Assert(alternateLookupDelegate is not null);
                 Set = set;
+                _alternateLookupDelegate = alternateLookupDelegate;
             }
 
             /// <summary>Gets the <see cref="FrozenSet{T}"/> against which this instance performs operations.</summary>
@@ -90,15 +94,15 @@ internal AlternateLookup(FrozenSet<T> set)
             /// <summary>Determines whether a set contains the specified element.</summary>
             /// <param name="item">The element to locate in the set.</param>
             /// <returns>true if the set contains the specified element; otherwise, false.</returns>
-            public bool Contains(TAlternate item) => Set.FindItemIndex(item) >= 0;
+            public bool Contains(TAlternate item) => _alternateLookupDelegate(Set, item) >= 0;
 
             /// <summary>Searches the set for a given value and returns the equal value it finds, if any.</summary>
             /// <param name="equalValue">The value to search for.</param>
             /// <param name="actualValue">The value from the set that the search found, or the default value of <typeparamref name="T"/> when the search yielded no match.</param>
             /// <returns>A value indicating whether the search was successful.</returns>
             public bool TryGetValue(TAlternate equalValue, [MaybeNullWhen(false)] out T actualValue)
             {
-                int index = Set.FindItemIndex(equalValue);
+                int index = _alternateLookupDelegate(Set, equalValue);
                 if (index >= 0)
                 {
                     actualValue = Set.Items[index];

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/FrozenSet.cs

@@ -329,18 +329,56 @@ public bool TryGetValue(T equalValue, [MaybeNullWhen(false)] out T actualValue)
         /// <returns>The index of the value, or -1 if not found.</returns>
         private protected abstract int FindItemIndex(T item);
 
-        /// <summary>Finds the index of a specific value in a set.</summary>
-        /// <param name="item">The value to lookup.</param>
-        /// <returns>The index of the value, or -1 if not found.</returns>
-        private protected virtual int FindItemIndex<TAlternate>(TAlternate item)
+        /// <summary>
+        /// Retrieves a delegate which calls a method equivalent to <see cref="FindItemIndex"/>
+        /// for the <typeparamref name="TAlternate"/>.
+        /// </summary>
+        /// <remarks>
+        /// This is virtual rather than abstract because only some implementations need to support this, e.g. implementations that
+        /// are only ever used with the default comparer won't ever hit code paths that use this, at least not
+        /// until/if we make `EqualityComparer{string}.Default` implement `IAlternateEqualityComparer{ReadOnlySpan{char}, string}`.
+        ///
+        /// Generic Virtual method invocation is slower than delegate invocation and could negate
+        /// much of the benefit of using Alternate Keys. By retrieving the delegate up-front when
+        /// the lookup is created, we only pay for generic virtual method invocation once.
+        /// </remarks>
+        private protected virtual AlternateLookupDelegate<TAlternate> GetAlternateLookupDelegate<TAlternate>()
 #if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
             // This method will only ever be used on .NET 9+. However, because of how everything is structured,
             // and to avoid a proliferation of conditional files for many of the derived types (in particular
             // for the OrdinalString* implementations), we still build this method into all builds, even though
             // it'll be unused. But we can't use the allows ref struct constraint downlevel, hence the #if.
             where TAlternate : allows ref struct
+#pragma warning restore SA1001
+#endif
+            => AlternateLookupDelegateHolder<TAlternate>.ReturnsNullRef;
+
+        /// <summary>
+        /// Invokes a method equivalent to <see cref="FindItemIndex"/>
+        /// for the <typeparamref name="TAlternate"/>.
+        /// </summary>
+        internal delegate int AlternateLookupDelegate<TAlternate>(FrozenSet<T> set, TAlternate key)
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            where TAlternate : allows ref struct
+#pragma warning restore SA1001
 #endif
-            => -1;
+            ;
+
+        /// <summary>
+        /// Holds an implementation of <see cref="AlternateLookupDelegate{TAlternate}"/> which always returns -1.
+        /// </summary>
+        private static class AlternateLookupDelegateHolder<TAlternate>
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            where TAlternate : allows ref struct
+#pragma warning restore SA1001
+#endif
+        {
+            public static readonly AlternateLookupDelegate<TAlternate> ReturnsNullRef = (_, _) => -1;
+        }
+
 
         /// <summary>Returns an enumerator that iterates through the set.</summary>
         /// <returns>An enumerator that iterates through the set.</returns>

src/libraries/System.Collections.Immutable/src/System/Collections/Frozen/Int32/Int32FrozenDictionary.AlternateLookup.cs

@@ -11,7 +11,33 @@ namespace System.Collections.Frozen
     internal sealed partial class Int32FrozenDictionary<TValue>
     {
         /// <inheritdoc/>
-        private protected override ref readonly TValue GetValueRefOrNullRefCore<TAlternateKey>(TAlternateKey key)
+        private protected override AlternateLookupDelegate<TAlternateKey> GetAlternateLookupDelegate<TAlternateKey>()
+            => AlternateLookupDelegateHolder<TAlternateKey>.Instance;
+
+        private static class AlternateLookupDelegateHolder<TAlternateKey>
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
+        {
+            /// <summary>
+            /// Invokes <see cref="GetValueRefOrNullRefCoreAlternate{TAlternate}(TAlternate)"/>
+            /// on instances known to be of type <see cref="Int32FrozenDictionary{TValue}"/>.
+            /// </summary>
+            public static readonly AlternateLookupDelegate<TAlternateKey> Instance = (dictionary, key)
+                => ref ((Int32FrozenDictionary<TValue>)dictionary).GetValueRefOrNullRefCoreAlternate(key);
+        }
+
+        /// <inheritdoc cref="GetValueRefOrNullRefCore(int)" />
+        private ref readonly TValue GetValueRefOrNullRefCoreAlternate<TAlternateKey>(TAlternateKey key)
+            where TAlternateKey : notnull
+#if NET9_0_OR_GREATER
+#pragma warning disable SA1001 // Commas should be spaced correctly
+            , allows ref struct
+#pragma warning restore SA1001
+#endif
         {
             IAlternateEqualityComparer<TAlternateKey, int> comparer = GetAlternateEqualityComparer<TAlternateKey>();