Unity-Tools/Assets/Plugins/UniTask/Runtime/UniTask.WhenEach.cs

using Cysharp.Threading.Tasks.Internal;
using System;
using System.Collections.Generic;
using System.Runtime.ExceptionServices;
using System.Threading;

namespace Cysharp.Threading.Tasks
{
    public partial struct UniTask
    {
        public static IUniTaskAsyncEnumerable<WhenEachResult<T>> WhenEach<T>(IEnumerable<UniTask<T>> tasks)
        {
            return new WhenEachEnumerable<T>(tasks);
        }

        public static IUniTaskAsyncEnumerable<WhenEachResult<T>> WhenEach<T>(params UniTask<T>[] tasks)
        {
            return new WhenEachEnumerable<T>(tasks);
        }
    }

    public readonly struct WhenEachResult<T>
    {
        public T Result { get; }
        public Exception Exception { get; }

        //[MemberNotNullWhen(false, nameof(Exception))]
        public bool IsCompletedSuccessfully => Exception == null;

        //[MemberNotNullWhen(true, nameof(Exception))]
        public bool IsFaulted => Exception != null;

        public WhenEachResult(T result)
        {
            this.Result = result;
            this.Exception = null;
        }

        public WhenEachResult(Exception exception)
        {
            if (exception == null) throw new ArgumentNullException(nameof(exception));
            this.Result = default;
            this.Exception = exception;
        }

        public void TryThrow()
        {
            if (IsFaulted)
            {
                ExceptionDispatchInfo.Capture(Exception).Throw();
            }
        }

        public T GetResult()
        {
            if (IsFaulted)
            {
                ExceptionDispatchInfo.Capture(Exception).Throw();
            }
            return Result;
        }

        public override string ToString()
        {
            if (IsCompletedSuccessfully)
            {
                return Result?.ToString() ?? "";
            }
            else
            {
                return $"Exception{{{Exception.Message}}}";
            }
        }
    }

    internal enum WhenEachState : byte
    {
        NotRunning,
        Running,
        Completed
    }

    internal sealed class WhenEachEnumerable<T> : IUniTaskAsyncEnumerable<WhenEachResult<T>>
    {
        IEnumerable<UniTask<T>> source;

        public WhenEachEnumerable(IEnumerable<UniTask<T>> source)
        {
            this.source = source;
        }

        public IUniTaskAsyncEnumerator<WhenEachResult<T>> GetAsyncEnumerator(CancellationToken cancellationToken = default)
        {
            return new Enumerator(source, cancellationToken);
        }

        sealed class Enumerator : IUniTaskAsyncEnumerator<WhenEachResult<T>>
        {
            readonly IEnumerable<UniTask<T>> source;
            CancellationToken cancellationToken;

            Channel<WhenEachResult<T>> channel;
            IUniTaskAsyncEnumerator<WhenEachResult<T>> channelEnumerator;
            int completeCount;
            WhenEachState state;

            public Enumerator(IEnumerable<UniTask<T>> source, CancellationToken cancellationToken)
            {
                this.source = source;
                this.cancellationToken = cancellationToken;
            }

            public WhenEachResult<T> Current => channelEnumerator.Current;

            public UniTask<bool> MoveNextAsync()
            {
                cancellationToken.ThrowIfCancellationRequested();

                if (state == WhenEachState.NotRunning)
                {
                    state = WhenEachState.Running;
                    channel = Channel.CreateSingleConsumerUnbounded<WhenEachResult<T>>();
                    channelEnumerator = channel.Reader.ReadAllAsync().GetAsyncEnumerator(cancellationToken);

                    if (source is UniTask<T>[] array)
                    {
                        ConsumeAll(this, array, array.Length);
                    }
                    else
                    {
                        using (var rentArray = ArrayPoolUtil.Materialize(source))
                        {
                            ConsumeAll(this, rentArray.Array, rentArray.Length);
                        }
                    }
                }

                return channelEnumerator.MoveNextAsync();
            }

            static void ConsumeAll(Enumerator self, UniTask<T>[] array, int length)
            {
                for (int i = 0; i < length; i++)
                {
                    RunWhenEachTask(self, array[i], length).Forget();
                }
            }

            static async UniTaskVoid RunWhenEachTask(Enumerator self, UniTask<T> task, int length)
            {
                try
                {
                    var result = await task;
                    self.channel.Writer.TryWrite(new WhenEachResult<T>(result));
                }
                catch (Exception ex)
                {
                    self.channel.Writer.TryWrite(new WhenEachResult<T>(ex));
                }

                if (Interlocked.Increment(ref self.completeCount) == length)
                {
                    self.state = WhenEachState.Completed;
                    self.channel.Writer.TryComplete();
                }
            }

            public async UniTask DisposeAsync()
            {
                if (channelEnumerator != null)
                {
                    await channelEnumerator.DisposeAsync();
                }

                if (state != WhenEachState.Completed)
                {
                    state = WhenEachState.Completed;
                    channel.Writer.TryComplete(new OperationCanceledException());
                }
            }
        }
    }
}