From e025c739fab26c3aa379423253df132223f6ad2f Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Thu, 1 Feb 2024 21:06:55 -0500
Subject: [PATCH] Avoid allocating HashSet in Distinct() for some counts

If we can get the count for the underlying source and it's 0 or 1, we can avoid allocating the HashSet, as distinctness only matters when there are multiple elements.
---
 .../src/System/Linq/Distinct.SpeedOpt.cs      | 45 +++++++++++++++++--
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/libraries/System.Linq/src/System/Linq/Distinct.SpeedOpt.cs b/src/libraries/System.Linq/src/System/Linq/Distinct.SpeedOpt.cs
index 70e96b7ed68fac..c565e9c4c253e1 100644
--- a/src/libraries/System.Linq/src/System/Linq/Distinct.SpeedOpt.cs
+++ b/src/libraries/System.Linq/src/System/Linq/Distinct.SpeedOpt.cs
@@ -9,11 +9,50 @@ public static partial class Enumerable
     {
         private sealed partial class DistinctIterator<TSource> : IIListProvider<TSource>
         {
-            public TSource[] ToArray() => Enumerable.HashSetToArray(new HashSet<TSource>(_source, _comparer));
+            public TSource[] ToArray()
+            {
+                if (TryGetNonEnumeratedCount(_source, out int count) && count < 2)
+                {
+                    if (count == 1)
+                    {
+                        return [_source.First()];
+                    }
 
-            public List<TSource> ToList() => new List<TSource>(new HashSet<TSource>(_source, _comparer));
+                    return [];
+                }
 
-            public int GetCount(bool onlyIfCheap) => onlyIfCheap ? -1 : new HashSet<TSource>(_source, _comparer).Count;
+                return HashSetToArray(new HashSet<TSource>(_source, _comparer));
+            }
+
+            public List<TSource> ToList()
+            {
+                if (TryGetNonEnumeratedCount(_source, out int count) && count < 2)
+                {
+                    if (count == 1)
+                    {
+                        return new List<TSource>(1) { _source.First() };
+                    }
+
+                    return [];
+                }
+
+                return new List<TSource>(new HashSet<TSource>(_source, _comparer));
+            }
+
+            public int GetCount(bool onlyIfCheap)
+            {
+                if (TryGetNonEnumeratedCount(_source, out int count) && count < 2)
+                {
+                    return count;
+                }
+
+                if (onlyIfCheap)
+                {
+                    return -1;
+                }
+
+                return new HashSet<TSource>(_source, _comparer).Count;
+            }
         }
     }
 }