TE6 打飞机Demo

TE6 打飞机Demo

UnityProject/Packages/UniTask/Runtime/UniTask.WaitUntil.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.Tracing;
 using System.Threading;
 using Cysharp.Threading.Tasks.Internal;
 
@@ -14,11 +15,21 @@ namespace Cysharp.Threading.Tasks
             return new UniTask(WaitUntilPromise.Create(predicate, timing, cancellationToken, cancelImmediately, out var token), token);
         }
 
+        public static UniTask WaitUntil<T>(T state, Func<T, bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
+        {
+            return new UniTask(WaitUntilPromise<T>.Create(state, predicate, timing, cancellationToken, cancelImmediately, out var token), token);
+        }
+
         public static UniTask WaitWhile(Func<bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
         {
             return new UniTask(WaitWhilePromise.Create(predicate, timing, cancellationToken, cancelImmediately, out var token), token);
         }
 
+        public static UniTask WaitWhile<T>(T state, Func<T, bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
+        {
+            return new UniTask(WaitWhilePromise<T>.Create(state, predicate, timing, cancellationToken, cancelImmediately, out var token), token);
+        }
+
         public static UniTask WaitUntilCanceled(CancellationToken cancellationToken, PlayerLoopTiming timing = PlayerLoopTiming.Update, bool completeImmediately = false)
         {
             return new UniTask(WaitUntilCanceledPromise.Create(cancellationToken, timing, completeImmediately, out var token), token);
@@ -49,6 +60,7 @@ namespace Cysharp.Threading.Tasks
             Func<bool> predicate;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -70,6 +82,7 @@ namespace Cysharp.Threading.Tasks
 
                 result.predicate = predicate;
                 result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
 
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
@@ -96,7 +109,14 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(this);
+                    }
                 }
             }
 
@@ -147,6 +167,136 @@ namespace Cysharp.Threading.Tasks
                 predicate = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
+                return pool.TryPush(this);
+            }
+        }
+
+        sealed class WaitUntilPromise<T> : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitUntilPromise<T>>
+        {
+            static TaskPool<WaitUntilPromise<T>> pool;
+            WaitUntilPromise<T> nextNode;
+            public ref WaitUntilPromise<T> NextNode => ref nextNode;
+
+            static WaitUntilPromise()
+            {
+                TaskPool.RegisterSizeGetter(typeof(WaitUntilPromise<T>), () => pool.Size);
+            }
+
+            Func<T, bool> predicate;
+            T argument;
+            CancellationToken cancellationToken;
+            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
+
+            UniTaskCompletionSourceCore<object> core;
+
+            WaitUntilPromise()
+            {
+            }
+
+            public static IUniTaskSource Create(T argument, Func<T, bool> predicate, PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                if (!pool.TryPop(out var result))
+                {
+                    result = new WaitUntilPromise<T>();
+                }
+
+                result.predicate = predicate;
+                result.argument = argument;
+                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
+
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var promise = (WaitUntilPromise<T>)state;
+                        promise.core.TrySetCanceled(promise.cancellationToken);
+                    }, result);
+                }
+
+                TaskTracker.TrackActiveTask(result, 3);
+
+                PlayerLoopHelper.AddAction(timing, result);
+
+                token = result.core.Version;
+                return result;
+            }
+
+            public void GetResult(short token)
+            {
+                try
+                {
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(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.TrySetCanceled(cancellationToken);
+                    return false;
+                }
+
+                try
+                {
+                    if (!predicate(argument))
+                    {
+                        return true;
+                    }
+                }
+                catch (Exception ex)
+                {
+                    core.TrySetException(ex);
+                    return false;
+                }
+
+                core.TrySetResult(null);
+                return false;
+            }
+
+            bool TryReturn()
+            {
+                TaskTracker.RemoveTracking(this);
+                core.Reset();
+                predicate = default;
+                argument = default;
+                cancellationToken = default;
+                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -165,6 +315,7 @@ namespace Cysharp.Threading.Tasks
             Func<bool> predicate;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -186,7 +337,8 @@ namespace Cysharp.Threading.Tasks
 
                 result.predicate = predicate;
                 result.cancellationToken = cancellationToken;
-                
+                result.cancelImmediately = cancelImmediately;
+
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@@ -212,7 +364,14 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(this);
+                    }
                 }
             }
 
@@ -263,6 +422,136 @@ namespace Cysharp.Threading.Tasks
                 predicate = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
+                return pool.TryPush(this);
+            }
+        }
+
+        sealed class WaitWhilePromise<T> : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitWhilePromise<T>>
+        {
+            static TaskPool<WaitWhilePromise<T>> pool;
+            WaitWhilePromise<T> nextNode;
+            public ref WaitWhilePromise<T> NextNode => ref nextNode;
+
+            static WaitWhilePromise()
+            {
+                TaskPool.RegisterSizeGetter(typeof(WaitWhilePromise<T>), () => pool.Size);
+            }
+
+            Func<T, bool> predicate;
+            T argument;
+            CancellationToken cancellationToken;
+            CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
+
+            UniTaskCompletionSourceCore<object> core;
+
+            WaitWhilePromise()
+            {
+            }
+
+            public static IUniTaskSource Create(T argument, Func<T, bool> predicate, PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+                }
+
+                if (!pool.TryPop(out var result))
+                {
+                    result = new WaitWhilePromise<T>();
+                }
+
+                result.predicate = predicate;
+                result.argument = argument;
+                result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
+
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
+                {
+                    result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
+                    {
+                        var promise = (WaitWhilePromise<T>)state;
+                        promise.core.TrySetCanceled(promise.cancellationToken);
+                    }, result);
+                }
+
+                TaskTracker.TrackActiveTask(result, 3);
+
+                PlayerLoopHelper.AddAction(timing, result);
+
+                token = result.core.Version;
+                return result;
+            }
+
+            public void GetResult(short token)
+            {
+                try
+                {
+                    core.GetResult(token);
+                }
+                finally
+                {
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(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.TrySetCanceled(cancellationToken);
+                    return false;
+                }
+
+                try
+                {
+                    if (predicate(argument))
+                    {
+                        return true;
+                    }
+                }
+                catch (Exception ex)
+                {
+                    core.TrySetException(ex);
+                    return false;
+                }
+
+                core.TrySetResult(null);
+                return false;
+            }
+
+            bool TryReturn()
+            {
+                TaskTracker.RemoveTracking(this);
+                core.Reset();
+                predicate = default;
+                argument = default;
+                cancellationToken = default;
+                cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -280,6 +569,7 @@ namespace Cysharp.Threading.Tasks
 
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -287,7 +577,7 @@ namespace Cysharp.Threading.Tasks
             {
             }
 
-            public static IUniTaskSource Create(CancellationToken cancellationToken, PlayerLoopTiming timing, bool completeImmediately, out short token)
+            public static IUniTaskSource Create(CancellationToken cancellationToken, PlayerLoopTiming timing, bool cancelImmediately, out short token)
             {
                 if (cancellationToken.IsCancellationRequested)
                 {
@@ -300,8 +590,9 @@ namespace Cysharp.Threading.Tasks
                 }
 
                 result.cancellationToken = cancellationToken;
+                result.cancelImmediately = cancelImmediately;
 
-                if (completeImmediately && cancellationToken.CanBeCanceled)
+                if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
                     {
@@ -326,7 +617,14 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(this);
+                    }
                 }
             }
 
@@ -362,6 +660,7 @@ namespace Cysharp.Threading.Tasks
                 core.Reset();
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -385,6 +684,7 @@ namespace Cysharp.Threading.Tasks
             IEqualityComparer<U> equalityComparer;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<U> core;
 
@@ -410,7 +710,8 @@ namespace Cysharp.Threading.Tasks
                 result.currentValue = monitorFunction(target);
                 result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
                 result.cancellationToken = cancellationToken;
-                
+                result.cancelImmediately = cancelImmediately;
+
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@@ -436,7 +737,14 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(this);
+                    }
                 }
             }
 
@@ -497,6 +805,7 @@ namespace Cysharp.Threading.Tasks
                 equalityComparer = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }
@@ -519,6 +828,7 @@ namespace Cysharp.Threading.Tasks
             IEqualityComparer<U> equalityComparer;
             CancellationToken cancellationToken;
             CancellationTokenRegistration cancellationTokenRegistration;
+            bool cancelImmediately;
 
             UniTaskCompletionSourceCore<U> core;
 
@@ -543,7 +853,8 @@ namespace Cysharp.Threading.Tasks
                 result.currentValue = monitorFunction(target);
                 result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
                 result.cancellationToken = cancellationToken;
-                
+                result.cancelImmediately = cancelImmediately;
+
                 if (cancelImmediately && cancellationToken.CanBeCanceled)
                 {
                     result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@@ -569,7 +880,14 @@ namespace Cysharp.Threading.Tasks
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
+                    {
+                        TryReturn();
+                    }
+                    else
+                    {
+                        TaskTracker.RemoveTracking(this);
+                    }
                 }
             }
 
@@ -630,6 +948,7 @@ namespace Cysharp.Threading.Tasks
                 equalityComparer = default;
                 cancellationToken = default;
                 cancellationTokenRegistration.Dispose();
+                cancelImmediately = default;
                 return pool.TryPush(this);
             }
         }