Implementing more overloads of GroupBy.

Ix.NET/Source/System.Linq.Async/System/Linq/Operators/GroupBy.cs

@@ -3,7 +3,6 @@
 // See the LICENSE file in the project root for more information. 
 
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -20,15 +19,13 @@ namespace System.Linq
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> GroupBy<TSource, TKey>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector) =>
             new GroupedAsyncEnumerableWithTask<TSource, TKey>(source, keySelector, comparer: null);
 
-#if !NO_DEEP_CANCELLATION
-        public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> GroupBy<TSource, TKey>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector) =>
-            new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey>(source, keySelector, comparer: null);
-#endif
-
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> GroupBy<TSource, TKey>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, IEqualityComparer<TKey> comparer) =>
             new GroupedAsyncEnumerableWithTask<TSource, TKey>(source, keySelector, comparer);
 
 #if !NO_DEEP_CANCELLATION
+        public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> GroupBy<TSource, TKey>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector) =>
+            new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey>(source, keySelector, comparer: null);
+
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TSource>> GroupBy<TSource, TKey>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, IEqualityComparer<TKey> comparer) =>
             new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey>(source, keySelector, comparer);
 #endif
@@ -42,15 +39,13 @@ namespace System.Linq
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TElement>> GroupBy<TSource, TKey, TElement>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector) =>
             new GroupedAsyncEnumerableWithTask<TSource, TKey, TElement>(source, keySelector, elementSelector, comparer: null);
 
-#if !NO_DEEP_CANCELLATION
-        public static IAsyncEnumerable<IAsyncGrouping<TKey, TElement>> GroupBy<TSource, TKey, TElement>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector) =>
-            new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement>(source, keySelector, elementSelector, comparer: null);
-#endif
-
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TElement>> GroupBy<TSource, TKey, TElement>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, IEqualityComparer<TKey> comparer) =>
             new GroupedAsyncEnumerableWithTask<TSource, TKey, TElement>(source, keySelector, elementSelector, comparer);
 
 #if !NO_DEEP_CANCELLATION
+        public static IAsyncEnumerable<IAsyncGrouping<TKey, TElement>> GroupBy<TSource, TKey, TElement>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector) =>
+            new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement>(source, keySelector, elementSelector, comparer: null);
+
         public static IAsyncEnumerable<IAsyncGrouping<TKey, TElement>> GroupBy<TSource, TKey, TElement>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, IEqualityComparer<TKey> comparer) =>
             new GroupedAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement>(source, keySelector, elementSelector, comparer);
 #endif
@@ -64,105 +59,35 @@ namespace System.Linq
         public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TKey, IAsyncEnumerable<TSource>, ValueTask<TResult>> resultSelector) =>
             new GroupedResultAsyncEnumerableWithTask<TSource, TKey, TResult>(source, keySelector, resultSelector, comparer: null);
 
-#if !NO_DEEP_CANCELLATION
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TKey, IAsyncEnumerable<TSource>, CancellationToken, ValueTask<TResult>> resultSelector) =>
-            new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TResult>(source, keySelector, resultSelector, comparer: null);
-#endif
-
         public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TKey, IAsyncEnumerable<TSource>, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer) =>
             new GroupedResultAsyncEnumerableWithTask<TSource, TKey, TResult>(source, keySelector, resultSelector, comparer);
 
 #if !NO_DEEP_CANCELLATION
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TKey, IAsyncEnumerable<TSource>, CancellationToken, ValueTask<TResult>> resultSelector) =>
+            new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TResult>(source, keySelector, resultSelector, comparer: null);
+
         public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TKey, IAsyncEnumerable<TSource>, CancellationToken, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer) =>
             new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TResult>(source, keySelector, resultSelector, comparer);
 #endif
 
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, Func<TSource, TElement> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, TResult> resultSelector)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy(keySelector, elementSelector, comparer: null).Select(g => resultSelector(g.Key, g));
-        }
-
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, Func<TSource, TElement> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, TResult> resultSelector, IEqualityComparer<TKey> comparer)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy(keySelector, elementSelector, comparer).Select(g => resultSelector(g.Key, g));
-        }
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, Func<TSource, TElement> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, TResult> resultSelector) =>
+            new GroupedResultAsyncEnumerable<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer: null);
 
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> resultSelector)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy<TSource, TKey, TElement>(keySelector, elementSelector, comparer: null).Select(g => resultSelector(g.Key, g));
-        }
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, Func<TSource, TElement> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, TResult> resultSelector, IEqualityComparer<TKey> comparer) =>
+            new GroupedResultAsyncEnumerable<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer);
 
-#if !NO_DEEP_CANCELLATION
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> resultSelector)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy(keySelector, elementSelector, comparer: null).Select((g, ct) => resultSelector(g.Key, g, ct));
-        }
-#endif
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> resultSelector) =>
+            new GroupedResultAsyncEnumerableWithTask<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer: null);
 
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy(keySelector, elementSelector, comparer).Select(g => resultSelector(g.Key, g));
-        }
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer) =>
+            new GroupedResultAsyncEnumerableWithTask<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer);
 
 #if !NO_DEEP_CANCELLATION
-        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer)
-        {
-            if (source == null)
-                throw Error.ArgumentNull(nameof(source));
-            if (keySelector == null)
-                throw Error.ArgumentNull(nameof(keySelector));
-            if (elementSelector == null)
-                throw Error.ArgumentNull(nameof(elementSelector));
-            if (resultSelector == null)
-                throw Error.ArgumentNull(nameof(resultSelector));
-
-            return source.GroupBy(keySelector, elementSelector, comparer).Select((g, ct) => resultSelector(g.Key, g, ct));
-        }
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> resultSelector) =>
+            new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer: null);
+
+        public static IAsyncEnumerable<TResult> GroupBy<TSource, TKey, TElement, TResult>(this IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer) =>
+            new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement, TResult>(source, keySelector, elementSelector, resultSelector, comparer);
 #endif
 
         private sealed class GroupedResultAsyncEnumerable<TSource, TKey, TResult> : AsyncIterator<TResult>, IAsyncIListProvider<TResult>
@@ -295,7 +220,7 @@ namespace System.Linq
                 {
                     case AsyncIteratorState.Allocated:
                         _lookup = await Internal.LookupWithTask<TKey, TSource>.CreateAsync(_source, _keySelector, _comparer, _cancellationToken).ConfigureAwait(false);
-                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken);
+                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
                         _state = AsyncIteratorState.Iterating;
                         goto case AsyncIteratorState.Iterating;
 
@@ -385,7 +310,7 @@ namespace System.Linq
                 {
                     case AsyncIteratorState.Allocated:
                         _lookup = await Internal.LookupWithTask<TKey, TSource>.CreateAsync(_source, _keySelector, _comparer, _cancellationToken).ConfigureAwait(false);
-                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g, _cancellationToken).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken);
+                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g, _cancellationToken).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
                         _state = AsyncIteratorState.Iterating;
                         goto case AsyncIteratorState.Iterating;
 
@@ -434,6 +359,281 @@ namespace System.Linq
         }
 #endif
 
+        private sealed class GroupedResultAsyncEnumerable<TSource, TKey, TElement, TResult> : AsyncIterator<TResult>, IAsyncIListProvider<TResult>
+        {
+            private readonly IAsyncEnumerable<TSource> _source;
+            private readonly Func<TSource, TKey> _keySelector;
+            private readonly Func<TSource, TElement> _elementSelector;
+            private readonly Func<TKey, IAsyncEnumerable<TElement>, TResult> _resultSelector;
+            private readonly IEqualityComparer<TKey> _comparer;
+
+            private Internal.Lookup<TKey, TElement> _lookup;
+            private IEnumerator<TResult> _enumerator;
+
+            public GroupedResultAsyncEnumerable(IAsyncEnumerable<TSource> source, Func<TSource, TKey> keySelector, Func<TSource, TElement> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, TResult> resultSelector, IEqualityComparer<TKey> comparer)
+            {
+                _source = source ?? throw Error.ArgumentNull(nameof(source));
+                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
+                _elementSelector = elementSelector ?? throw Error.ArgumentNull(nameof(elementSelector));
+                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
+                _comparer = comparer;
+            }
+
+            public override AsyncIteratorBase<TResult> Clone()
+            {
+                return new GroupedResultAsyncEnumerable<TSource, TKey, TElement, TResult>(_source, _keySelector, _elementSelector, _resultSelector, _comparer);
+            }
+
+            public override async ValueTask DisposeAsync()
+            {
+                if (_enumerator != null)
+                {
+                    _enumerator.Dispose();
+                    _enumerator = null;
+                    _lookup = null;
+                }
+
+                await base.DisposeAsync().ConfigureAwait(false);
+            }
+
+            protected override async ValueTask<bool> MoveNextCore()
+            {
+                switch (_state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        _lookup = await Internal.Lookup<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, _cancellationToken).ConfigureAwait(false);
+                        _enumerator = _lookup.ApplyResultSelector(_resultSelector).GetEnumerator();
+                        _state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+                    case AsyncIteratorState.Iterating:
+                        if (_enumerator.MoveNext())
+                        {
+                            _current = _enumerator.Current;
+                            return true;
+                        }
+
+                        await DisposeAsync().ConfigureAwait(false);
+                        break;
+                }
+
+                return false;
+            }
+
+            public async ValueTask<TResult[]> ToArrayAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.Lookup<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return l.ToArray(_resultSelector);
+            }
+
+            public async ValueTask<List<TResult>> ToListAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.Lookup<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return l.ToList(_resultSelector);
+            }
+
+            public ValueTask<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
+            {
+                if (onlyIfCheap)
+                {
+                    return new ValueTask<int>(-1);
+                }
+
+                return Core();
+
+                async ValueTask<int> Core()
+                {
+                    var l = await Internal.Lookup<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+
+                    return l.Count;
+                }
+            }
+        }
+
+        private sealed class GroupedResultAsyncEnumerableWithTask<TSource, TKey, TElement, TResult> : AsyncIterator<TResult>, IAsyncIListProvider<TResult>
+        {
+            private readonly IAsyncEnumerable<TSource> _source;
+            private readonly Func<TSource, ValueTask<TKey>> _keySelector;
+            private readonly Func<TSource, ValueTask<TElement>> _elementSelector;
+            private readonly Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> _resultSelector;
+            private readonly IEqualityComparer<TKey> _comparer;
+
+            private Internal.LookupWithTask<TKey, TElement> _lookup;
+            private IAsyncEnumerator<TResult> _enumerator;
+
+            public GroupedResultAsyncEnumerableWithTask(IAsyncEnumerable<TSource> source, Func<TSource, ValueTask<TKey>> keySelector, Func<TSource, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer)
+            {
+                _source = source ?? throw Error.ArgumentNull(nameof(source));
+                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
+                _elementSelector = elementSelector ?? throw Error.ArgumentNull(nameof(elementSelector));
+                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
+                _comparer = comparer;
+            }
+
+            public override AsyncIteratorBase<TResult> Clone()
+            {
+                return new GroupedResultAsyncEnumerableWithTask<TSource, TKey, TElement, TResult>(_source, _keySelector, _elementSelector, _resultSelector, _comparer);
+            }
+
+            public override async ValueTask DisposeAsync()
+            {
+                if (_enumerator != null)
+                {
+                    await _enumerator.DisposeAsync().ConfigureAwait(false);
+                    _enumerator = null;
+                    _lookup = null;
+                }
+
+                await base.DisposeAsync().ConfigureAwait(false);
+            }
+
+            protected override async ValueTask<bool> MoveNextCore()
+            {
+                switch (_state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        _lookup = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, _cancellationToken).ConfigureAwait(false);
+                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
+                        _state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+                    case AsyncIteratorState.Iterating:
+                        if (await _enumerator.MoveNextAsync().ConfigureAwait(false))
+                        {
+                            _current = _enumerator.Current;
+                            return true;
+                        }
+
+                        await DisposeAsync().ConfigureAwait(false);
+                        break;
+                }
+
+                return false;
+            }
+
+            public async ValueTask<TResult[]> ToArrayAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return await l.ToArray(_resultSelector).ConfigureAwait(false);
+            }
+
+            public async ValueTask<List<TResult>> ToListAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return await l.ToList(_resultSelector).ConfigureAwait(false);
+            }
+
+            public ValueTask<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
+            {
+                if (onlyIfCheap)
+                {
+                    return new ValueTask<int>(-1);
+                }
+
+                return Core();
+
+                async ValueTask<int> Core()
+                {
+                    var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+
+                    return l.Count;
+                }
+            }
+        }
+
+#if !NO_DEEP_CANCELLATION
+        private sealed class GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement, TResult> : AsyncIterator<TResult>, IAsyncIListProvider<TResult>
+        {
+            private readonly IAsyncEnumerable<TSource> _source;
+            private readonly Func<TSource, CancellationToken, ValueTask<TKey>> _keySelector;
+            private readonly Func<TSource, CancellationToken, ValueTask<TElement>> _elementSelector;
+            private readonly Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> _resultSelector;
+            private readonly IEqualityComparer<TKey> _comparer;
+
+            private Internal.LookupWithTask<TKey, TElement> _lookup;
+            private IAsyncEnumerator<TResult> _enumerator;
+
+            public GroupedResultAsyncEnumerableWithTaskAndCancellation(IAsyncEnumerable<TSource> source, Func<TSource, CancellationToken, ValueTask<TKey>> keySelector, Func<TSource, CancellationToken, ValueTask<TElement>> elementSelector, Func<TKey, IAsyncEnumerable<TElement>, CancellationToken, ValueTask<TResult>> resultSelector, IEqualityComparer<TKey> comparer)
+            {
+                _source = source ?? throw Error.ArgumentNull(nameof(source));
+                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
+                _elementSelector = elementSelector ?? throw Error.ArgumentNull(nameof(elementSelector));
+                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
+                _comparer = comparer;
+            }
+
+            public override AsyncIteratorBase<TResult> Clone()
+            {
+                return new GroupedResultAsyncEnumerableWithTaskAndCancellation<TSource, TKey, TElement, TResult>(_source, _keySelector, _elementSelector, _resultSelector, _comparer);
+            }
+
+            public override async ValueTask DisposeAsync()
+            {
+                if (_enumerator != null)
+                {
+                    await _enumerator.DisposeAsync().ConfigureAwait(false);
+                    _enumerator = null;
+                    _lookup = null;
+                }
+
+                await base.DisposeAsync().ConfigureAwait(false);
+            }
+
+            protected override async ValueTask<bool> MoveNextCore()
+            {
+                switch (_state)
+                {
+                    case AsyncIteratorState.Allocated:
+                        _lookup = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, _cancellationToken).ConfigureAwait(false);
+                        _enumerator = _lookup.Select(async g => await _resultSelector(g.Key, g, _cancellationToken).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
+                        _state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+
+                    case AsyncIteratorState.Iterating:
+                        if (await _enumerator.MoveNextAsync().ConfigureAwait(false))
+                        {
+                            _current = _enumerator.Current;
+                            return true;
+                        }
+
+                        await DisposeAsync().ConfigureAwait(false);
+                        break;
+                }
+
+                return false;
+            }
+
+            public async ValueTask<TResult[]> ToArrayAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return await l.ToArray(_resultSelector, cancellationToken).ConfigureAwait(false);
+            }
+
+            public async ValueTask<List<TResult>> ToListAsync(CancellationToken cancellationToken)
+            {
+                var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+                return await l.ToList(_resultSelector, cancellationToken).ConfigureAwait(false);
+            }
+
+            public ValueTask<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
+            {
+                if (onlyIfCheap)
+                {
+                    return new ValueTask<int>(-1);
+                }
+
+                return Core();
+
+                async ValueTask<int> Core()
+                {
+                    var l = await Internal.LookupWithTask<TKey, TElement>.CreateAsync(_source, _keySelector, _elementSelector, _comparer, cancellationToken).ConfigureAwait(false);
+
+                    return l.Count;
+                }
+            }
+        }
+#endif
+
         private sealed class GroupedAsyncEnumerable<TSource, TKey, TElement> : AsyncIterator<IAsyncGrouping<TKey, TElement>>, IAsyncIListProvider<IAsyncGrouping<TKey, TElement>>
         {
             private readonly IAsyncEnumerable<TSource> _source;