Support IAlternateEqualityComparer with FrozenDictionary/Set (#103191)

* Support IAlternateEqualityComparer with FrozenDictionary/Set

* Address PR feedback

src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.csproj

@@ -15,6 +15,10 @@
     <Compile Include="$(CoreLibSharedDir)System\Runtime\CompilerServices\CollectionBuilderAttribute.cs" Link="System\Runtime\CompilerServices\CollectionBuilderAttribute.cs" />
   </ItemGroup>
 
+  <ItemGroup Condition="$([MSBuild]::IsTargetFrameworkCompatible('$(TargetFramework)', 'net9.0'))">
+    <Compile Include="System.Collections.Immutable.net9.cs" />
+  </ItemGroup>
+
   <ItemGroup Condition="'$(TargetFramework)' == '$(NetCoreAppCurrent)'">
     <ProjectReference Include="$(LibrariesProjectRoot)System.Collections\ref\System.Collections.csproj" />
     <ProjectReference Include="$(LibrariesProjectRoot)System.Runtime\ref\System.Runtime.csproj" />

src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.net9.cs

@@ -0,0 +1,36 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// ------------------------------------------------------------------------------
+// Changes to this file must follow the https://aka.ms/api-review process.
+// ------------------------------------------------------------------------------
+
+namespace System.Collections.Frozen
+{
+    public abstract partial class FrozenDictionary<TKey, TValue>
+    {
+        public System.Collections.Frozen.FrozenDictionary<TKey,TValue>.AlternateLookup<TAlternateKey> GetAlternateLookup<TAlternateKey>() where TAlternateKey : notnull, allows ref struct { throw null; }
+        public bool TryGetAlternateLookup<TAlternateKey>(out System.Collections.Frozen.FrozenDictionary<TKey, TValue>.AlternateLookup<TAlternateKey> lookup) where TAlternateKey : notnull, allows ref struct { throw null; }
+        public readonly partial struct AlternateLookup<TAlternateKey> where TAlternateKey : notnull, allows ref struct
+        {
+            private readonly object _dummy;
+            private readonly int _dummyPrimitive;
+            public System.Collections.Frozen.FrozenDictionary<TKey, TValue> Dictionary { get { throw null; } }
+            public TValue this[TAlternateKey key] { get { throw null; } }
+            public bool ContainsKey(TAlternateKey key) { throw null; }
+            public bool TryGetValue(TAlternateKey key, [System.Diagnostics.CodeAnalysis.MaybeNullWhenAttribute(false)] out TValue value) { throw null; }
+        }
+    }
+    public abstract partial class FrozenSet<T>
+    {
+        public System.Collections.Frozen.FrozenSet<T>.AlternateLookup<TAlternate> GetAlternateLookup<TAlternate>() { throw null; }
+        public bool TryGetAlternateLookup<TAlternate>(out System.Collections.Frozen.FrozenSet<T>.AlternateLookup<TAlternate> lookup) { throw null; }
+        public readonly partial struct AlternateLookup<TAlternate>
+        {
+            private readonly object _dummy;
+            private readonly int _dummyPrimitive;
+            public System.Collections.Frozen.FrozenSet<T> Set { get { throw null; } }
+            public bool Contains(TAlternate item) { throw null; }
+            public bool TryGetValue(TAlternate equalValue, [System.Diagnostics.CodeAnalysis.MaybeNullWhenAttribute(false)] out T actualValue) { throw null; }
+        }
+    }
+}

src/libraries/System.Collections.Immutable/src/Resources/Strings.resx

@@ -105,4 +105,7 @@
   <data name="ArgumentOutOfRange_NeedNonNegNum" xml:space="preserve">
     <value>Non-negative number required.</value>
   </data>
+  <data name="InvalidOperation_IncompatibleComparer" xml:space="preserve">
+    <value>The collection, in conjunction with its comparer, does not support the specified alternate key type.</value>
+  </data>
 </root>

src/libraries/System.Collections.Immutable/src/System.Collections.Immutable.csproj

@@ -146,7 +146,21 @@ The System.Collections.Immutable library is built-in as part of the shared frame
     <Compile Include="System\Runtime.InteropServices\ImmutableCollectionsMarshal.cs" />
     <Compile Include="Validation\Requires.cs" />
     <Compile Include="$(CommonPath)System\Runtime\Versioning\NonVersionableAttribute.cs" Link="Common\System\Runtime\Versioning\NonVersionableAttribute.cs" />
-    <None Include="Interfaces.cd" />
+  </ItemGroup>
+
+  <ItemGroup Condition="$([MSBuild]::IsTargetFrameworkCompatible('$(TargetFramework)', 'net9.0'))">
+    <Compile Include="System\Collections\Frozen\DefaultFrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\DefaultFrozenSet.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\FrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\FrozenSet.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\SmallFrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\SmallFrozenSet.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\Int32\Int32FrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\Int32\Int32FrozenSet.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\String\LengthBucketsFrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\String\LengthBucketsFrozenSet.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\String\OrdinalStringFrozenDictionary.AlternateLookup.cs" />
+    <Compile Include="System\Collections\Frozen\String\OrdinalStringFrozenSet.AlternateLookup.cs" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFrameworkIdentifier)' != '.NETCoreApp'">

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

@@ -0,0 +1,33 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
+
+namespace System.Collections.Frozen
+{
+    internal sealed partial class DefaultFrozenDictionary<TKey, TValue>
+    {
+        /// <inheritdoc/>
+        private protected override ref readonly TValue GetValueRefOrNullRefCore<TAlternateKey>(TAlternateKey key)
+        {
+            IAlternateEqualityComparer<TAlternateKey, TKey> comparer = GetAlternateEqualityComparer<TAlternateKey>();
+
+            int hashCode = comparer.GetHashCode(key);
+            _hashTable.FindMatchingEntries(hashCode, out int index, out int endIndex);
+
+            while (index <= endIndex)
+            {
+                if (hashCode == _hashTable.HashCodes[index] && comparer.Equals(key, _keys[index]))
+                {
+                    return ref _values[index];
+                }
+
+                index++;
+            }
+
+            return ref Unsafe.NullRef<TValue>();
+        }
+    }
+}

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

@@ -9,7 +9,7 @@ namespace System.Collections.Frozen
     /// <summary>Provides the default <see cref="FrozenDictionary{TKey, TValue}"/> implementation to use when no other special-cases apply.</summary>
     /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
     /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
-    internal sealed class DefaultFrozenDictionary<TKey, TValue> : KeysAndValuesFrozenDictionary<TKey, TValue>, IDictionary<TKey, TValue>
+    internal sealed partial class DefaultFrozenDictionary<TKey, TValue> : KeysAndValuesFrozenDictionary<TKey, TValue>, IDictionary<TKey, TValue>
         where TKey : notnull
     {
         internal DefaultFrozenDictionary(Dictionary<TKey, TValue> source)

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

@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+
+namespace System.Collections.Frozen
+{
+    internal sealed partial class DefaultFrozenSet<T>
+    {
+        /// <inheritdoc />
+        private protected override int FindItemIndex<TAlternate>(TAlternate item)
+        {
+            IAlternateEqualityComparer<TAlternate, T> comparer = GetAlternateEqualityComparer<TAlternate>();
+
+            int hashCode = item is null ? 0 : comparer.GetHashCode(item);
+            _hashTable.FindMatchingEntries(hashCode, out int index, out int endIndex);
+
+            while (index <= endIndex)
+            {
+                if (hashCode == _hashTable.HashCodes[index] && comparer.Equals(item, _items[index]))
+                {
+                    return index;
+                }
+
+                index++;
+            }
+
+            return -1;
+        }
+    }
+}

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

@@ -7,7 +7,7 @@ namespace System.Collections.Frozen
 {
     /// <summary>Provides the default <see cref="FrozenSet{T}"/> implementation to use when no other special-cases apply.</summary>
     /// <typeparam name="T">The type of the values in the set.</typeparam>
-    internal sealed class DefaultFrozenSet<T> : ItemsFrozenSet<T, DefaultFrozenSet<T>.GSW>
+    internal sealed partial class DefaultFrozenSet<T> : ItemsFrozenSet<T, DefaultFrozenSet<T>.GSW>
     {
         internal DefaultFrozenSet(HashSet<T> source)
             : base(source)

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

@@ -0,0 +1,141 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+namespace System.Collections.Frozen
+{
+    public abstract partial class FrozenDictionary<TKey, TValue>
+    {
+        /// <summary>
+        /// Gets an instance of a type that may be used to perform operations on a <see cref="FrozenDictionary{TKey, TValue}"/>
+        /// using a <typeparamref name="TAlternateKey"/> as a key instead of a <typeparamref name="TKey"/>.
+        /// </summary>
+        /// <typeparam name="TAlternateKey">The alternate type of a key for performing lookups.</typeparam>
+        /// <returns>The created lookup instance.</returns>
+        /// <exception cref="InvalidOperationException">This instance's comparer is not compatible with <typeparamref name="TAlternateKey"/>.</exception>
+        /// <remarks>
+        /// This instance must be using a comparer that implements <see cref="IAlternateEqualityComparer{TAlternateKey, TKey}"/> with
+        /// <typeparamref name="TAlternateKey"/> and <typeparamref name="TKey"/>. If it doesn't, an exception will be thrown.
+        /// </remarks>
+        public AlternateLookup<TAlternateKey> GetAlternateLookup<TAlternateKey>() where TAlternateKey : notnull, allows ref struct
+        {
+            if (!TryGetAlternateLookup(out AlternateLookup<TAlternateKey> lookup))
+            {
+                ThrowHelper.ThrowIncompatibleComparer();
+            }
+
+            return lookup;
+        }
+
+        /// <summary>
+        /// Gets an instance of a type that may be used to perform operations on a <see cref="FrozenDictionary{TKey, TValue}"/>
+        /// using a <typeparamref name="TAlternateKey"/> as a key instead of a <typeparamref name="TKey"/>.
+        /// </summary>
+        /// <typeparam name="TAlternateKey">The alternate type of a key for performing lookups.</typeparam>
+        /// <param name="lookup">The created lookup instance when the method returns true, or a default instance that should not be used if the method returns false.</param>
+        /// <returns>true if a lookup could be created; otherwise, false.</returns>
+        /// <remarks>
+        /// This instance must be using a comparer that implements <see cref="IAlternateEqualityComparer{TAlternateKey, TKey}"/> with
+        /// <typeparamref name="TAlternateKey"/> and <typeparamref name="TKey"/>. If it doesn't, the method will return false.
+        /// </remarks>
+        public bool TryGetAlternateLookup<TAlternateKey>(out AlternateLookup<TAlternateKey> lookup) where TAlternateKey : notnull, allows ref struct
+        {
+            // The comparer must support the specified TAlternateKey. If it doesn't we can't create a lookup.
+            // Some implementations where TKey is string rely on the length of the input and use it as part of the storage scheme.
+            // That means we can only support TAlternateKeys that have a length we can check, which means we have to special-case
+            // it. Since which implementation we pick is based on a heuristic and can't be predicted by the consumer, we don't
+            // just have this requirement in that one implementation but for all implementations that might be picked for string.
+            // As such, if the key is a string, we only support ReadOnlySpan<char> as the alternate key.
+            if (Comparer is IAlternateEqualityComparer<TAlternateKey, TKey> &&
+                (typeof(TKey) != typeof(string) || typeof(TAlternateKey) == typeof(ReadOnlySpan<char>)))
+            {
+                lookup = new AlternateLookup<TAlternateKey>(this);
+                return true;
+            }
+
+            lookup = default;
+            return false;
+        }
+
+        /// <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, TKey> GetAlternateEqualityComparer<TAlternateKey>() where TAlternateKey : notnull, allows ref struct
+        {
+            Debug.Assert(Comparer is IAlternateEqualityComparer<TAlternateKey, TKey>, "Must have already been verified");
+            return Unsafe.As<IAlternateEqualityComparer<TAlternateKey, TKey>>(Comparer);
+        }
+
+        /// <summary>
+        /// Provides a type that may be used to perform operations on a <see cref="FrozenDictionary{TKey, TValue}"/>
+        /// using a <typeparamref name="TAlternateKey"/> as a key instead of a <typeparamref name="TKey"/>.
+        /// </summary>
+        /// <typeparam name="TAlternateKey">The alternate type of a key for performing lookups.</typeparam>
+        public readonly struct AlternateLookup<TAlternateKey> where TAlternateKey : notnull, allows ref struct
+        {
+            /// <summary>Initialize the instance. The dictionary must have already been verified to have a compatible comparer.</summary>
+            internal AlternateLookup(FrozenDictionary<TKey, TValue> dictionary)
+            {
+                Debug.Assert(dictionary is not null);
+                Debug.Assert(dictionary.Comparer is IAlternateEqualityComparer<TAlternateKey, TKey>);
+                Dictionary = dictionary;
+            }
+
+            /// <summary>Gets the <see cref="FrozenDictionary{TKey, TValue}"/> against which this instance performs operations.</summary>
+            public FrozenDictionary<TKey, TValue> Dictionary { get; }
+
+            /// <summary>Gets or sets the value associated with the specified alternate key.</summary>
+            /// <param name="key">The alternate key of the value to get or set.</param>
+            /// <value>
+            /// The value associated with the specified alternate key. If the specified alternate key is not found, a get operation throws
+            /// a <see cref="KeyNotFoundException"/>, and a set operation creates a new element with the specified key.
+            /// </value>
+            /// <exception cref="ArgumentNullException"><paramref name="key"/> is <see langword="null"/>.</exception>
+            /// <exception cref="KeyNotFoundException">The alternate key does not exist in the collection.</exception>
+            public TValue this[TAlternateKey key]
+            {
+                get
+                {
+                    ref readonly TValue valueRef = ref Dictionary.GetValueRefOrNullRefCore(key);
+                    if (Unsafe.IsNullRef(in valueRef))
+                    {
+                        ThrowHelper.ThrowKeyNotFoundException();
+                    }
+
+                    return valueRef;
+                }
+            }
+
+            /// <summary>Determines whether the <see cref="FrozenDictionary{TKey, TValue}"/> contains the specified alternate key.</summary>
+            /// <param name="key">The alternate key to check.</param>
+            /// <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));
+
+            /// <summary>Gets the value associated with the specified alternate key.</summary>
+            /// <param name="key">The alternate key of the value to get.</param>
+            /// <param name="value">
+            /// When this method returns, contains the value associated with the specified key, if the key is found;
+            /// otherwise, the default value for the type of the value parameter.
+            /// </param>
+            /// <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);
+
+                if (!Unsafe.IsNullRef(in valueRef))
+                {
+                    value = valueRef;
+                    return true;
+                }
+
+                value = default;
+                return false;
+            }
+        }
+    }
+}