Move WithCancellation to separate file and reimplement.

Ix.NET/Source/System.Linq.Async/System/Linq/AsyncEnumerator.WithCancellation.cs

@@ -0,0 +1,54 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information. 
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.Collections.Generic
+{
+    public static partial class AsyncEnumerator
+    {
+        /// <summary>
+        /// Wraps the specified enumerator with an enumerator that checks for cancellation upon every invocation
+        /// of the <see cref="IAsyncEnumerator{T}.MoveNextAsync"/> method.
+        /// </summary>
+        /// <typeparam name="T">The type of the elements returned by the enumerator.</typeparam>
+        /// <param name="source">The enumerator to augment with cancellation support.</param>
+        /// <param name="cancellationToken">The cancellation token to observe.</param>
+        /// <returns>An enumerator that honors cancellation requests.</returns>
+        public static IAsyncEnumerator<T> WithCancellation<T>(this IAsyncEnumerator<T> source, CancellationToken cancellationToken)
+        {
+            if (source == null)
+                throw Error.ArgumentNull(nameof(source));
+
+            if (cancellationToken == default)
+                return source;
+
+            return new WithCancellationAsyncEnumerator<T>(source, cancellationToken);
+        }
+
+        private sealed class WithCancellationAsyncEnumerator<T> : IAsyncEnumerator<T>
+        {
+            private readonly IAsyncEnumerator<T> _source;
+            private readonly CancellationToken _cancellationToken;
+
+            public WithCancellationAsyncEnumerator(IAsyncEnumerator<T> source, CancellationToken cancellationToken)
+            {
+                _source = source;
+                _cancellationToken = cancellationToken;
+            }
+
+            public T Current => _source.Current;
+
+            public ValueTask DisposeAsync() => _source.DisposeAsync();
+
+            public ValueTask<bool> MoveNextAsync()
+            {
+                _cancellationToken.ThrowIfCancellationRequested();
+
+                return _source.MoveNextAsync(); // REVIEW: Signal cancellation through task or synchronously?
+            }
+        }
+    }
+}

Ix.NET/Source/System.Linq.Async/System/Linq/AsyncEnumerator.cs

@@ -12,7 +12,7 @@ namespace System.Collections.Generic
     /// <summary>
     /// Provides a set of extension methods for <see cref="IAsyncEnumerator{T}"/>.
     /// </summary>
-    public static class AsyncEnumerator
+    public static partial class AsyncEnumerator
     {
         /// <summary>
         /// Creates a new enumerator using the specified delegates implementing the members of <see cref="IAsyncEnumerator{T}"/>.
@@ -52,30 +52,6 @@ namespace System.Collections.Generic
             return source.MoveNextAsync();
         }
 
-        /// <summary>
-        /// Wraps the specified enumerator with an enumerator that checks for cancellation upon every invocation
-        /// of the <see cref="IAsyncEnumerator{T}.MoveNextAsync"/> method.
-        /// </summary>
-        /// <typeparam name="T">The type of the elements returned by the enumerator.</typeparam>
-        /// <param name="source">The enumerator to augment with cancellation support.</param>
-        /// <param name="cancellationToken">The cancellation token to observe.</param>
-        /// <returns>An enumerator that honors cancellation requests.</returns>
-        public static IAsyncEnumerator<T> WithCancellation<T>(this IAsyncEnumerator<T> source, CancellationToken cancellationToken)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-
-            return new AnonymousAsyncIterator<T>(
-                moveNext: () =>
-                {
-                    cancellationToken.ThrowIfCancellationRequested();
-                    return source.MoveNextAsync();
-                },
-                currentFunc: () => source.Current,
-                dispose: source.DisposeAsync
-            );
-        }
-
         internal static IAsyncEnumerator<T> Create<T>(Func<TaskCompletionSource<bool>, ValueTask<bool>> moveNext, Func<T> current, Func<ValueTask> dispose)
         {
             return new AnonymousAsyncIterator<T>(