// 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.Collections.Generic; using System.Diagnostics; using System.Threading; using System.Threading.Tasks; namespace System.Linq { public static partial class AsyncEnumerable { public static IAsyncEnumerable Where(this IAsyncEnumerable source, Func predicate) { if (source == null) throw Error.ArgumentNull(nameof(source)); if (predicate == null) throw Error.ArgumentNull(nameof(predicate)); if (source is AsyncIterator iterator) { return iterator.Where(predicate); } // TODO: Can we add array/list optimizations here, does it make sense? return new WhereEnumerableAsyncIterator(source, predicate); } public static IAsyncEnumerable Where(this IAsyncEnumerable source, Func predicate) { if (source == null) throw Error.ArgumentNull(nameof(source)); if (predicate == null) throw Error.ArgumentNull(nameof(predicate)); return new WhereEnumerableWithIndexAsyncIterator(source, predicate); } public static IAsyncEnumerable Where(this IAsyncEnumerable source, Func> predicate) { if (source == null) throw Error.ArgumentNull(nameof(source)); if (predicate == null) throw Error.ArgumentNull(nameof(predicate)); if (source is AsyncIterator iterator) { return iterator.Where(predicate); } // TODO: Can we add array/list optimizations here, does it make sense? return new WhereEnumerableAsyncIteratorWithTask(source, predicate); } public static IAsyncEnumerable Where(this IAsyncEnumerable source, Func> predicate) { if (source == null) throw Error.ArgumentNull(nameof(source)); if (predicate == null) throw Error.ArgumentNull(nameof(predicate)); return new WhereEnumerableWithIndexAsyncIteratorWithTask(source, predicate); } private static Func CombinePredicates(Func predicate1, Func predicate2) { return x => predicate1(x) && predicate2(x); } private static Func> CombinePredicates(Func> predicate1, Func> predicate2) { return async x => await predicate1(x).ConfigureAwait(false) && await predicate2(x).ConfigureAwait(false); } internal sealed class WhereEnumerableAsyncIterator : AsyncIterator { private readonly Func _predicate; private readonly IAsyncEnumerable _source; private IAsyncEnumerator _enumerator; public WhereEnumerableAsyncIterator(IAsyncEnumerable source, Func predicate) { Debug.Assert(source != null); Debug.Assert(predicate != null); _source = source; _predicate = predicate; } public override AsyncIterator Clone() { return new WhereEnumerableAsyncIterator(_source, _predicate); } public override async ValueTask DisposeAsync() { if (_enumerator != null) { await _enumerator.DisposeAsync().ConfigureAwait(false); _enumerator = null; } await base.DisposeAsync().ConfigureAwait(false); } public override IAsyncEnumerable Select(Func selector) { return new WhereSelectEnumerableAsyncIterator(_source, _predicate, selector); } public override IAsyncEnumerable Where(Func predicate) { return new WhereEnumerableAsyncIterator(_source, CombinePredicates(_predicate, predicate)); } protected override async ValueTask MoveNextCore() { switch (_state) { case AsyncIteratorState.Allocated: _enumerator = _source.GetAsyncEnumerator(_cancellationToken); _state = AsyncIteratorState.Iterating; goto case AsyncIteratorState.Iterating; case AsyncIteratorState.Iterating: while (await _enumerator.MoveNextAsync().ConfigureAwait(false)) { var item = _enumerator.Current; if (_predicate(item)) { _current = item; return true; } } await DisposeAsync().ConfigureAwait(false); break; } return false; } } internal sealed class WhereEnumerableWithIndexAsyncIterator : AsyncIterator { private readonly Func _predicate; private readonly IAsyncEnumerable _source; private IAsyncEnumerator _enumerator; private int _index; public WhereEnumerableWithIndexAsyncIterator(IAsyncEnumerable source, Func predicate) { Debug.Assert(source != null); Debug.Assert(predicate != null); _source = source; _predicate = predicate; } public override AsyncIterator Clone() { return new WhereEnumerableWithIndexAsyncIterator(_source, _predicate); } public override async ValueTask DisposeAsync() { if (_enumerator != null) { await _enumerator.DisposeAsync().ConfigureAwait(false); _enumerator = null; } await base.DisposeAsync().ConfigureAwait(false); } protected override async ValueTask MoveNextCore() { switch (_state) { case AsyncIteratorState.Allocated: _enumerator = _source.GetAsyncEnumerator(_cancellationToken); _index = -1; _state = AsyncIteratorState.Iterating; goto case AsyncIteratorState.Iterating; case AsyncIteratorState.Iterating: while (await _enumerator.MoveNextAsync().ConfigureAwait(false)) { checked { _index++; } var item = _enumerator.Current; if (_predicate(item, _index)) { _current = item; return true; } } await DisposeAsync().ConfigureAwait(false); break; } return false; } } internal sealed class WhereEnumerableAsyncIteratorWithTask : AsyncIterator { private readonly Func> _predicate; private readonly IAsyncEnumerable _source; private IAsyncEnumerator _enumerator; public WhereEnumerableAsyncIteratorWithTask(IAsyncEnumerable source, Func> predicate) { Debug.Assert(source != null); Debug.Assert(predicate != null); _source = source; _predicate = predicate; } public override AsyncIterator Clone() { return new WhereEnumerableAsyncIteratorWithTask(_source, _predicate); } public override async ValueTask DisposeAsync() { if (_enumerator != null) { await _enumerator.DisposeAsync().ConfigureAwait(false); _enumerator = null; } await base.DisposeAsync().ConfigureAwait(false); } public override IAsyncEnumerable Where(Func> predicate) { return new WhereEnumerableAsyncIteratorWithTask(_source, CombinePredicates(_predicate, predicate)); } protected override async ValueTask MoveNextCore() { switch (_state) { case AsyncIteratorState.Allocated: _enumerator = _source.GetAsyncEnumerator(_cancellationToken); _state = AsyncIteratorState.Iterating; goto case AsyncIteratorState.Iterating; case AsyncIteratorState.Iterating: while (await _enumerator.MoveNextAsync().ConfigureAwait(false)) { var item = _enumerator.Current; if (await _predicate(item).ConfigureAwait(false)) { _current = item; return true; } } await DisposeAsync().ConfigureAwait(false); break; } return false; } } internal sealed class WhereEnumerableWithIndexAsyncIteratorWithTask : AsyncIterator { private readonly Func> _predicate; private readonly IAsyncEnumerable _source; private IAsyncEnumerator _enumerator; private int _index; public WhereEnumerableWithIndexAsyncIteratorWithTask(IAsyncEnumerable source, Func> predicate) { Debug.Assert(source != null); Debug.Assert(predicate != null); _source = source; _predicate = predicate; } public override AsyncIterator Clone() { return new WhereEnumerableWithIndexAsyncIteratorWithTask(_source, _predicate); } public override async ValueTask DisposeAsync() { if (_enumerator != null) { await _enumerator.DisposeAsync().ConfigureAwait(false); _enumerator = null; } await base.DisposeAsync().ConfigureAwait(false); } protected override async ValueTask MoveNextCore() { switch (_state) { case AsyncIteratorState.Allocated: _enumerator = _source.GetAsyncEnumerator(_cancellationToken); _index = -1; _state = AsyncIteratorState.Iterating; goto case AsyncIteratorState.Iterating; case AsyncIteratorState.Iterating: while (await _enumerator.MoveNextAsync().ConfigureAwait(false)) { checked { _index++; } var item = _enumerator.Current; if (await _predicate(item, _index).ConfigureAwait(false)) { _current = item; return true; } } await DisposeAsync().ConfigureAwait(false); break; } return false; } } internal sealed class WhereSelectEnumerableAsyncIterator : AsyncIterator { private readonly Func _predicate; private readonly Func _selector; private readonly IAsyncEnumerable _source; private IAsyncEnumerator _enumerator; public WhereSelectEnumerableAsyncIterator(IAsyncEnumerable source, Func predicate, Func selector) { Debug.Assert(source != null); Debug.Assert(predicate != null); Debug.Assert(selector != null); _source = source; _predicate = predicate; _selector = selector; } public override AsyncIterator Clone() { return new WhereSelectEnumerableAsyncIterator(_source, _predicate, _selector); } public override async ValueTask DisposeAsync() { if (_enumerator != null) { await _enumerator.DisposeAsync().ConfigureAwait(false); _enumerator = null; } await base.DisposeAsync().ConfigureAwait(false); } public override IAsyncEnumerable Select(Func selector) { return new WhereSelectEnumerableAsyncIterator(_source, _predicate, CombineSelectors(_selector, selector)); } protected override async ValueTask MoveNextCore() { switch (_state) { case AsyncIteratorState.Allocated: _enumerator = _source.GetAsyncEnumerator(_cancellationToken); _state = AsyncIteratorState.Iterating; goto case AsyncIteratorState.Iterating; case AsyncIteratorState.Iterating: while (await _enumerator.MoveNextAsync().ConfigureAwait(false)) { var item = _enumerator.Current; if (_predicate(item)) { _current = _selector(item); return true; } } await DisposeAsync().ConfigureAwait(false); break; } return false; } } } }