UniTask extension methods for addressables AsyncOperationHandle

UnityAsyncOperationHandleExtensions.cs

#if CSHARP_7_OR_LATER || (UNITY_2018_3_OR_NEWER && (NET_STANDARD_2_0 || NET_4_6))
#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
using UnityEngine.ResourceManagement.AsyncOperations;
using System;
using System.Threading;
using UniRx.Async.Internal;

namespace UniRx.Async
{
    public static class UnityAsyncOperationHandleExtensions
    {
        public static AsyncOperationHandleAwaiter GetAwaiter(this AsyncOperationHandle asyncOperationHandle)
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");
            return new AsyncOperationHandleAwaiter(asyncOperationHandle);
        }

        public static UniTask<object> ToUniTask(this AsyncOperationHandle asyncOperationHandle)
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");
            return new UniTask<object>(new AsyncOperationHandleAwaiter(asyncOperationHandle));
        }

        public static UniTask<object> ConfigureAwait(this AsyncOperationHandle asyncOperationHandle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellation = default(CancellationToken))
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");

            var awaiter = new AsyncOperationHandleConfiguredAwaiter(asyncOperationHandle, progress, cancellation);
            if (!awaiter.IsCompleted)
            {
                PlayerLoopHelper.AddAction(timing, awaiter);
            }
            return new UniTask<object>(awaiter);
        }

        public static AsyncOperationHandleAwaiter<T> GetAwaiter<T>(this AsyncOperationHandle<T> asyncOperationHandle)
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");
            return new AsyncOperationHandleAwaiter<T>(asyncOperationHandle);
        }

        public static UniTask<T> ToUniTask<T>(this AsyncOperationHandle<T> asyncOperationHandle)
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");
            return new UniTask<T>(new AsyncOperationHandleAwaiter<T>(asyncOperationHandle));
        }

        public static UniTask<T> ConfigureAwait<T>(this AsyncOperationHandle<T> asyncOperationHandle, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellation = default(CancellationToken))
        {
            if (!asyncOperationHandle.IsValid()) throw new InvalidOperationException("Invalid AsyncOperationHandle.");

            var awaiter = new AsyncOperationHandleConfiguredAwaiter<T>(asyncOperationHandle, progress, cancellation);
            if (!awaiter.IsCompleted)
            {
                PlayerLoopHelper.AddAction(timing, awaiter);
            }
            return new UniTask<T>(awaiter);
        }
        
        public struct AsyncOperationHandleAwaiter : IAwaiter<object>
        {
            AsyncOperationHandle asyncOperation;
            Action<AsyncOperationHandle> continuationAction;
            AwaiterStatus status;
            object result;

            public AsyncOperationHandleAwaiter(AsyncOperationHandle asyncOperation)
            {
                this.status = asyncOperation.IsDone ? AwaiterStatus.Succeeded : AwaiterStatus.Pending;
                this.asyncOperation = (this.status.IsCompleted()) ? default : asyncOperation;
                this.result = (this.status.IsCompletedSuccessfully()) ? asyncOperation.Result : default;
                this.continuationAction = null;
            }

            public bool IsCompleted => status.IsCompleted();
            public AwaiterStatus Status => status;

            public object GetResult()
            {
                if (status == AwaiterStatus.Succeeded) return this.result;

                if (status == AwaiterStatus.Pending)
                {
                    // first timing of call
                    if (asyncOperation.IsDone)
                    {
                        status = AwaiterStatus.Succeeded;
                    }
                    else
                    {
                        Error.ThrowNotYetCompleted();
                    }
                }

                this.result = asyncOperation.Result;

                if (continuationAction != null)
                {
                    asyncOperation.Completed -= continuationAction;
                    asyncOperation = default; // remove reference.
                    continuationAction = null;
                }
                else
                {
                    asyncOperation = default; // remove reference.
                }

                return this.result;
            }

            void IAwaiter.GetResult() => GetResult();

            public void OnCompleted(Action continuation)
            {
                UnsafeOnCompleted(continuation);
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(continuationAction);
                continuationAction = continuation.AsFuncOfT<AsyncOperationHandle>();
                asyncOperation.Completed += continuationAction;
            }
        }

        class AsyncOperationHandleConfiguredAwaiter : IAwaiter<object>, IPlayerLoopItem
        {
            AsyncOperationHandle asyncOperation;
            IProgress<float> progress;
            CancellationToken cancellationToken;
            AwaiterStatus status;
            Action continuation;
            object result;

            public AsyncOperationHandleConfiguredAwaiter(AsyncOperationHandle asyncOperation, IProgress<float> progress, CancellationToken cancellationToken)
            {
                this.status = cancellationToken.IsCancellationRequested ? AwaiterStatus.Canceled
                            : asyncOperation.IsDone ? AwaiterStatus.Succeeded
                            : AwaiterStatus.Pending;

                if (this.status.IsCompletedSuccessfully()) this.result = asyncOperation.Result;
                if (this.status.IsCompleted()) return;

                this.asyncOperation = asyncOperation;
                this.progress = progress;
                this.cancellationToken = cancellationToken;
                this.continuation = null;
                this.result = null;

                TaskTracker.TrackActiveTask(this, 2);
            }

            public bool IsCompleted => status.IsCompleted();
            public AwaiterStatus Status => status;
            void IAwaiter.GetResult() => GetResult();

            public object GetResult()
            {
                if (status == AwaiterStatus.Succeeded) return this.result;

                if (status == AwaiterStatus.Canceled)
                {
                    Error.ThrowOperationCanceledException();
                }

                return Error.ThrowNotYetCompleted<object>();
            }

            public bool MoveNext()
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    InvokeContinuation(AwaiterStatus.Canceled);
                    return false;
                }

                if (progress != null)
                {
                    progress.Report(asyncOperation.PercentComplete);
                }

                if (asyncOperation.IsDone)
                {
                    this.result = asyncOperation.Result;
                    InvokeContinuation(AwaiterStatus.Succeeded);
                    return false;
                }

                return true;
            }

            void InvokeContinuation(AwaiterStatus status)
            {
                this.status = status;
                var cont = this.continuation;

                // cleanup
                TaskTracker.RemoveTracking(this);
                this.continuation = null;
                this.cancellationToken = CancellationToken.None;
                this.progress = null;
                this.asyncOperation = default;

                if (cont != null) cont.Invoke();
            }

            public void OnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(this.continuation);
                this.continuation = continuation;
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(this.continuation);
                this.continuation = continuation;
            }
        }
        
        public struct AsyncOperationHandleAwaiter<T> : IAwaiter<T>
        {
            AsyncOperationHandle<T> asyncOperation;
            Action<AsyncOperationHandle<T>> continuationAction;
            AwaiterStatus status;
            T result;

            public AsyncOperationHandleAwaiter(AsyncOperationHandle<T> asyncOperation)
            {
                this.status = asyncOperation.IsDone ? AwaiterStatus.Succeeded : AwaiterStatus.Pending;
                this.asyncOperation = (this.status.IsCompleted()) ? default : asyncOperation;
                this.result = (this.status.IsCompletedSuccessfully()) ? asyncOperation.Result : default;
                this.continuationAction = null;
            }

            public bool IsCompleted => status.IsCompleted();
            public AwaiterStatus Status => status;

            public T GetResult()
            {
                if (status == AwaiterStatus.Succeeded) return this.result;

                if (status == AwaiterStatus.Pending)
                {
                    // first timing of call
                    if (asyncOperation.IsDone)
                    {
                        status = AwaiterStatus.Succeeded;
                    }
                    else
                    {
                        Error.ThrowNotYetCompleted();
                    }
                }

                this.result = asyncOperation.Result;

                if (continuationAction != null)
                {
                    asyncOperation.Completed -= continuationAction;
                    asyncOperation = default; // remove reference.
                    continuationAction = null;
                }
                else
                {
                    asyncOperation = default; // remove reference.
                }

                return this.result;
            }

            void IAwaiter.GetResult() => GetResult();

            public void OnCompleted(Action continuation)
            {
                UnsafeOnCompleted(continuation);
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(continuationAction);
                continuationAction = continuation.AsFuncOfT<AsyncOperationHandle<T>>();
                asyncOperation.Completed += continuationAction;
            }
        }

        class AsyncOperationHandleConfiguredAwaiter<T> : IAwaiter<T>, IPlayerLoopItem
        {
            AsyncOperationHandle<T> asyncOperation;
            IProgress<float> progress;
            CancellationToken cancellationToken;
            AwaiterStatus status;
            Action continuation;
            T result;

            public AsyncOperationHandleConfiguredAwaiter(AsyncOperationHandle<T> asyncOperation, IProgress<float> progress, CancellationToken cancellationToken)
            {
                this.status = cancellationToken.IsCancellationRequested ? AwaiterStatus.Canceled
                            : asyncOperation.IsDone ? AwaiterStatus.Succeeded
                            : AwaiterStatus.Pending;

                if (this.status.IsCompletedSuccessfully()) this.result = asyncOperation.Result;
                if (this.status.IsCompleted()) return;

                this.asyncOperation = asyncOperation;
                this.progress = progress;
                this.cancellationToken = cancellationToken;
                this.continuation = null;
                this.result = default;

                TaskTracker.TrackActiveTask(this, 2);
            }

            public bool IsCompleted => status.IsCompleted();
            public AwaiterStatus Status => status;
            void IAwaiter.GetResult() => GetResult();

            public T GetResult()
            {
                if (status == AwaiterStatus.Succeeded) return this.result;

                if (status == AwaiterStatus.Canceled)
                {
                    Error.ThrowOperationCanceledException();
                }

                return Error.ThrowNotYetCompleted<T>();
            }

            public bool MoveNext()
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    InvokeContinuation(AwaiterStatus.Canceled);
                    return false;
                }

                if (progress != null)
                {
                    progress.Report(asyncOperation.PercentComplete);
                }

                if (asyncOperation.IsDone)
                {
                    this.result = asyncOperation.Result;
                    InvokeContinuation(AwaiterStatus.Succeeded);
                    return false;
                }

                return true;
            }

            void InvokeContinuation(AwaiterStatus status)
            {
                this.status = status;
                var cont = this.continuation;

                // cleanup
                TaskTracker.RemoveTracking(this);
                this.continuation = null;
                this.cancellationToken = CancellationToken.None;
                this.progress = null;
                this.asyncOperation = default;

                if (cont != null) cont.Invoke();
            }

            public void OnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(this.continuation);
                this.continuation = continuation;
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(this.continuation);
                this.continuation = continuation;
            }
        }
    }
}
#endif