From d2880a818f86ebbc70f3cd566cc0436774a952c1 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Mon, 8 Jun 2020 01:44:31 +0900
Subject: [PATCH] Add UniTask.NextFrame, UniTask.WaitForEndOfFrame

---
 .../Plugins/UniTask/Runtime/UniTask.Delay.cs  | 127 +++++++++++++++++-
 1 file changed, 126 insertions(+), 1 deletion(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs
index cf4b6d4..30fe4bb 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Delay.cs
@@ -3,7 +3,6 @@
 using System;
 using System.Runtime.CompilerServices;
 using System.Threading;
-using Cysharp.Threading.Tasks.Internal;
 using UnityEngine;
 
 namespace Cysharp.Threading.Tasks
@@ -21,6 +20,30 @@ namespace Cysharp.Threading.Tasks
             return new UniTask(YieldPromise.Create(timing, cancellationToken, out var token), token);
         }
 
+        /// <summary>
+        /// Similar as UniTask.Yield but guaranteed run on next frame.
+        /// </summary>
+        public static UniTask NextFrame(PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default)
+        {
+            return new UniTask(NextFramePromise.Create(timing, cancellationToken, out var token), token);
+        }
+
+        /// <summary>
+        /// Same as UniTask.Yield(PlayerLoopTiming.LastPostLateUpdate).
+        /// </summary>
+        public static YieldAwaitable WaitForEndOfFrame()
+        {
+            return UniTask.Yield(PlayerLoopTiming.LastPostLateUpdate);
+        }
+
+        /// <summary>
+        /// Same as UniTask.Yield(PlayerLoopTiming.LastPostLateUpdate, cancellationToken).
+        /// </summary>
+        public static UniTask WaitForEndOfFrame(CancellationToken cancellationToken)
+        {
+            return UniTask.Yield(PlayerLoopTiming.LastPostLateUpdate, cancellationToken);
+        }
+
         public static UniTask DelayFrame(int delayFrameCount, PlayerLoopTiming delayTiming = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken))
         {
             if (delayFrameCount < 0)
@@ -152,6 +175,108 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
+        sealed class NextFramePromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<NextFramePromise>
+        {
+            static TaskPool<NextFramePromise> pool;
+            public NextFramePromise NextNode { get; set; }
+
+            static NextFramePromise()
+            {
+                TaskPool.RegisterSizeGetter(typeof(NextFramePromise), () => pool.Size);
+            }
+
+            int frameCount;
+            CancellationToken cancellationToken;
+            UniTaskCompletionSourceCore<AsyncUnit> core;
+
+            NextFramePromise()
+            {
+            }
+
+            public static IUniTaskSource Create(PlayerLoopTiming timing, CancellationToken cancellationToken, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                if (!pool.TryPop(out var result))
+                {
+                    result = new NextFramePromise();
+                }
+
+                result.frameCount = Time.frameCount;
+                result.cancellationToken = cancellationToken;
+
+                TaskTracker.TrackActiveTask(result, 3);
+
+                PlayerLoopHelper.AddAction(timing, result);
+
+                token = result.core.Version;
+                return result;
+            }
+
+            public void GetResult(short token)
+            {
+                try
+                {
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    TryReturn();
+                }
+            }
+
+            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)
+                {
+                    core.TrySetCanceled(cancellationToken);
+                    return false;
+                }
+
+                if (frameCount == Time.frameCount)
+                {
+                    return true;
+                }
+
+                core.TrySetResult(AsyncUnit.Default);
+                return false;
+            }
+
+            bool TryReturn()
+            {
+                TaskTracker.RemoveTracking(this);
+                core.Reset();
+                cancellationToken = default;
+                return pool.TryPush(this);
+            }
+
+            ~NextFramePromise()
+            {
+                if (TryReturn())
+                {
+                    GC.ReRegisterForFinalize(this);
+                }
+            }
+        }
+
         sealed class DelayFramePromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<DelayFramePromise>
         {
             static TaskPool<DelayFramePromise> pool;