Fix Contains for Distinct/Union with custom comparer (#112815)

* Fix Contains for Distinct/Union with custom comparer

The recently-added optimization needs to factor in a non-default comparer.

* Address feedback

---------

Co-authored-by: Eirik Tsarpalis <[email protected]>

Author: stephentoub

src/libraries/System.Linq/src/System/Linq/Distinct.SpeedOpt.cs

@@ -17,7 +17,13 @@ private sealed partial class DistinctIterator<TSource>
 
             public override TSource? TryGetFirst(out bool found) => _source.TryGetFirst(out found);
 
-            public override bool Contains(TSource value) => _source.Contains(value);
+            public override bool Contains(TSource value) =>
+                // If we're using the default comparer, then source.Distinct().Contains(value) is no different from
+                // source.Contains(value), as the Distinct() won't remove anything that could have caused
+                // Contains to return true. If, however, there is a custom comparer, Distinct might remove
+                // the elements that would have matched, and thus we can't skip it.
+                _comparer is null ? _source.Contains(value) :
+                base.Contains(value);
         }
     }
 }

src/libraries/System.Linq/src/System/Linq/Union.SpeedOpt.cs

@@ -48,16 +48,26 @@ private HashSet<TSource> FillSet()
 
             public override bool Contains(TSource value)
             {
-                IEnumerable<TSource>? source;
-                for (int i = 0; (source = GetEnumerable(i)) is not null; i++)
+                // If there's no comparer, then source1.Union(source2).Contains(value) is no different from
+                // source1.Contains(value) || source2.Contains(value), as Union's set semantics won't remove
+                // anything from either that could have matched. However, if there is a comparer, it's possible
+                // the Union could end up removing items that would have matched, and thus we can't skip it.
+                if (_comparer is null || _comparer == EqualityComparer<TSource>.Default)
                 {
-                    if (source.Contains(value))
+                    IEnumerable<TSource>? source;
+                    for (int i = 0; (source = GetEnumerable(i)) is not null; i++)
                     {
-                        return true;
+                        if (source.Contains(value))
+                        {
+                            return true;
+                        }
                     }
+
+                    return false;
                 }
 
-                return default;
+
+                return base.Contains(value);
             }
         }
     }

src/libraries/System.Linq/tests/ContainsTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
+using System.Diagnostics;
 using Xunit;
 
 namespace System.Linq.Tests
@@ -189,8 +190,9 @@ public void FollowingVariousOperators()
                 Assert.True(transform(source.Concat(source)).Distinct().Contains(2));
                 Assert.False(transform(source.Concat(source)).Distinct().Contains(4));
                 Assert.True(transform(source.Concat(source)).Distinct().Contains(1));
-                Assert.True(transform(source.Concat(source)).Distinct(EqualityComparer<int>.Create((x, y) => true)).Contains(2));
-                Assert.False(transform(source.Concat(source)).Distinct(EqualityComparer<int>.Create((x, y) => true)).Contains(0));
+                Assert.True(transform(source.Concat(source)).Distinct(EqualityComparer<int>.Create((x, y) => true, x => 0)).Contains(1));
+                Assert.False(transform(source.Concat(source)).Distinct(EqualityComparer<int>.Create((x, y) => true, x => 0)).Contains(2));
+                Assert.False(transform(source.Concat(source)).Distinct(EqualityComparer<int>.Create((x, y) => true, x => 0)).Contains(0));
 
                 // OrderBy
                 Assert.True(transformedSource.OrderBy(x => x).Contains(2));
@@ -246,13 +248,21 @@ public void FollowingVariousOperators()
 
                 // Union
                 Assert.True(transformedSource.Union(transform([4])).Contains(4));
+                Assert.True(transformedSource.Union(transform([4]), EqualityComparer<int>.Create((x, y) => true, x => 0)).Contains(1));
+                Assert.False(transformedSource.Union(transform([4]), EqualityComparer<int>.Create((x, y) => true, x => 0)).Contains(4));
                 Assert.False(transformedSource.Union(transform([3])).Contains(4));
             }
 
             // DefaultIfEmpty
             Assert.True(Enumerable.Empty<int>().DefaultIfEmpty(1).Contains(1));
             Assert.False(Enumerable.Empty<int>().DefaultIfEmpty(1).Contains(0));
 
+            // Distinct
+            Assert.True(new string[] { "a", "A" }.Distinct().Contains("a"));
+            Assert.True(new string[] { "a", "A" }.Distinct().Contains("A"));
+            Assert.True(new string[] { "a", "A" }.Distinct(StringComparer.OrdinalIgnoreCase).Contains("a"));
+            Assert.False(new string[] { "a", "A" }.Distinct(StringComparer.OrdinalIgnoreCase).Contains("A"));
+
             // Repeat
             Assert.True(Enumerable.Repeat(1, 5).Contains(1));
             Assert.False(Enumerable.Repeat(1, 5).Contains(2));
@@ -268,6 +278,12 @@ public void FollowingVariousOperators()
             Assert.False(new object[] { 1, "2", 3 }.OfType<int>().Contains(2));
             Assert.True(new object[] { 1, "2", 3 }.OfType<string>().Contains("2"));
             Assert.False(new object[] { 1, "2", 3 }.OfType<string>().Contains("4"));
+
+            // Union
+            Assert.True(new string[] { "a" }.Union(new string[] { "A" }).Contains("a"));
+            Assert.True(new string[] { "a" }.Union(new string[] { "A" }).Contains("A"));
+            Assert.True(new string[] { "a" }.Union(new string[] { "A" }, StringComparer.OrdinalIgnoreCase).Contains("a"));
+            Assert.False(new string[] { "a" }.Union(new string[] { "A" }, StringComparer.OrdinalIgnoreCase).Contains("A"));
         }
     }
 }