/// <summary>
/// Completes a consumer with an exception.
/// </summary>
internal static void Fail <T>(this AsyncEnumerator <T> .AsyncConsumer consumer, Exception exception)
{
var src = new TaskCompletionSource <int>();
src.SetException(exception);
consumer.OnComplete(src.Task);
}
/// <summary>
/// Completes a consumer.
/// </summary>
internal static void Complete <T>(this AsyncEnumerator <T> .AsyncConsumer consumer)
{
var src = new TaskCompletionSource <int>();
src.SetResult(1);
consumer.OnComplete(src.Task);
}
internal static IAsyncEnumerable <T> Empty <T>()
{
return(new DelegatingAsyncEnumerable <T>(() =>
{
var enumerator = new AsyncEnumerator <T>();
enumerator.Consumer.Complete();
return enumerator;
}));
}
private static async void PopulateAsync <T>(this AsyncEnumerator <T> enumerator, Func <IAsyncConsumer <T>, Task> producer)
{
var consumer = enumerator.Consumer;
await consumer.WaitForStartAsync();
// Instead of awaiting the producer task, we use an explicit continuation here
// because we want to propagate any outcome, especially exceptions, to the consumer.
producer(consumer).ContinueWith(consumer.OnComplete);
}
/// <summary>
/// Creates an <see cref="IAsyncEnumerable{T}"/> instance from this iterator function.
/// </summary>
/// <typeparam name="T">The type of the elements in the returned enumerable.</typeparam>
/// <param name="producer">The asynchronous function used to populate the returned enumerable.</param>
/// <returns>An instance of <c>IAsyncEnumerable</c> that represents the items yielded by the function provided in <c>producer</c>.</returns>
public static IAsyncEnumerable <T> Enumerate <T>(Func <IAsyncConsumer <T>, Task> producer)
{
if (producer == null)
{
throw new ArgumentNullException("producer");
}
return(new DelegatingAsyncEnumerable <T>(() =>
{
var enumerator = new AsyncEnumerator <T>();
enumerator.PopulateAsync(producer);
return enumerator;
}));
}
/// <summary>
/// Groups the elements of a sequence according to a specified key selector function.
/// </summary>
/// <typeparam name="TSource">The type of the elements of <c>source</c>.</typeparam>
/// <typeparam name="TKey">The type of the key returned by <c>keySelector</c>.</typeparam>
/// <typeparam name="TElement"></typeparam>
/// <typeparam name="TResult"></typeparam>
/// <param name="source">An IEnumerable{T} whose elements to group.</param>
/// <param name="keySelector">A function to extract the key for each element.</param>
/// <param name="elementSelector">A function to create an intermediate item for each element.</param>
/// <param name="resultSelector">A function that creates an element in the returned sequence from a key and an intermediate item.</param>
/// <param name="comparer">An IEqualityComparer{T} to compare keys with.</param>
/// <returns>An IAsyncEnumerable{TResult} where each item represents a group of source items and an according key.</returns>
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 new ArgumentNullException("source");
}
if (keySelector == null)
{
throw new ArgumentNullException("keySelector");
}
if (elementSelector == null)
{
throw new ArgumentNullException("elementSelector");
}
if (resultSelector == null)
{
throw new ArgumentNullException("resultSelector");
}
return(Enumerate <TResult>(async consumer =>
{
// consumers of inner grouping enumerators by key
var innerConsumersByKeys = new Dictionary <TKey, AsyncEnumerator <TElement> .AsyncConsumer>(comparer);
// The following dictionaries hold tasks for awaiting demand in inner enumerators and according keys
// keys by task
var keysByWaitingForInnerMoveNextTask = new Dictionary <Task, TKey>();
// tasks by key
var waitingForInnerMoveNextTasksByKey = new Dictionary <TKey, Task>(comparer);
// Waiter used to await demand for either the outer enumerator or inner grouping enumerators
var waiter = new MultiWaiter <Task>();
// Keys that currently exists demanded for. These keys are not currently awaited by waiter
var pendingKeys = new HashSet <TKey>(comparer);
// Task for awaiting demand for the next grouping (outer enumerator)
Task waitingOuterMoveNextTask = null;
try
{
var sourceEnumerator = source.GetEnumerator();
bool completed = false;
do
{
do
{
completed = !await sourceEnumerator.MoveNextAsync();
if (completed)
{
break;
}
var sourceItem = sourceEnumerator.Current;
var intermediateItem = elementSelector(sourceItem);
var key = keySelector(sourceItem);
AsyncEnumerator <TElement> .AsyncConsumer innerConsumer;
if (!innerConsumersByKeys.TryGetValue(key, out innerConsumer))
{
// So we have a new key and therefore a new grouping
var groupEnumerator = new AsyncEnumerator <TElement>();
innerConsumer = groupEnumerator.Consumer;
innerConsumersByKeys.Add(key, innerConsumer);
if (waitingOuterMoveNextTask != null)
{
// If we are currently waiting for demand of the outer enumerator, we remove the task from the waiter
// as we will get a new waiter shortly.
waiter.Remove(waitingOuterMoveNextTask);
}
waitingOuterMoveNextTask = consumer.YieldAsync(resultSelector(key, groupEnumerator.AsCachedEnumerable()));
if (waitingOuterMoveNextTask.IsCompleted)
{
waitingOuterMoveNextTask = null;
}
else
{
waiter.Add(waitingOuterMoveNextTask);
// It's a new key, so we definitely have demand for it. However, we will satisfy
// the demand inside this iteration.
pendingKeys.Add(key);
}
}
else
{
// So we created the grouping already.
Task oldTask;
if (waitingForInnerMoveNextTasksByKey.TryGetValue(key, out oldTask))
{
// remove this group from pending waiters because we will get a new waiter task shortly
waitingForInnerMoveNextTasksByKey.Remove(key);
keysByWaitingForInnerMoveNextTask.Remove(oldTask);
waiter.Remove(oldTask);
}
else
{
// The only reason why we wouldn't be waiting for demand is, that
// the group is already pending.
Debug.Assert(pendingKeys.Contains(key));
}
}
var newTask = innerConsumer.YieldAsync(intermediateItem);
if (newTask.IsCompleted)
{
// If the task was completed synchronously, we just need to add the key
// to the set of pending keys but don't need to remember the completed task.
pendingKeys.Add(key);
}
else
{
waitingForInnerMoveNextTasksByKey.Add(key, newTask);
keysByWaitingForInnerMoveNextTask.Add(newTask, key);
waiter.Add(newTask);
pendingKeys.Remove(key);
}
} while ((pendingKeys.Count > 0) || (waitingOuterMoveNextTask == null));
if (completed)
{
break;
}
// We use a MultiWaiter to await demand instead of using TaskEx.WhenAny because
// of performance issues ( O(N log(N)) instead of O(N^2) )
Task nextTask = await waiter.WaitNextAsync();
if (nextTask == waitingOuterMoveNextTask)
{
waitingOuterMoveNextTask = null;
}
else
{
pendingKeys.Add(keysByWaitingForInnerMoveNextTask[nextTask]);
}
} while (!completed);
}
catch (Exception e)
{
foreach (var innerConsumer in innerConsumersByKeys.Values)
{
innerConsumer.Fail(e);
}
throw;
}
finally
{
waiter.Dispose();
}
foreach (var innerConsumer in innerConsumersByKeys.Values)
{
innerConsumer.Complete();
}
}));
}
public AsyncConsumer(AsyncEnumerator <T> parent)
{
this.parent = parent;
}