From bbd5686816c1105585fdbc825d27f7c5f3dd31b4 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Tue, 19 May 2020 01:35:16 +0900
Subject: [PATCH] Add UniTask.WaitUntilCanceled

---
 .../Plugins/UniTask/Runtime/Progress.cs       | 13 ++-
 .../UniTask/Runtime/UniTask.WaitUntil.cs      | 91 ++++++++++++++++++-
 src/UniTask/Assets/Scenes/SandboxMain.cs      | 11 ++-
 3 files changed, 103 insertions(+), 12 deletions(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/Progress.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/Progress.cs
index c154c3d..ed11290 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/Progress.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/Progress.cs
@@ -1,7 +1,4 @@
-#if NET_4_6 || NET_STANDARD_2_0 || CSHARP_7_OR_LATER
-#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
-
-using System;
+using System;
 using System.Collections.Generic;
 using Cysharp.Threading.Tasks.Internal;
 
@@ -21,7 +18,11 @@ namespace Cysharp.Threading.Tasks
         public static IProgress<T> CreateOnlyValueChanged<T>(Action<T> handler, IEqualityComparer<T> comparer = null)
         {
             if (handler == null) return NullProgress<T>.Instance;
+#if UNITY_2018_3_OR_NEWER
             return new OnlyValueChangedProgress<T>(handler, comparer ?? UnityEqualityComparer.GetDefault<T>());
+#else
+            return new OnlyValueChangedProgress<T>(handler, comparer ?? EqualityComparer<T>.Default);
+#endif
         }
 
         sealed class NullProgress<T> : IProgress<T>
@@ -83,6 +84,4 @@ namespace Cysharp.Threading.Tasks
             }
         }
     }
-}
-
-#endif
\ No newline at end of file
+}
\ No newline at end of file
diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
index f11f264..6c2a37e 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
@@ -19,6 +19,11 @@ namespace Cysharp.Threading.Tasks
             return new UniTask(WaitWhilePromise.Create(predicate, timing, cancellationToken, out var token), token);
         }
 
+        public static UniTask WaitUntilCanceled(CancellationToken cancellationToken, PlayerLoopTiming timing = PlayerLoopTiming.Update)
+        {
+            return new UniTask(WaitUntilCanceledPromise.Create(cancellationToken, timing, out var token), token);
+        }
+
         public static UniTask<U> WaitUntilValueChanged<T, U>(T target, Func<T, U> monitorFunction, PlayerLoopTiming monitorTiming = PlayerLoopTiming.Update, IEqualityComparer<U> equalityComparer = null, CancellationToken cancellationToken = default(CancellationToken))
           where T : class
         {
@@ -234,6 +239,91 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
+        sealed class WaitUntilCanceledPromise : IUniTaskSource, IPlayerLoopItem, IPromisePoolItem
+        {
+            static readonly PromisePool<WaitUntilCanceledPromise> pool = new PromisePool<WaitUntilCanceledPromise>();
+
+            CancellationToken cancellationToken;
+
+            UniTaskCompletionSourceCore<object> core;
+
+            WaitUntilCanceledPromise()
+            {
+            }
+
+            public static IUniTaskSource Create(CancellationToken cancellationToken, PlayerLoopTiming timing, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                var result = pool.TryRent() ?? new WaitUntilCanceledPromise();
+
+                result.cancellationToken = cancellationToken;
+
+                TaskTracker.TrackActiveTask(result, 3);
+
+                PlayerLoopHelper.AddAction(timing, result);
+
+                token = result.core.Version;
+                return result;
+            }
+
+            public void GetResult(short token)
+            {
+                try
+                {
+                    TaskTracker.RemoveTracking(this);
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    pool.TryReturn(this);
+                }
+            }
+
+            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.TrySetResult(null);
+                    return false;
+                }
+
+                return true;
+            }
+
+            public void Reset()
+            {
+                core.Reset();
+                cancellationToken = default;
+            }
+
+            ~WaitUntilCanceledPromise()
+            {
+                if (pool.TryReturn(this))
+                {
+                    GC.ReRegisterForFinalize(this);
+                }
+            }
+        }
+
         // where T : UnityEngine.Object, can not add constraint
         sealed class WaitUntilValueChangedUnityObjectPromise<T, U> : IUniTaskSource<U>, IPlayerLoopItem, IPromisePoolItem
         {
@@ -353,7 +443,6 @@ namespace Cysharp.Threading.Tasks
             }
         }
 
-
         sealed class WaitUntilValueChangedStandardObjectPromise<T, U> : IUniTaskSource<U>, IPlayerLoopItem, IPromisePoolItem
             where T : class
         {
diff --git a/src/UniTask/Assets/Scenes/SandboxMain.cs b/src/UniTask/Assets/Scenes/SandboxMain.cs
index 6ba3ae5..dc42d22 100644
--- a/src/UniTask/Assets/Scenes/SandboxMain.cs
+++ b/src/UniTask/Assets/Scenes/SandboxMain.cs
@@ -203,16 +203,19 @@ public class SandboxMain : MonoBehaviour
         //await UniTaskAsyncEnumerable.EveryUpdate().Take(10).ForEachAsync((x, i) => rp.Value = i);
 
         //rp.Dispose();
+        var cts = new CancellationTokenSource();
+        Running(cts.Token).Forget();
 
-        await this.GetAsyncUpdateTrigger().ForEachAsync(x => Debug.Log("yeah"));
-        Debug.Log("DONE");
+
+
+        cts.CancelAfterSlim(TimeSpan.FromSeconds(3));
 
     }
 
-    async UniTaskVoid Running(IReadOnlyAsyncReactiveProperty<int> rp)
+    async UniTaskVoid Running(CancellationToken ct)
     {
         Debug.Log("BEGIN");
-        await rp.ForEachAsync(x => Debug.Log("AR:" + x));
+        await UniTask.WaitUntilCanceled(ct);
         Debug.Log("DONE");
     }