UniTaskXR/Runtime/UnityAsyncExtensions.tt

<#@ template debug="false" hostspecific="false" language="C#" #>
<#@ assembly name="System.Core" #>
<#@ import namespace="System.Linq" #>
<#@ import namespace="System.Text" #>
<#@ import namespace="System.Collections.Generic" #>
<#@ output extension=".cs" #>
<#
    var types = new (string typeName, string returnType, string returnField)[]
    {
        ("AsyncOperation", "void", null),
        ("ResourceRequest", "UnityEngine.Object", "asset"),
        ("AssetBundleRequest", "UnityEngine.Object", "asset"), // allAssets?
        ("AssetBundleCreateRequest", "AssetBundle", "assetBundle"),
        ("UnityWebRequestAsyncOperation", "UnityWebRequest", "webRequest") // -> #if ENABLE_UNITYWEBREQUEST && (!UNITY_2019_1_OR_NEWER || UNITASK_WEBREQUEST_SUPPORT)
    };

    Func<string, string> ToUniTaskReturnType = x => (x == "void") ? "UniTask" : $"UniTask<{x}>";
    Func<string, string> ToIUniTaskSourceReturnType = x => (x == "void") ? "IUniTaskSource" : $"IUniTaskSource<{x}>";
    Func<(string typeName, string returnType, string returnField), bool> IsUnityWebRequest = x => x.returnType == "UnityWebRequest";
    Func<(string typeName, string returnType, string returnField), bool> IsAssetBundleModule = x => x.typeName == "AssetBundleRequest" || x.typeName == "AssetBundleCreateRequest";
    Func<(string typeName, string returnType, string returnField), bool> IsVoid = x => x.returnType == "void";
#>
#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member

using System;
using System.Runtime.CompilerServices;
using System.Threading;
using UnityEngine;
using Cysharp.Threading.Tasks.Internal;
#if ENABLE_UNITYWEBREQUEST && (!UNITY_2019_1_OR_NEWER || UNITASK_WEBREQUEST_SUPPORT)
using UnityEngine.Networking;
#endif

namespace Cysharp.Threading.Tasks
{
    public static partial class UnityAsyncExtensions
    {
<# foreach(var t in types) { #>
<# if(IsUnityWebRequest(t)) { #>
#if ENABLE_UNITYWEBREQUEST && (!UNITY_2019_1_OR_NEWER || UNITASK_WEBREQUEST_SUPPORT)
<# } else if(IsAssetBundleModule(t)) { #>
#if UNITASK_ASSETBUNDLE_SUPPORT
<# } #>
        #region <#= t.typeName #>

        public static <#= t.typeName #>Awaiter GetAwaiter(this <#= t.typeName #> asyncOperation)
        {
            Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation));
            return new <#= t.typeName #>Awaiter(asyncOperation);
        }

        public static <#= ToUniTaskReturnType(t.returnType) #> WithCancellation(this <#= t.typeName #> asyncOperation, CancellationToken cancellationToken)
        {
            return ToUniTask(asyncOperation, cancellationToken: cancellationToken);
        }

        public static <#= ToUniTaskReturnType(t.returnType) #> ToUniTask(this <#= t.typeName #> asyncOperation, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken))
        {
            Error.ThrowArgumentNullException(asyncOperation, nameof(asyncOperation));
            if (cancellationToken.IsCancellationRequested) return UniTask.FromCanceled<#= IsVoid(t) ? "" : "<" + t.returnType + ">" #>(cancellationToken);
<# if(IsUnityWebRequest(t)) { #>
            if (asyncOperation.isDone)
            {
                if (asyncOperation.webRequest.IsError())
                {
                    return UniTask.FromException<UnityWebRequest>(new UnityWebRequestException(asyncOperation.webRequest));
                }
                return UniTask.FromResult(asyncOperation.webRequest);
            }
<# } else { #>
            if (asyncOperation.isDone) return <#= IsVoid(t) ? "UniTask.CompletedTask" : $"UniTask.FromResult(asyncOperation.{t.returnField})" #>;
<# } #>
            return new <#= ToUniTaskReturnType(t.returnType) #>(<#= t.typeName #>ConfiguredSource.Create(asyncOperation, timing, progress, cancellationToken, out var token), token);
        }

        public struct <#= t.typeName #>Awaiter : ICriticalNotifyCompletion
        {
            <#= t.typeName #> asyncOperation;
            Action<AsyncOperation> continuationAction;

            public <#= t.typeName #>Awaiter(<#= t.typeName #> asyncOperation)
            {
                this.asyncOperation = asyncOperation;
                this.continuationAction = null;
            }

            public bool IsCompleted => asyncOperation.isDone;

            public <#= t.returnType #> GetResult()
            {
                if (continuationAction != null)
                {
                    asyncOperation.completed -= continuationAction;
                    continuationAction = null;
<# if (!IsVoid(t)) { #>
                    var result = <#= $"asyncOperation.{t.returnField}" #>;
                    asyncOperation = null;
<# if(IsUnityWebRequest(t)) { #>
                    if (result.IsError())
                    {
                        throw new UnityWebRequestException(result);
                    }
<# } #>
                    return result;
<# } else { #>
                    asyncOperation = null;
<# } #>
                }
                else
                {
<# if (!IsVoid(t)) { #>
                    var result = <#= $"asyncOperation.{t.returnField}" #>;
                    asyncOperation = null;
<# if(IsUnityWebRequest(t)) { #>
                    if (result.IsError())
                    {
                        throw new UnityWebRequestException(result);
                    }
<# } #>
                    return result;
<# } else { #>
                    asyncOperation = null;
<# } #>
                }
            }

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

            public void UnsafeOnCompleted(Action continuation)
            {
                Error.ThrowWhenContinuationIsAlreadyRegistered(continuationAction);
                continuationAction = PooledDelegate<AsyncOperation>.Create(continuation);
                asyncOperation.completed += continuationAction;
            }
        }

        sealed class <#= t.typeName #>ConfiguredSource : <#= ToIUniTaskSourceReturnType(t.returnType) #>, IPlayerLoopItem, ITaskPoolNode<<#= t.typeName #>ConfiguredSource>
        {
            static TaskPool<<#= t.typeName #>ConfiguredSource> pool;
            <#= t.typeName #>ConfiguredSource nextNode;
            public ref <#= t.typeName #>ConfiguredSource NextNode => ref nextNode;

            static <#= t.typeName #>ConfiguredSource()
            {
                TaskPool.RegisterSizeGetter(typeof(<#= t.typeName #>ConfiguredSource), () => pool.Size);
            }

            <#= t.typeName #> asyncOperation;
            IProgress<float> progress;
            CancellationToken cancellationToken;

            UniTaskCompletionSourceCore<<#= IsVoid(t) ? "AsyncUnit" : t.returnType #>> core;

            <#= t.typeName #>ConfiguredSource()
            {

            }

            public static <#= ToIUniTaskSourceReturnType(t.returnType) #> Create(<#= t.typeName #> asyncOperation, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, out short token)
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    return AutoResetUniTaskCompletionSource<#= IsVoid(t) ? "" : $"<{t.returnType}>" #>.CreateFromCanceled(cancellationToken, out token);
                }

                if (!pool.TryPop(out var result))
                {
                    result = new <#= t.typeName #>ConfiguredSource();
                }

                result.asyncOperation = asyncOperation;
                result.progress = progress;
                result.cancellationToken = cancellationToken;

                TaskTracker.TrackActiveTask(result, 3);

                PlayerLoopHelper.AddAction(timing, result);

                token = result.core.Version;
                return result;
            }

            public <#= t.returnType #> GetResult(short token)
            {
                try
                {
<# if (!IsVoid(t)) { #>
                    return core.GetResult(token);
<# } else { #>
                    core.GetResult(token);
<# } #>
                }
                finally
                {
                    TryReturn();
                }
            }

<# if (!IsVoid(t)) { #>
            void IUniTaskSource.GetResult(short token)
            {
                GetResult(token);
            }
<# } #>

            public UniTaskStatus GetStatus(short token)
            {
                return core.GetStatus(token);
            }

            public UniTaskStatus UnsafeGetStatus()
            {
                return core.UnsafeGetStatus();
            }

            public void OnCompleted(Action<object> continuation, object state, short token)
            {
                core.OnCompleted(continuation, state, token);
            }

            public bool MoveNext()
            {
                if (cancellationToken.IsCancellationRequested)
                {
<# if(IsUnityWebRequest(t)) { #>
                    asyncOperation.webRequest.Abort();
<# } #>
                    core.TrySetCanceled(cancellationToken);
                    return false;
                }

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

                if (asyncOperation.isDone)
                {
<# if(IsUnityWebRequest(t)) { #>
                    if (asyncOperation.webRequest.IsError())
                    {
                        core.TrySetException(new UnityWebRequestException(asyncOperation.webRequest));
                    }
                    else
                    {
                        core.TrySetResult(asyncOperation.webRequest);
                    }
<# } else { #>
                    core.TrySetResult(<#= IsVoid(t) ? "AsyncUnit.Default" : $"asyncOperation.{t.returnField}" #>);
<# } #>
                    return false;
                }

                return true;
            }

            bool TryReturn()
            {
                TaskTracker.RemoveTracking(this);
                core.Reset();
                asyncOperation = default;
                progress = default;
                cancellationToken = default;
                return pool.TryPush(this);
            }
        }

        #endregion
<# if(IsUnityWebRequest(t) || IsAssetBundleModule(t)) { #>
#endif
<# } #>

<# } #>
    }
}