Add back IComparable-based optimization to FrozenDictionary/Set

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

@@ -29,6 +29,8 @@ The System.Collections.Immutable library is built-in as part of the shared frame
     <Compile Include="System\Collections\Frozen\KeysAndValuesFrozenDictionary.cs" />
     <Compile Include="System\Collections\Frozen\SmallFrozenDictionary.cs" />
     <Compile Include="System\Collections\Frozen\SmallFrozenSet.cs" />
+    <Compile Include="System\Collections\Frozen\SmallValueTypeComparableFrozenDictionary.cs" />
+    <Compile Include="System\Collections\Frozen\SmallValueTypeComparableFrozenSet.cs" />
     <Compile Include="System\Collections\Frozen\SmallValueTypeDefaultComparerFrozenDictionary.cs" />
     <Compile Include="System\Collections\Frozen\SmallValueTypeDefaultComparerFrozenSet.cs" />
     <Compile Include="System\Collections\Frozen\ValueTypeDefaultComparerFrozenDictionary.cs" />

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

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Text;
+
 namespace System.Collections.Frozen
 {
     /// <summary>
@@ -28,5 +30,53 @@ internal static class Constants
         /// hashing ValueTypeDefaultComparerFrozenDictionary/Set.
         /// </remarks>
         public const int MaxItemsInSmallValueTypeFrozenCollection = 10;
+
+        /// <summary>
+        /// Whether the <typeparamref name="T"/> is known to implement <see cref="IComparable{T}"/> safely and efficiently,
+        /// such that its comparison operations should be used in searching for types in small collections.
+        /// </summary>
+        /// <remarks>
+        /// This does not automatically return true for any type that implements <see cref="IComparable{T}"/>.
+        /// Doing so leads to problems for container types (e.g. ValueTuple{T1, T2}) where the
+        /// container implements <see cref="IComparable{T}"/> to delegate to its contained items' implementation
+        /// but those then don't provide such support.
+        /// </remarks>
+        public static bool IsKnownComparable<T>() =>
+            // This list covers all of the IComparable<T> value types in Corelib that aren't containers (like ValueTuple).
+            typeof(T) == typeof(bool) ||
+            typeof(T) == typeof(sbyte) ||
+            typeof(T) == typeof(byte) ||
+            typeof(T) == typeof(char) ||
+            typeof(T) == typeof(short) ||
+            typeof(T) == typeof(ushort) ||
+            typeof(T) == typeof(int) ||
+            typeof(T) == typeof(uint) ||
+            typeof(T) == typeof(long) ||
+            typeof(T) == typeof(ulong) ||
+            typeof(T) == typeof(nint) ||
+            typeof(T) == typeof(nuint) ||
+            typeof(T) == typeof(decimal) ||
+            typeof(T) == typeof(float) ||
+            typeof(T) == typeof(double) ||
+            typeof(T) == typeof(decimal) ||
+            typeof(T) == typeof(TimeSpan) ||
+            typeof(T) == typeof(DateTime) ||
+            typeof(T) == typeof(DateTimeOffset) ||
+            typeof(T) == typeof(Guid) ||
+#if NETCOREAPP3_0_OR_GREATER
+            typeof(T) == typeof(Rune) ||
+#endif
+#if NET5_0_OR_GREATER
+            typeof(T) == typeof(Half) ||
+#endif
+#if NET6_0_OR_GREATER
+            typeof(T) == typeof(DateOnly) ||
+            typeof(T) == typeof(TimeOnly) ||
+#endif
+#if NET7_0_OR_GREATER
+            typeof(T) == typeof(Int128) ||
+            typeof(T) == typeof(UInt128) ||
+#endif
+            typeof(T).IsEnum;
     }
 }

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

@@ -6,7 +6,6 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
-using System.Numerics;
 using System.Runtime.CompilerServices;
 
 namespace System.Collections.Frozen
@@ -202,14 +201,26 @@ private static FrozenDictionary<TKey, TValue> ChooseImplementationOptimizedForRe
                 {
                     if (source.Count <= Constants.MaxItemsInSmallValueTypeFrozenCollection)
                     {
+                        // If the key is a something we know we can efficiently compare, use a specialized implementation
+                        // that will enable quickly ruling out values outside of the range of keys stored.
+                        if (Constants.IsKnownComparable<TKey>())
+                        {
+                            return (FrozenDictionary<TKey, TValue>)(object)new SmallValueTypeComparableFrozenDictionary<TKey, TValue>(source);
+                        }
+
+                        // Otherwise, use an implementation optimized for a small number of value types using the default comparer.
                         return (FrozenDictionary<TKey, TValue>)(object)new SmallValueTypeDefaultComparerFrozenDictionary<TKey, TValue>(source);
                     }
 
+                    // Use a hash-based implementation.
+
+                    // For Int32 keys, we can reuse the key storage as the hash storage, saving on space and extra indirection.
                     if (typeof(TKey) == typeof(int))
                     {
                         return (FrozenDictionary<TKey, TValue>)(object)new Int32FrozenDictionary<TValue>((Dictionary<int, TValue>)(object)source);
                     }
 
+                    // Fallback to an implementation usable with any value type and the default comparer.
                     return new ValueTypeDefaultComparerFrozenDictionary<TKey, TValue>(source);
                 }
             }

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

@@ -133,14 +133,26 @@ private static FrozenSet<T> ChooseImplementationOptimizedForReading<T>(HashSet<T
                 {
                     if (source.Count <= Constants.MaxItemsInSmallValueTypeFrozenCollection)
                     {
+                        // If the type is a something we know we can efficiently compare, use a specialized implementation
+                        // that will enable quickly ruling out values outside of the range of keys stored.
+                        if (Constants.IsKnownComparable<T>())
+                        {
+                            return (FrozenSet<T>)(object)new SmallValueTypeComparableFrozenSet<T>(source);
+                        }
+
+                        // Otherwise, use an implementation optimized for a small number of value types using the default comparer.
                         return (FrozenSet<T>)(object)new SmallValueTypeDefaultComparerFrozenSet<T>(source);
                     }
 
+                    // Use a hash-based implementation.
+
+                    // For Int32 values, we can reuse the item storage as the hash storage, saving on space and extra indirection.
                     if (typeof(T) == typeof(int))
                     {
                         return (FrozenSet<T>)(object)new Int32FrozenSet((HashSet<int>)(object)source);
                     }
 
+                    // Fallback to an implementation usable with any value type and the default comparer.
                     return new ValueTypeDefaultComparerFrozenSet<T>(source);
                 }
             }

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

@@ -0,0 +1,70 @@
+// 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.Linq;
+using System.Runtime.CompilerServices;
+
+namespace System.Collections.Frozen
+{
+    /// <summary>Provides a frozen dictionary to use when the key is a value type, the default comparer is used, and the item count is small.</summary>
+    /// <remarks>
+    /// While not constrained in this manner, the <typeparamref name="TKey"/> must be an <see cref="IComparable{T}"/>.
+    /// This implementation is only used for a set of types that have a known-good <see cref="IComparable{T}"/> implementation; it's not
+    /// used for an <see cref="IComparable{T}"/> as we can't know for sure whether it's valid, e.g. if the TKey is a ValueTuple`2, it itself
+    /// is comparable, but its items might not be such that trying to compare it will result in exception.
+    /// </remarks>
+    internal sealed class SmallValueTypeComparableFrozenDictionary<TKey, TValue> : FrozenDictionary<TKey, TValue>
+        where TKey : notnull
+    {
+        private readonly TKey[] _keys;
+        private readonly TValue[] _values;
+        private readonly TKey _max;
+
+        internal SmallValueTypeComparableFrozenDictionary(Dictionary<TKey, TValue> source) : base(EqualityComparer<TKey>.Default)
+        {
+            Debug.Assert(default(TKey) is IComparable<TKey>);
+            Debug.Assert(default(TKey) is not null);
+            Debug.Assert(typeof(TKey).IsValueType);
+
+            Debug.Assert(source.Count != 0);
+            Debug.Assert(ReferenceEquals(source.Comparer, EqualityComparer<TKey>.Default));
+
+            _keys = source.Keys.ToArray();
+            _values = source.Values.ToArray();
+
+            Array.Sort(_keys, _values);
+            _max = _keys[_keys.Length - 1];
+        }
+
+        private protected override TKey[] KeysCore => _keys;
+        private protected override TValue[] ValuesCore => _values;
+        private protected override Enumerator GetEnumeratorCore() => new Enumerator(_keys, _values);
+        private protected override int CountCore => _keys.Length;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private protected override ref readonly TValue GetValueRefOrNullRefCore(TKey key)
+        {
+            if (Comparer<TKey>.Default.Compare(key, _max) <= 0)
+            {
+                TKey[] keys = _keys;
+                for (int i = 0; i < keys.Length; i++)
+                {
+                    int c = Comparer<TKey>.Default.Compare(key, keys[i]);
+                    if (c <= 0)
+                    {
+                        if (c == 0)
+                        {
+                            return ref _values[i];
+                        }
+
+                        break;
+                    }
+                }
+            }
+
+            return ref Unsafe.NullRef<TValue>();
+        }
+    }
+}

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

@@ -0,0 +1,76 @@
+// 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.Diagnostics;
+using System.Linq;
+
+namespace System.Collections.Frozen
+{
+    /// <summary>Provides a frozen set to use when the item is a value type, the default comparer is used, and the item count is small.</summary>
+    /// <remarks>
+    /// While not constrained in this manner, the <typeparamref name="T"/> must be an <see cref="IComparable{T}"/>.
+    /// This implementation is only used for a set of types that have a known-good <see cref="IComparable{T}"/> implementation; it's not
+    /// used for an <see cref="IComparable{T}"/> as we can't know for sure whether it's valid, e.g. if the T is a ValueTuple`2, it itself
+    /// is comparable, but its items might not be such that trying to compare it will result in exception.
+    /// </remarks>
+    internal sealed class SmallValueTypeComparableFrozenSet<T> : FrozenSetInternalBase<T, SmallValueTypeComparableFrozenSet<T>.GSW>
+    {
+        private readonly T[] _items;
+        private readonly T _max;
+
+        internal SmallValueTypeComparableFrozenSet(HashSet<T> source) : base(EqualityComparer<T>.Default)
+        {
+            Debug.Assert(default(T) is IComparable<T>);
+            Debug.Assert(default(T) is not null);
+            Debug.Assert(typeof(T).IsValueType);
+
+            Debug.Assert(source.Count != 0);
+            Debug.Assert(ReferenceEquals(source.Comparer, EqualityComparer<T>.Default));
+
+            _items = source.ToArray();
+
+            Array.Sort(_items);
+            _max = _items[_items.Length - 1];
+        }
+
+        private protected override T[] ItemsCore => _items;
+        private protected override Enumerator GetEnumeratorCore() => new Enumerator(_items);
+        private protected override int CountCore => _items.Length;
+
+        private protected override int FindItemIndex(T item)
+        {
+            if (Comparer<T>.Default.Compare(item, _max) <= 0)
+            {
+                T[] items = _items;
+                for (int i = 0; i < items.Length; i++)
+                {
+                    int c = Comparer<T>.Default.Compare(item, items[i]);
+                    if (c <= 0)
+                    {
+                        if (c == 0)
+                        {
+                            return i;
+                        }
+
+                        break;
+                    }
+                }
+            }
+
+            return -1;
+        }
+
+        internal struct GSW : IGenericSpecializedWrapper
+        {
+            private SmallValueTypeComparableFrozenSet<T> _set;
+            public void Store(FrozenSet<T> set) => _set = (SmallValueTypeComparableFrozenSet<T>)set;
+
+            public int Count => _set.Count;
+            public IEqualityComparer<T> Comparer => _set.Comparer;
+            public int FindItemIndex(T item) => _set.FindItemIndex(item);
+            public Enumerator GetEnumerator() => _set.GetEnumerator();
+        }
+    }
+}