From 4af685159026c8533e1d8e91481c89ee21055e57 Mon Sep 17 00:00:00 2001
From: Iaroslav Lakhin <iaroslav.lakhin@gmail.com>
Date: Sun, 18 Feb 2024 00:44:52 +0300
Subject: [PATCH] Add AsyncInstantiate support

---
 .../AsyncInstantiateOperationExtensions.cs    | 173 ++++++++++++++++++
 ...syncInstantiateOperationExtensions.cs.meta |   3 +
 2 files changed, 176 insertions(+)
 create mode 100644 src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs
 create mode 100644 src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs.meta

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs
new file mode 100644
index 0000000..658d4f4
--- /dev/null
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs
@@ -0,0 +1,173 @@
+#if UNITASK_ASYNCINSTANTIATE_SUPPORT
+using System;
+using System.Threading;
+using UnityEngine;
+using Object = UnityEngine.Object;
+
+namespace Cysharp.Threading.Tasks
+{
+    public static class AsyncInstantiateOperationExtensions
+    {
+        public static UniTask<T[]>.Awaiter GetAwaiter<T>(this AsyncInstantiateOperation<T> operation)
+            where T : Object
+        {
+            return ToUniTask(operation).GetAwaiter();
+        }
+
+        public static UniTask<T[]> WithCancellation<T>(this AsyncInstantiateOperation<T> operation, CancellationToken cancellationToken)
+            where T : Object
+        {
+            return ToUniTask(operation, cancellationToken: cancellationToken);
+        }
+
+        public static UniTask<T[]> ToUniTask<T>(this AsyncInstantiateOperation<T> operation, IProgress<float> progress = null, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default)
+            where T : Object
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                operation.Cancel();
+
+                return UniTask.FromCanceled<T[]>(cancellationToken);
+            }
+
+            if (operation.isDone)
+            {
+                return UniTask.FromResult(operation.Result);
+            }
+
+            return new UniTask<T[]>(AsyncInstantiateOperationConfiguredSource<T>.Create(operation, timing, progress, cancellationToken, out var token), token);
+        }
+
+        private sealed class AsyncInstantiateOperationConfiguredSource<T> : IUniTaskSource<T[]>, IPlayerLoopItem, ITaskPoolNode<AsyncInstantiateOperationConfiguredSource<T>>
+            where T : Object
+        {
+            private static TaskPool<AsyncInstantiateOperationConfiguredSource<T>> pool;
+            private AsyncInstantiateOperationConfiguredSource<T> nextNode;
+
+            private AsyncInstantiateOperation<T> operation;
+            private IProgress<float> progress;
+            private CancellationToken cancellationToken;
+
+            private bool completed;
+            private UniTaskCompletionSourceCore<T[]> core;
+
+            static AsyncInstantiateOperationConfiguredSource()
+            {
+                TaskPool.RegisterSizeGetter(typeof(AsyncInstantiateOperationConfiguredSource<T>), () => pool.Size);
+            }
+
+            public static IUniTaskSource<T[]> Create(AsyncInstantiateOperation<T> operation, PlayerLoopTiming timing, IProgress<float> progress, CancellationToken cancellationToken, out short token)
+            {
+                if (!pool.TryPop(out var source))
+                {
+                    source = new AsyncInstantiateOperationConfiguredSource<T>();
+                }
+
+                source.operation = operation;
+                source.progress = progress;
+                source.cancellationToken = cancellationToken;
+                source.completed = false;
+
+                TaskTracker.TrackActiveTask(source, 3);
+                PlayerLoopHelper.AddAction(timing, source);
+
+                operation.completed += source.OnCompleted;
+                token = source.core.Version;
+
+                return source;
+            }
+
+            private void OnCompleted(AsyncOperation asyncOperation)
+            {
+                operation.completed -= OnCompleted;
+
+                if (completed)
+                {
+                    Destroy();
+                }
+                else
+                {
+                    completed = true;
+
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        operation.Cancel();
+                        core.TrySetCanceled(cancellationToken);
+                    }
+                    else
+                    {
+                        core.TrySetResult(operation.Result);
+                    }
+                }
+            }
+
+            private void Destroy()
+            {
+                TaskTracker.RemoveTracking(this);
+
+                core.Reset();
+                operation = default;
+                progress = default;
+                cancellationToken = default;
+
+                pool.TryPush(this);
+            }
+
+            void IUniTaskSource.GetResult(short token)
+            {
+                core.GetResult(token);
+            }
+
+            UniTaskStatus IUniTaskSource.GetStatus(short token)
+            {
+                return core.GetStatus(token);
+            }
+
+            void IUniTaskSource.OnCompleted(Action<object> continuation, object state, short token)
+            {
+                core.OnCompleted(continuation, state, token);
+            }
+
+            UniTaskStatus IUniTaskSource.UnsafeGetStatus()
+            {
+                return core.UnsafeGetStatus();
+            }
+
+            T[] IUniTaskSource<T[]>.GetResult(short token)
+            {
+                return core.GetResult(token);
+            }
+
+            void IUniTaskSource<T[]>.OnCompleted(Action<object> continuation, object state, short token)
+            {
+                core.OnCompleted(continuation, state, token);
+            }
+
+            bool IPlayerLoopItem.MoveNext()
+            {
+                if (completed)
+                {
+                    Destroy();
+
+                    return false;
+                }
+
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    completed = true;
+                    operation.Cancel();
+                    core.TrySetCanceled(cancellationToken);
+
+                    return false;
+                }
+
+                progress?.Report(operation.progress);
+
+                return true;
+            }
+
+            ref AsyncInstantiateOperationConfiguredSource<T> ITaskPoolNode<AsyncInstantiateOperationConfiguredSource<T>>.NextNode => ref nextNode;
+        }
+    }
+}
+#endif
diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs.meta b/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs.meta
new file mode 100644
index 0000000..a23ba0e
--- /dev/null
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/AsyncInstantiateOperationExtensions.cs.meta
@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 2ebf81ddeba24644b2946dcef4f16d16
+timeCreated: 1708199291
\ No newline at end of file