Add back IComparable-based optimization to FrozenDictionary/Set (#84301)

* Add back IComparable-based optimization to FrozenDictionary/Set

We previously had an optimization in FrozenDictionary/Set that special-cased a small number of value types when the default comparer is used... if that type was IComparable, we would sort the types, which would then a) enable us to immediately reject items larger than the known max, and b) enable us to stop searching once we hit an item larger than the one for which we were searching (which then implicitly also immediately rules out items smaller than the known min).

We removed that optimization in general because some prominent IComparable implementations don't always work, in particular container types like ValueTuple that implement IComparable but then it's only functional if the contained types are also comparable.  This commit puts it back for an allow-list of types.

* Address PR feedback

Author: stephentoub

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();
+        }
+    }
+}