Special-case empty enumerables in AsyncEnumerable (#112321)

The goal here is to handle known empty inputs in a way that reduces unnecessary allocation, e.g. avoiding an iterator allocation or avoiding an intermediate collection. This does not modify operators where there isn't such overhead.
dotnet · Feb 10, 2025 · 6070e26 · 6070e26
1 parent d05d6c2
commit 6070e26
Show file tree

Hide file tree

Showing 74 changed files with 920 additions and 186 deletions.
diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/AggregateBy.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/AggregateBy.cs
@@ -42,7 +42,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> AggregateBy<TSou
             ThrowHelper.ThrowIfNull(keySelector);
             ThrowHelper.ThrowIfNull(func);
 
-            return Impl(source, keySelector, seed, func, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, TAccumulate>>() :
+                Impl(source, keySelector, seed, func, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> Impl(
                 IAsyncEnumerable<TSource> source,
@@ -117,7 +119,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> AggregateBy<TSou
             ThrowHelper.ThrowIfNull(keySelector);
             ThrowHelper.ThrowIfNull(func);
 
-            return Impl(source, keySelector, seed, func, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, TAccumulate>>() :
+                Impl(source, keySelector, seed, func, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> Impl(
                 IAsyncEnumerable<TSource> source,
@@ -188,7 +192,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> AggregateBy<TSou
             ThrowHelper.ThrowIfNull(seedSelector);
             ThrowHelper.ThrowIfNull(func);
 
-            return Impl(source, keySelector, seedSelector, func, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, TAccumulate>>() :
+                Impl(source, keySelector, seedSelector, func, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> Impl(
                 IAsyncEnumerable<TSource> source,
@@ -264,7 +270,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> AggregateBy<TSou
             ThrowHelper.ThrowIfNull(seedSelector);
             ThrowHelper.ThrowIfNull(func);
 
-            return Impl(source, keySelector, seedSelector, func, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, TAccumulate>>() :
+                Impl(source, keySelector, seedSelector, func, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> Impl(
                 IAsyncEnumerable<TSource> source,
@@ -277,28 +285,26 @@ static async IAsyncEnumerable<KeyValuePair<TKey, TAccumulate>> Impl(
                 IAsyncEnumerator<TSource> enumerator = source.GetAsyncEnumerator(cancellationToken);
                 try
                 {
-                    if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
+                    if (await enumerator.MoveNextAsync().ConfigureAwait(false))
                     {
-                        yield break;
-                    }
+                        Dictionary<TKey, TAccumulate> dict = new(keyComparer);
 
-                    Dictionary<TKey, TAccumulate> dict = new(keyComparer);
-
-                    do
-                    {
-                        TSource value = enumerator.Current;
-                        TKey key = await keySelector(value, cancellationToken).ConfigureAwait(false);
+                        do
+                        {
+                            TSource value = enumerator.Current;
+                            TKey key = await keySelector(value, cancellationToken).ConfigureAwait(false);
 
-                        dict[key] = await func(
-                            dict.TryGetValue(key, out TAccumulate? acc) ? acc : await seedSelector(key, cancellationToken).ConfigureAwait(false),
-                            value,
-                            cancellationToken).ConfigureAwait(false);
-                    }
-                    while (await enumerator.MoveNextAsync().ConfigureAwait(false));
+                            dict[key] = await func(
+                                dict.TryGetValue(key, out TAccumulate? acc) ? acc : await seedSelector(key, cancellationToken).ConfigureAwait(false),
+                                value,
+                                cancellationToken).ConfigureAwait(false);
+                        }
+                        while (await enumerator.MoveNextAsync().ConfigureAwait(false));
 
-                    foreach (KeyValuePair<TKey, TAccumulate> countBy in dict)
-                    {
-                        yield return countBy;
+                        foreach (KeyValuePair<TKey, TAccumulate> countBy in dict)
+                        {
+                            yield return countBy;
+                        }
                     }
                 }
                 finally

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Cast.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Cast.cs
@@ -26,8 +26,9 @@ public static IAsyncEnumerable<TResult> Cast<TResult>( // satisfies the C# query
         {
             ThrowHelper.ThrowIfNull(source);
 
-            return source is IAsyncEnumerable<TResult> result ?
-                result :
+            return
+                source.IsKnownEmpty() ? Empty<TResult>() :
+                source as IAsyncEnumerable<TResult> ??
                 Impl(source, default);
 
             static async IAsyncEnumerable<TResult> Impl(

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Chunk.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Chunk.cs
@@ -30,7 +30,9 @@ public static IAsyncEnumerable<TSource[]> Chunk<TSource>(
             ThrowHelper.ThrowIfNull(source);
             ThrowHelper.ThrowIfNegativeOrZero(size);
 
-            return Chunk(source, size, default);
+            return
+                source.IsKnownEmpty() ? Empty<TSource[]>() :
+                Chunk(source, size, default);
 
             async static IAsyncEnumerable<TSource[]> Chunk(
                 IAsyncEnumerable<TSource> source,

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Concat.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Concat.cs
@@ -23,7 +23,10 @@ public static IAsyncEnumerable<TSource> Concat<TSource>(
             ThrowHelper.ThrowIfNull(first);
             ThrowHelper.ThrowIfNull(second);
 
-            return Impl(first, second, default);
+            return
+                first.IsKnownEmpty() ? second :
+                second.IsKnownEmpty() ? first :
+                Impl(first, second, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> first,

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/CountBy.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/CountBy.cs
@@ -28,7 +28,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, int>> CountBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(source);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(source, keySelector, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, int>>() :
+                Impl(source, keySelector, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, int>> Impl(
                 IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, IEqualityComparer<TKey>? keyComparer, [EnumeratorCancellation] CancellationToken cancellationToken)
@@ -83,7 +85,9 @@ public static IAsyncEnumerable<KeyValuePair<TKey, int>> CountBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(source);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(source, keySelector, keyComparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<KeyValuePair<TKey, int>>() :
+                Impl(source, keySelector, keyComparer, default);
 
             static async IAsyncEnumerable<KeyValuePair<TKey, int>> Impl(
                 IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, IEqualityComparer<TKey>? keyComparer, [EnumeratorCancellation] CancellationToken cancellationToken)

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Distinct.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Distinct.cs
@@ -21,7 +21,9 @@ public static IAsyncEnumerable<TSource> Distinct<TSource>(
         {
             ThrowHelper.ThrowIfNull(source);
 
-            return Impl(source, comparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(source, comparer, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> source,

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/DistinctBy.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/DistinctBy.cs
@@ -32,7 +32,9 @@ public static IAsyncEnumerable<TSource> DistinctBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(source);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(source, keySelector, comparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(source, keySelector, comparer, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> source,
@@ -86,7 +88,9 @@ public static IAsyncEnumerable<TSource> DistinctBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(source);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(source, keySelector, comparer, default);
+            return
+                source.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(source, keySelector, comparer, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> source,

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Empty.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Empty.cs
@@ -16,9 +16,14 @@ public static partial class AsyncEnumerable
         /// <returns>An empty <see cref="IAsyncEnumerable{T}"/> whose type argument is <typeparamref name="TResult"/>.</returns>
         public static IAsyncEnumerable<TResult> Empty<TResult>() => EmptyAsyncEnumerable<TResult>.Instance;
 
-        private sealed class EmptyAsyncEnumerable<TResult> : IAsyncEnumerable<TResult>, IAsyncEnumerator<TResult>
+        /// <summary>Determines whether <paramref name="source"/> is known to be an always-empty enumerable.</summary>
+        private static bool IsKnownEmpty<TResult>(this IAsyncEnumerable<TResult> source) =>
+            ReferenceEquals(source, EmptyAsyncEnumerable<TResult>.Instance);
+
+        private sealed class EmptyAsyncEnumerable<TResult> :
+            IAsyncEnumerable<TResult>, IAsyncEnumerator<TResult>, IOrderedAsyncEnumerable<TResult>
         {
-            public static EmptyAsyncEnumerable<TResult> Instance { get; } = new EmptyAsyncEnumerable<TResult>();
+            public static readonly EmptyAsyncEnumerable<TResult> Instance = new();
 
             public IAsyncEnumerator<TResult> GetAsyncEnumerator(CancellationToken cancellationToken = default) => this;
 
@@ -27,6 +32,18 @@ private sealed class EmptyAsyncEnumerable<TResult> : IAsyncEnumerable<TResult>,
             public TResult Current => default!;
 
             public ValueTask DisposeAsync() => default;
+
+            public IOrderedAsyncEnumerable<TResult> CreateOrderedAsyncEnumerable<TKey>(Func<TResult, TKey> keySelector, IComparer<TKey>? comparer, bool descending)
+            {
+                ThrowHelper.ThrowIfNull(keySelector);
+                return this;
+            }
+
+            public IOrderedAsyncEnumerable<TResult> CreateOrderedAsyncEnumerable<TKey>(Func<TResult, CancellationToken, ValueTask<TKey>> keySelector, IComparer<TKey>? comparer, bool descending)
+            {
+                ThrowHelper.ThrowIfNull(keySelector);
+                return this;
+            }
         }
     }
 }
diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Except.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/Except.cs
@@ -26,27 +26,44 @@ public static IAsyncEnumerable<TSource> Except<TSource>(
             ThrowHelper.ThrowIfNull(first);
             ThrowHelper.ThrowIfNull(second);
 
-            return Impl(first, second, comparer, default);
+            return
+                first.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(first, second, comparer, default);
 
             async static IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> first,
                 IAsyncEnumerable<TSource> second,
                 IEqualityComparer<TSource>? comparer,
                 [EnumeratorCancellation] CancellationToken cancellationToken)
             {
-                HashSet<TSource> set = new(comparer);
-
-                await foreach (TSource element in second.WithCancellation(cancellationToken).ConfigureAwait(false))
+                IAsyncEnumerator<TSource> firstEnumerator = first.GetAsyncEnumerator(cancellationToken);
+                try
                 {
-                    set.Add(element);
-                }
+                    if (!await firstEnumerator.MoveNextAsync().ConfigureAwait(false))
+                    {
+                        yield break;
+                    }
 
-                await foreach (TSource element in first.WithCancellation(cancellationToken).ConfigureAwait(false))
-                {
-                    if (set.Add(element))
+                    HashSet<TSource> set = new(comparer);
+
+                    await foreach (TSource element in second.WithCancellation(cancellationToken).ConfigureAwait(false))
                     {
-                        yield return element;
+                        set.Add(element);
                     }
+
+                    do
+                    {
+                        TSource firstElement = firstEnumerator.Current;
+                        if (set.Add(firstElement))
+                        {
+                            yield return firstElement;
+                        }
+                    }
+                    while (await firstEnumerator.MoveNextAsync().ConfigureAwait(false));
+                }
+                finally
+                {
+                    await firstEnumerator.DisposeAsync().ConfigureAwait(false);
                 }
             }
         }

diff --git a/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/ExceptBy.cs b/src/libraries/System.Linq.AsyncEnumerable/src/System/Linq/ExceptBy.cs
@@ -33,7 +33,9 @@ public static IAsyncEnumerable<TSource> ExceptBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(second);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(first, second, keySelector, comparer, default);
+            return
+                first.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(first, second, keySelector, comparer, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> first,
@@ -42,19 +44,34 @@ static async IAsyncEnumerable<TSource> Impl(
                 IEqualityComparer<TKey>? comparer,
                 [EnumeratorCancellation] CancellationToken cancellationToken)
             {
-                HashSet<TKey> set = new(comparer);
-
-                await foreach (TKey key in second.WithCancellation(cancellationToken).ConfigureAwait(false))
+                IAsyncEnumerator<TSource> firstEnumerator = first.GetAsyncEnumerator(cancellationToken);
+                try
                 {
-                    set.Add(key);
-                }
+                    if (!await firstEnumerator.MoveNextAsync().ConfigureAwait(false))
+                    {
+                        yield break;
+                    }
 
-                await foreach (TSource element in first.WithCancellation(cancellationToken).ConfigureAwait(false))
-                {
-                    if (set.Add(keySelector(element)))
+                    HashSet<TKey> set = new(comparer);
+
+                    await foreach (TKey key in second.WithCancellation(cancellationToken).ConfigureAwait(false))
                     {
-                        yield return element;
+                        set.Add(key);
                     }
+
+                    do
+                    {
+                        TSource firstElement = firstEnumerator.Current;
+                        if (set.Add(keySelector(firstElement)))
+                        {
+                            yield return firstElement;
+                        }
+                    }
+                    while (await firstEnumerator.MoveNextAsync().ConfigureAwait(false));
+                }
+                finally
+                {
+                    await firstEnumerator.DisposeAsync().ConfigureAwait(false);
                 }
             }
         }
@@ -82,7 +99,9 @@ public static IAsyncEnumerable<TSource> ExceptBy<TSource, TKey>(
             ThrowHelper.ThrowIfNull(second);
             ThrowHelper.ThrowIfNull(keySelector);
 
-            return Impl(first, second, keySelector, comparer, default);
+            return
+                first.IsKnownEmpty() ? Empty<TSource>() :
+                Impl(first, second, keySelector, comparer, default);
 
             static async IAsyncEnumerable<TSource> Impl(
                 IAsyncEnumerable<TSource> first,
@@ -91,19 +110,34 @@ static async IAsyncEnumerable<TSource> Impl(
                 IEqualityComparer<TKey>? comparer,
                 [EnumeratorCancellation] CancellationToken cancellationToken)
             {
-                HashSet<TKey> set = new(comparer);
-
-                await foreach (TKey key in second.WithCancellation(cancellationToken).ConfigureAwait(false))
+                IAsyncEnumerator<TSource> firstEnumerator = first.GetAsyncEnumerator(cancellationToken);
+                try
                 {
-                    set.Add(key);
-                }
+                    if (!await firstEnumerator.MoveNextAsync().ConfigureAwait(false))
+                    {
+                        yield break;
+                    }
 
-                await foreach (TSource element in first.WithCancellation(cancellationToken).ConfigureAwait(false))
-                {
-                    if (set.Add(await keySelector(element, cancellationToken).ConfigureAwait(false)))
+                    HashSet<TKey> set = new(comparer);
+
+                    await foreach (TKey key in second.WithCancellation(cancellationToken).ConfigureAwait(false))
                     {
-                        yield return element;
+                        set.Add(key);
                     }
+
+                    do
+                    {
+                        TSource firstElement = firstEnumerator.Current;
+                        if (set.Add(await keySelector(firstElement, cancellationToken).ConfigureAwait(false)))
+                        {
+                            yield return firstElement;
+                        }
+                    }
+                    while (await firstEnumerator.MoveNextAsync().ConfigureAwait(false));
+                }
+                finally
+                {
+                    await firstEnumerator.DisposeAsync().ConfigureAwait(false);
                 }
             }
         }