Use nearest prime sizing for collection trimming

src/coreclr/nativeaot/System.Private.Reflection.Execution/src/System.Private.Reflection.Execution.csproj

@@ -67,7 +67,7 @@
     <Compile Include="$(CompilerCommonPath)\Internal\LowLevelLinq\LowLevelEnumerable.ToArray.cs">
       <Link>Internal\LowLevelLinq\LowLevelEnumerable.ToArray.cs</Link>
     </Compile>
-    <Compile Include="$(LibrariesProjectRoot)\System.Private.CoreLib\src\System\Collections\HashHelpers.cs">
+    <Compile Include="$(LibrariesProjectRoot)\Common\src\System\Collections\HashHelpers.cs">
       <Link>System\Collections\HashHelpers.cs</Link>
     </Compile>
   </ItemGroup>

src/libraries/System.Private.CoreLib/src/System/Collections/HashHelpers.cs → src/libraries/Common/src/System/Collections/HashHelpers.cs

@@ -14,6 +14,7 @@ internal static partial class HashHelpers
         public const int MaxPrimeArrayLength = 0x7FFFFFC3;
 
         public const int HashPrime = 101;
+        private const int MinPrime = 3;
 
         // Table of prime numbers to use as hash table sizes.
         // A typical resize algorithm would pick the smallest prime number in this array
@@ -39,40 +40,114 @@ internal static partial class HashHelpers
 
         public static bool IsPrime(int candidate)
         {
-            if ((candidate & 1) != 0)
+            // This only tests hash table capacities, whose minimum is MinPrime, so 2 is intentionally excluded.
+            Debug.Assert(candidate >= MinPrime);
+
+            if ((candidate & 1) == 0 || (uint)candidate % MinPrime == 0)
             {
-                int limit = (int)Math.Sqrt(candidate);
-                for (int divisor = 3; divisor <= limit; divisor += 2)
+                return candidate == MinPrime;
+            }
+
+            return HasNoPrimeDivisors(candidate, (int)Math.Sqrt(candidate));
+        }
+
+        private static bool HasNoPrimeDivisors(int candidate, int limit)
+        {
+            // Every prime greater than 3 is 6k - 1 or 6k + 1, so test both candidates in each group.
+            for (int divisor = 5; divisor <= limit; divisor += 6)
+            {
+                if (candidate % divisor == 0 || candidate % (divisor + 2) == 0)
                 {
-                    if ((candidate % divisor) == 0)
-                        return false;
+                    return false;
                 }
-                return true;
             }
-            return candidate == 2;
+
+            return true;
         }
 
         public static int GetPrime(int min)
         {
             if (min < 0)
+            {
                 throw new ArgumentException(SR.Arg_HTCapacityOverflow);
+            }
 
-            foreach (int prime in Primes)
+            if (min <= MinPrime)
             {
-                if (prime >= min)
-                    return prime;
+                return MinPrime;
             }
 
-            // Outside of our predefined table. Compute the hard way.
-            for (int i = (min | 1); i < int.MaxValue; i += 2)
+            // A short linear scan is faster for the common small capacities.
+            const int LinearSearchCount = 16;
+
+            ReadOnlySpan<int> primes = Primes;
+            if (min <= primes[LinearSearchCount - 1])
             {
-                if (IsPrime(i) && ((i - 1) % HashPrime != 0))
-                    return i;
+                for (int i = 1; i < LinearSearchCount; i++)
+                {
+                    if (primes[i] >= min)
+                    {
+                        return primes[i];
+                    }
+                }
+            }
+            else
+            {
+                int index = primes.Slice(LinearSearchCount).BinarySearch(min);
+                index = index < 0 ? ~index : index;
+                index += LinearSearchCount;
+                if ((uint)index < (uint)primes.Length)
+                {
+                    return primes[index];
+                }
+            }
+
+            return GetPrimeAtLeastCore(min);
+        }
+
+        public static int GetPrimeAtLeast(int min)
+        {
+            if (min < 0)
+            {
+                throw new ArgumentException(SR.Arg_HTCapacityOverflow);
+            }
+
+            return min <= MinPrime ? MinPrime : GetPrimeAtLeastCore(min);
+        }
+
+        private static int GetPrimeAtLeastCore(int min)
+        {
+            Debug.Assert(min > MinPrime);
+
+            int candidate = min | 1;
+            uint remainder = (uint)candidate % 6;
+            if (remainder == 3)
+            {
+                candidate += 2;
+            }
+
+            int increment = remainder == 1 ? 4 : 2;
+            int limit = (int)Math.Sqrt(candidate);
+            long nextLimitSquared = (long)(limit + 1) * (limit + 1);
+            while (true)
+            {
+                while (nextLimitSquared <= candidate)
+                {
+                    limit++;
+                    nextLimitSquared = (long)(limit + 1) * (limit + 1);
+                }
+
+                if ((uint)(candidate - 1) % HashPrime != 0 && HasNoPrimeDivisors(candidate, limit))
+                {
+                    return candidate;
+                }
+
+                candidate += increment;
+                increment = 6 - increment;
             }
-            return min;
         }
 
-        // Returns size of hashtable to grow to.
+        // Returns the size of the hashtable to grow to.
         public static int ExpandPrime(int oldSize)
         {
             int newSize = 2 * oldSize;

src/libraries/System.Collections.Concurrent/src/System.Collections.Concurrent.csproj

@@ -15,7 +15,7 @@
     <Compile Include="System\Collections\Concurrent\OrderablePartitioner.cs" />
     <Compile Include="System\Collections\Concurrent\Partitioner.cs" />
     <Compile Include="System\Collections\Concurrent\PartitionerStatic.cs" />
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs"
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs"
              Link="System\Collections\HashHelpers.cs" />
     <Compile Include="$(CoreLibSharedDir)System\Collections\Concurrent\IProducerConsumerCollectionDebugView.cs"
              Link="System\Collections\Concurrent\IProducerConsumerCollectionDebugView.cs" />

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

@@ -14,7 +14,7 @@ The System.Collections.Immutable library is built-in as part of the shared frame
   <ItemGroup>
     <Compile Include="System\Polyfills.cs" />
     <Compile Include="System\Collections\ThrowHelper.cs" />
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs" Link="System\Collections\HashHelpers.cs" />
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs" Link="System\Collections\HashHelpers.cs" />
     <Compile Include="$(CoreLibSharedDir)System\Collections\Generic\DebugViewDictionaryItem.cs" Link="Common\System\Collections\Generic\DebugViewDictionaryItem.cs" />
     <Compile Include="$(CoreLibSharedDir)System\Collections\Generic\IDictionaryDebugView.cs" Link="Common\System\Collections\Generic\IDictionaryDebugView.cs" />
     <Compile Include="System\Collections\Frozen\Constants.cs" />

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

@@ -29,7 +29,7 @@
     <Compile Include="System\Collections\Generic\StackDebugView.cs" />
     <Compile Include="System\Collections\StructuralComparisons.cs" />
     <!-- Shared Common -->
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
     <Compile Include="$(CommonPath)System\Collections\Generic\BitHelper.cs" Link="Common\System\Collections\Generic\BitHelper.cs" />
     <Compile Include="$(CommonPath)System\Collections\Generic\EnumerableHelpers.cs" Link="Common\System\Collections\Generic\EnumerableHelpers.cs" />
     <Compile Include="$(CommonPath)System\Obsoletions.cs" Link="Common\System\Obsoletions.cs" />

src/libraries/System.Collections/tests/Generic/Dictionary/Dictionary.Generic.Tests.cs

@@ -468,6 +468,16 @@ public void TrimExcess_Generic_LargeInitialCapacity_TrimReducesSize()
             Assert.Equal(7, dictionary.EnsureCapacity(0));
         }
 
+        [Theory]
+        [InlineData(132, 137)]
+        [InlineData(607, 613)]
+        public void TrimExcess_Generic_UsesNearestValidPrime(int requestedCapacity, int expectedCapacity)
+        {
+            var dictionary = new Dictionary<TKey, TValue>(1000);
+            dictionary.TrimExcess(requestedCapacity);
+            Assert.Equal(expectedCapacity, dictionary.Capacity);
+        }
+
         [Theory]
         [InlineData(20)]
         [InlineData(23)]

src/libraries/System.Collections/tests/Generic/HashSet/HashSet.Generic.Tests.cs

@@ -231,6 +231,16 @@ public void HashHet_Generic_TrimExcess_LargePopulatedHashSet_TrimReducesSize(int
             Assert.Equal(clone, set);
         }
 
+        [Theory]
+        [InlineData(132, 137)]
+        [InlineData(607, 613)]
+        public void HashSet_Generic_TrimExcess_UsesNearestValidPrime(int requestedCapacity, int expectedCapacity)
+        {
+            var set = new HashSet<T>(1000);
+            set.TrimExcess(requestedCapacity);
+            Assert.Equal(expectedCapacity, set.Capacity);
+        }
+
         [Theory]
         [InlineData(10, 20, 0)]
         [InlineData(10, 20, 7)]

src/libraries/System.Linq.AsyncEnumerable/src/System.Linq.AsyncEnumerable.csproj

@@ -13,7 +13,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
     <Compile Include="System\Linq\AggregateAsync.cs" />
     <Compile Include="System\Linq\AggregateBy.cs" />
     <Compile Include="System\Linq\AllAsync.cs" />

src/libraries/System.Linq.Parallel/src/System.Linq.Parallel.csproj

@@ -14,7 +14,7 @@
 
   <!-- Compiled Source Files -->
   <ItemGroup>
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
     <Compile Include="System\Linq\Parallel\Channels\AsynchronousChannel.cs" />
     <Compile Include="System\Linq\Parallel\Channels\SynchronousChannel.cs" />
     <Compile Include="System\Linq\Parallel\Enumerables\AggregationMinMaxHelpers.cs" />

src/libraries/System.Linq/src/System.Linq.csproj

@@ -6,7 +6,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <Compile Include="$(CoreLibSharedDir)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs" Link="Common\System\Collections\HashHelpers.cs" />
     <Compile Include="System\Linq\Aggregate.cs" />
     <Compile Include="System\Linq\AnyAll.cs" />
     <Compile Include="System\Linq\AppendPrepend.cs" />

src/libraries/System.Private.CoreLib/src/System.Private.CoreLib.Shared.projitems

@@ -235,7 +235,6 @@
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\Generic\RandomizedStringEqualityComparer.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\Generic\ReferenceEqualityComparer.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\Generic\NonRandomizedStringEqualityComparer.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)System\Collections\HashHelpers.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\HashHelpers.SerializationInfoTable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\Hashtable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Collections\ICollection.cs" />
@@ -1542,6 +1541,9 @@
     <Compile Include="$(CommonPath)System\SR.cs">
       <Link>Common\System\SR.cs</Link>
     </Compile>
+    <Compile Include="$(CommonPath)System\Collections\HashHelpers.cs">
+      <Link>Common\System\Collections\HashHelpers.cs</Link>
+    </Compile>
     <Compile Include="$(CommonPath)System\Collections\Concurrent\IProducerConsumerQueue.cs">
       <Link>System\Collections\Concurrent\IProducerConsumerQueue.cs</Link>
     </Compile>

src/libraries/System.Private.CoreLib/src/System/Collections/Generic/Dictionary.cs

@@ -1250,6 +1250,7 @@ private void Resize(int newSize, bool forceNewHashCodes)
             // Value types never rehash
             Debug.Assert(!forceNewHashCodes || !typeof(TKey).IsValueType);
             Debug.Assert(_entries != null, "_entries should be non-null");
+            Debug.Assert(HashHelpers.IsPrime(newSize));
             Debug.Assert(newSize >= _entries.Length);
 
             Entry[] entries = new Entry[newSize];
@@ -1696,21 +1697,30 @@ public void TrimExcess(int capacity)
                 ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.capacity);
             }
 
-            int newSize = HashHelpers.GetPrime(capacity);
+            int newSize = HashHelpers.GetPrimeAtLeast(capacity);
             Entry[]? oldEntries = _entries;
-            int currentCapacity = oldEntries == null ? 0 : oldEntries.Length;
-            if (newSize >= currentCapacity)
+            if (oldEntries is null || newSize >= oldEntries.Length)
             {
                 return;
             }
 
-            int oldCount = _count;
             _version++;
-            Initialize(newSize);
 
-            Debug.Assert(oldEntries is not null);
+            Debug.Assert(HashHelpers.IsPrime(newSize));
+            Debug.Assert(newSize >= Count);
+
+            int[] buckets = new int[newSize];
+            Entry[] entries = new Entry[newSize];
+
+            // Assign member variables after both arrays allocated to guard against corruption from OOM if second fails
+            _freeList = -1;
+#if TARGET_64BIT
+            _fastModMultiplier = HashHelpers.GetFastModMultiplier((uint)newSize);
+#endif
+            _buckets = buckets;
+            _entries = entries;
 
-            CopyEntries(oldEntries, oldCount);
+            CopyEntries(oldEntries, _count);
         }
 
         private void CopyEntries(Entry[] entries, int count)