using Cysharp.Threading.Tasks.Internal; using System; using System.Collections.Generic; using System.Linq; using System.Threading; namespace Cysharp.Threading.Tasks.Linq { public static partial class UniTaskAsyncEnumerable { public static IUniTaskAsyncEnumerable GroupJoin(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func outerKeySelector, Func innerKeySelector, Func, TResult> resultSelector) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); return new GroupJoin(outer, inner, outerKeySelector, innerKeySelector, resultSelector, EqualityComparer.Default); } public static IUniTaskAsyncEnumerable GroupJoin(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func outerKeySelector, Func innerKeySelector, Func, TResult> resultSelector, IEqualityComparer comparer) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); Error.ThrowArgumentNullException(comparer, nameof(comparer)); return new GroupJoin(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer); } public static IUniTaskAsyncEnumerable GroupJoinAwait(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, UniTask> resultSelector) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); return new GroupJoinAwait(outer, inner, outerKeySelector, innerKeySelector, resultSelector, EqualityComparer.Default); } public static IUniTaskAsyncEnumerable GroupJoinAwait(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, UniTask> resultSelector, IEqualityComparer comparer) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); Error.ThrowArgumentNullException(comparer, nameof(comparer)); return new GroupJoinAwait(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer); } public static IUniTaskAsyncEnumerable GroupJoinAwaitWithCancellation(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, CancellationToken, UniTask> resultSelector) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); return new GroupJoinAwaitWithCancellation(outer, inner, outerKeySelector, innerKeySelector, resultSelector, EqualityComparer.Default); } public static IUniTaskAsyncEnumerable GroupJoinAwaitWithCancellation(this IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, CancellationToken, UniTask> resultSelector, IEqualityComparer comparer) { Error.ThrowArgumentNullException(outer, nameof(outer)); Error.ThrowArgumentNullException(inner, nameof(inner)); Error.ThrowArgumentNullException(outerKeySelector, nameof(outerKeySelector)); Error.ThrowArgumentNullException(innerKeySelector, nameof(innerKeySelector)); Error.ThrowArgumentNullException(resultSelector, nameof(resultSelector)); Error.ThrowArgumentNullException(comparer, nameof(comparer)); return new GroupJoinAwaitWithCancellation(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer); } } internal sealed class GroupJoin : IUniTaskAsyncEnumerable { readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func outerKeySelector; readonly Func innerKeySelector; readonly Func, TResult> resultSelector; readonly IEqualityComparer comparer; public GroupJoin(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func outerKeySelector, Func innerKeySelector, Func, TResult> resultSelector, IEqualityComparer comparer) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; } public IUniTaskAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) { return new _GroupJoin(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer, cancellationToken); } sealed class _GroupJoin : MoveNextSource, IUniTaskAsyncEnumerator { static readonly Action MoveNextCoreDelegate = MoveNextCore; readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func outerKeySelector; readonly Func innerKeySelector; readonly Func, TResult> resultSelector; readonly IEqualityComparer comparer; CancellationToken cancellationToken; ILookup lookup; IUniTaskAsyncEnumerator enumerator; UniTask.Awaiter awaiter; public _GroupJoin(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func outerKeySelector, Func innerKeySelector, Func, TResult> resultSelector, IEqualityComparer comparer, CancellationToken cancellationToken) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; this.cancellationToken = cancellationToken; TaskTracker.TrackActiveTask(this, 3); } public TResult Current { get; private set; } public UniTask MoveNextAsync() { cancellationToken.ThrowIfCancellationRequested(); completionSource.Reset(); if (lookup == null) { CreateLookup().Forget(); } else { SourceMoveNext(); } return new UniTask(this, completionSource.Version); } async UniTaskVoid CreateLookup() { try { lookup = await inner.ToLookupAsync(innerKeySelector, comparer, cancellationToken); enumerator = outer.GetAsyncEnumerator(cancellationToken); } catch (Exception ex) { completionSource.TrySetException(ex); return; } SourceMoveNext(); } void SourceMoveNext() { try { awaiter = enumerator.MoveNextAsync().GetAwaiter(); if (awaiter.IsCompleted) { MoveNextCore(this); } else { awaiter.SourceOnCompleted(MoveNextCoreDelegate, this); } } catch (Exception ex) { completionSource.TrySetException(ex); } } static void MoveNextCore(object state) { var self = (_GroupJoin)state; if (self.TryGetResult(self.awaiter, out var result)) { if (result) { var outer = self.enumerator.Current; var key = self.outerKeySelector(outer); var values = self.lookup[key]; self.Current = self.resultSelector(outer, values); self.completionSource.TrySetResult(true); } else { self.completionSource.TrySetResult(false); } } } public UniTask DisposeAsync() { TaskTracker.RemoveTracking(this); if (enumerator != null) { return enumerator.DisposeAsync(); } return default; } } } internal sealed class GroupJoinAwait : IUniTaskAsyncEnumerable { readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func> outerKeySelector; readonly Func> innerKeySelector; readonly Func, UniTask> resultSelector; readonly IEqualityComparer comparer; public GroupJoinAwait(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, UniTask> resultSelector, IEqualityComparer comparer) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; } public IUniTaskAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) { return new _GroupJoinAwait(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer, cancellationToken); } sealed class _GroupJoinAwait : MoveNextSource, IUniTaskAsyncEnumerator { static readonly Action MoveNextCoreDelegate = MoveNextCore; readonly static Action ResultSelectCoreDelegate = ResultSelectCore; readonly static Action OuterKeySelectCoreDelegate = OuterKeySelectCore; readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func> outerKeySelector; readonly Func> innerKeySelector; readonly Func, UniTask> resultSelector; readonly IEqualityComparer comparer; CancellationToken cancellationToken; ILookup lookup; IUniTaskAsyncEnumerator enumerator; TOuter outerValue; UniTask.Awaiter awaiter; UniTask.Awaiter outerKeyAwaiter; UniTask.Awaiter resultAwaiter; public _GroupJoinAwait(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, UniTask> resultSelector, IEqualityComparer comparer, CancellationToken cancellationToken) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; this.cancellationToken = cancellationToken; TaskTracker.TrackActiveTask(this, 3); } public TResult Current { get; private set; } public UniTask MoveNextAsync() { cancellationToken.ThrowIfCancellationRequested(); completionSource.Reset(); if (lookup == null) { CreateLookup().Forget(); } else { SourceMoveNext(); } return new UniTask(this, completionSource.Version); } async UniTaskVoid CreateLookup() { try { lookup = await inner.ToLookupAwaitAsync(innerKeySelector, comparer, cancellationToken); enumerator = outer.GetAsyncEnumerator(cancellationToken); } catch (Exception ex) { completionSource.TrySetException(ex); return; } SourceMoveNext(); } void SourceMoveNext() { try { awaiter = enumerator.MoveNextAsync().GetAwaiter(); if (awaiter.IsCompleted) { MoveNextCore(this); } else { awaiter.SourceOnCompleted(MoveNextCoreDelegate, this); } } catch (Exception ex) { completionSource.TrySetException(ex); } } static void MoveNextCore(object state) { var self = (_GroupJoinAwait)state; if (self.TryGetResult(self.awaiter, out var result)) { if (result) { try { self.outerValue = self.enumerator.Current; self.outerKeyAwaiter = self.outerKeySelector(self.outerValue).GetAwaiter(); if (self.outerKeyAwaiter.IsCompleted) { OuterKeySelectCore(self); } else { self.outerKeyAwaiter.SourceOnCompleted(OuterKeySelectCoreDelegate, self); } } catch (Exception ex) { self.completionSource.TrySetException(ex); } } else { self.completionSource.TrySetResult(false); } } } static void OuterKeySelectCore(object state) { var self = (_GroupJoinAwait)state; if (self.TryGetResult(self.outerKeyAwaiter, out var result)) { try { var values = self.lookup[result]; self.resultAwaiter = self.resultSelector(self.outerValue, values).GetAwaiter(); if (self.resultAwaiter.IsCompleted) { ResultSelectCore(self); } else { self.resultAwaiter.SourceOnCompleted(ResultSelectCoreDelegate, self); } } catch (Exception ex) { self.completionSource.TrySetException(ex); } } } static void ResultSelectCore(object state) { var self = (_GroupJoinAwait)state; if (self.TryGetResult(self.resultAwaiter, out var result)) { self.Current = result; self.completionSource.TrySetResult(true); } } public UniTask DisposeAsync() { TaskTracker.RemoveTracking(this); if (enumerator != null) { return enumerator.DisposeAsync(); } return default; } } } internal sealed class GroupJoinAwaitWithCancellation : IUniTaskAsyncEnumerable { readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func> outerKeySelector; readonly Func> innerKeySelector; readonly Func, CancellationToken, UniTask> resultSelector; readonly IEqualityComparer comparer; public GroupJoinAwaitWithCancellation(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, CancellationToken, UniTask> resultSelector, IEqualityComparer comparer) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; } public IUniTaskAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) { return new _GroupJoinAwaitWithCancellation(outer, inner, outerKeySelector, innerKeySelector, resultSelector, comparer, cancellationToken); } sealed class _GroupJoinAwaitWithCancellation : MoveNextSource, IUniTaskAsyncEnumerator { static readonly Action MoveNextCoreDelegate = MoveNextCore; readonly static Action ResultSelectCoreDelegate = ResultSelectCore; readonly static Action OuterKeySelectCoreDelegate = OuterKeySelectCore; readonly IUniTaskAsyncEnumerable outer; readonly IUniTaskAsyncEnumerable inner; readonly Func> outerKeySelector; readonly Func> innerKeySelector; readonly Func, CancellationToken, UniTask> resultSelector; readonly IEqualityComparer comparer; CancellationToken cancellationToken; ILookup lookup; IUniTaskAsyncEnumerator enumerator; TOuter outerValue; UniTask.Awaiter awaiter; UniTask.Awaiter outerKeyAwaiter; UniTask.Awaiter resultAwaiter; public _GroupJoinAwaitWithCancellation(IUniTaskAsyncEnumerable outer, IUniTaskAsyncEnumerable inner, Func> outerKeySelector, Func> innerKeySelector, Func, CancellationToken, UniTask> resultSelector, IEqualityComparer comparer, CancellationToken cancellationToken) { this.outer = outer; this.inner = inner; this.outerKeySelector = outerKeySelector; this.innerKeySelector = innerKeySelector; this.resultSelector = resultSelector; this.comparer = comparer; this.cancellationToken = cancellationToken; TaskTracker.TrackActiveTask(this, 3); } public TResult Current { get; private set; } public UniTask MoveNextAsync() { cancellationToken.ThrowIfCancellationRequested(); completionSource.Reset(); if (lookup == null) { CreateLookup().Forget(); } else { SourceMoveNext(); } return new UniTask(this, completionSource.Version); } async UniTaskVoid CreateLookup() { try { lookup = await inner.ToLookupAwaitWithCancellationAsync(innerKeySelector, comparer, cancellationToken); enumerator = outer.GetAsyncEnumerator(cancellationToken); } catch (Exception ex) { completionSource.TrySetException(ex); return; } SourceMoveNext(); } void SourceMoveNext() { try { awaiter = enumerator.MoveNextAsync().GetAwaiter(); if (awaiter.IsCompleted) { MoveNextCore(this); } else { awaiter.SourceOnCompleted(MoveNextCoreDelegate, this); } } catch (Exception ex) { completionSource.TrySetException(ex); } } static void MoveNextCore(object state) { var self = (_GroupJoinAwaitWithCancellation)state; if (self.TryGetResult(self.awaiter, out var result)) { if (result) { try { self.outerValue = self.enumerator.Current; self.outerKeyAwaiter = self.outerKeySelector(self.outerValue, self.cancellationToken).GetAwaiter(); if (self.outerKeyAwaiter.IsCompleted) { OuterKeySelectCore(self); } else { self.outerKeyAwaiter.SourceOnCompleted(OuterKeySelectCoreDelegate, self); } } catch (Exception ex) { self.completionSource.TrySetException(ex); } } else { self.completionSource.TrySetResult(false); } } } static void OuterKeySelectCore(object state) { var self = (_GroupJoinAwaitWithCancellation)state; if (self.TryGetResult(self.outerKeyAwaiter, out var result)) { try { var values = self.lookup[result]; self.resultAwaiter = self.resultSelector(self.outerValue, values, self.cancellationToken).GetAwaiter(); if (self.resultAwaiter.IsCompleted) { ResultSelectCore(self); } else { self.resultAwaiter.SourceOnCompleted(ResultSelectCoreDelegate, self); } } catch (Exception ex) { self.completionSource.TrySetException(ex); } } } static void ResultSelectCore(object state) { var self = (_GroupJoinAwaitWithCancellation)state; if (self.TryGetResult(self.resultAwaiter, out var result)) { self.Current = result; self.completionSource.TrySetResult(true); } } public UniTask DisposeAsync() { TaskTracker.RemoveTracking(this); if (enumerator != null) { return enumerator.DisposeAsync(); } return default; } } } }