Add TransitionPromise task source.

Author: ChebanovDD

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/Interfaces/ITransitionPredicate.cs

@@ -7,7 +7,6 @@ namespace UnityMvvmToolkit.UniTask.Interfaces
     public interface ITransitionPredicate
     {
         bool TransitionEnd(TransitionEndEvent e);
-        bool TransitionCancel(TransitionCancelEvent e);
     }
 }
 

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/TransitionAsyncExtensions.cs

@@ -16,34 +16,31 @@ public static class TransitionAsyncExtensions
         private const char Dash = '-';
         private const int DefaultTimeoutMs = 2500;
 
-        public static async UniTask WaitForLongestTransitionEnd(this VisualElement element,
+        public static UniTask WaitForLongestTransitionEnd(this VisualElement element,
             CancellationToken cancellationToken = default)
         {
             var maxTransitionDelay = element.resolvedStyle.transitionDuration.Max(duration =>
                 duration.unit == TimeUnit.Millisecond ? duration.value : duration.value * 1000);
 
-            if (maxTransitionDelay > float.Epsilon)
-            {
-                await UniTask.Delay((int) maxTransitionDelay, cancellationToken: cancellationToken);
-            }
+            return maxTransitionDelay > float.Epsilon
+                ? UniTask.Delay((int) maxTransitionDelay, cancellationToken: cancellationToken)
+                : UniTask.CompletedTask;
         }
 
-        public static async UniTask WaitForTransitionEnd(this VisualElement element, int transitionIndex,
+        public static UniTask WaitForTransitionEnd(this VisualElement element, int transitionIndex,
             int timeoutMs = DefaultTimeoutMs, CancellationToken cancellationToken = default)
         {
-            if (GetTransitionDuration(element, transitionIndex, out var stylePropertyName) > float.Epsilon)
-            {
-                await WaitForTransition(element, stylePropertyName, timeoutMs, cancellationToken);
-            }
+            return GetTransitionDuration(element, transitionIndex, out var stylePropertyName) > float.Epsilon
+                ? WaitForTransition(element, stylePropertyName, timeoutMs, cancellationToken)
+                : UniTask.CompletedTask;
         }
 
-        public static async UniTask WaitForTransitionEnd(this VisualElement element, StylePropertyName stylePropertyName,
+        public static UniTask WaitForTransitionEnd(this VisualElement element, StylePropertyName stylePropertyName,
             int timeoutMs = DefaultTimeoutMs, CancellationToken cancellationToken = default)
         {
-            if (GetTransitionDuration(element, stylePropertyName) > float.Epsilon)
-            {
-                await WaitForTransition(element, stylePropertyName, timeoutMs, cancellationToken);
-            }
+            return GetTransitionDuration(element, stylePropertyName) > float.Epsilon
+                ? WaitForTransition(element, stylePropertyName, timeoutMs, cancellationToken)
+                : UniTask.CompletedTask;
         }
 
         /// <summary>
@@ -54,79 +51,42 @@ public static async UniTask WaitForTransitionEnd(this VisualElement element, Sty
         /// <param name="timeoutMs"></param>
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
-        public static async UniTask WaitForTransitionEnd(this VisualElement element, string propertyName,
+        public static UniTask WaitForTransitionEnd(this VisualElement element, string propertyName,
             int timeoutMs = DefaultTimeoutMs, CancellationToken cancellationToken = default)
         {
             var transitionData = GetTransitionData(element, propertyName);
-            if (transitionData is { Duration: { value: > float.Epsilon } })
-            {
-                await WaitForTransition(element, transitionData.Value.StylePropertyName, timeoutMs, cancellationToken);
-            }
+
+            return transitionData is { Duration: { value: > float.Epsilon } }
+                ? WaitForTransition(element, transitionData.Value.StylePropertyName, timeoutMs, cancellationToken)
+                : UniTask.CompletedTask;
         }
 
-        public static async UniTask WaitForAnyTransitionEnd(this VisualElement element, int timeoutMs = DefaultTimeoutMs,
+        public static UniTask WaitForAnyTransitionEnd(this VisualElement element, int timeoutMs = DefaultTimeoutMs,
             CancellationToken cancellationToken = default)
         {
-            if (AnyTransitionHasDuration(element))
-            {
-                await WaitForTransitionEnd(element, new TransitionAnyPredicate(), timeoutMs, cancellationToken);
-            }
+            return AnyTransitionHasDuration(element)
+                ? WaitForTransitionEnd(element, new TransitionAnyPredicate(), timeoutMs, cancellationToken)
+                : UniTask.CompletedTask;
         }
 
-        public static async UniTask WaitForAllTransitionsEnd(this VisualElement element, int timeoutMs = DefaultTimeoutMs,
+        public static UniTask WaitForAllTransitionsEnd(this VisualElement element, int timeoutMs = DefaultTimeoutMs,
             CancellationToken cancellationToken = default)
         {
             var transitionsCount = GetTransitionsWithDurationCount(element);
-            if (transitionsCount == 0)
-            {
-                return;
-            }
 
-            await WaitForTransitionEnd(element, new TransitionCounterPredicate(transitionsCount), timeoutMs,
-                cancellationToken);
+            return transitionsCount == 0
+                ? UniTask.CompletedTask
+                : WaitForTransitionEnd(element, new TransitionCounterPredicate(transitionsCount), timeoutMs,
+                    cancellationToken);
         }
 
-        public static async UniTask WaitForTransitionEnd<T>(this VisualElement element, T transitionPredicate,
+        public static UniTask WaitForTransitionEnd<T>(this VisualElement element, T transitionPredicate,
             int timeoutMs = DefaultTimeoutMs, CancellationToken cancellationToken = default)
             where T : ITransitionPredicate
         {
-            var taskCompletionSource = new UniTaskCompletionSource();
-
-            void OnTransitionEnd(TransitionEndEvent e)
-            {
-                if (transitionPredicate.TransitionEnd(e))
-                {
-                    taskCompletionSource.TrySetResult();
-                }
-            }
-
-            void OnTransitionCancel(TransitionCancelEvent e)
-            {
-                if (transitionPredicate.TransitionCancel(e))
-                {
-                    taskCompletionSource.TrySetCanceled();
-                }
-            }
-
-            try
-            {
-                element.RegisterCallback<TransitionCancelEvent>(OnTransitionCancel);
-                element.RegisterCallback<TransitionEndEvent>(OnTransitionEnd);
-
-                var transitionTask = taskCompletionSource.Task;
-                var timeoutTask = UniTask.Delay(timeoutMs, cancellationToken: cancellationToken);
-
-                var task = await UniTask.WhenAny(transitionTask, timeoutTask).SuppressCancellationThrow();
-                if (task.IsCanceled || task.Result == 1)
-                {
-                    taskCompletionSource.TrySetCanceled();
-                }
-            }
-            finally
-            {
-                element.UnregisterCallback<TransitionCancelEvent>(OnTransitionCancel);
-                element.UnregisterCallback<TransitionEndEvent>(OnTransitionEnd);
-            }
+            return new UniTask(
+                TransitionPromise<T>.Create(element, transitionPredicate, timeoutMs, PlayerLoopTiming.Update,
+                    cancellationToken, out var token), token);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/TransitionPredicates/TransitionAnyPredicate.cs

@@ -8,7 +8,6 @@ namespace UnityMvvmToolkit.UniTask.TransitionPredicates
     public readonly struct TransitionAnyPredicate : ITransitionPredicate
     {
         public bool TransitionEnd(TransitionEndEvent e) => true;
-        public bool TransitionCancel(TransitionCancelEvent e) => true;
     }
 }
 

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/TransitionPredicates/TransitionCounterPredicate.cs

@@ -15,7 +15,6 @@ public TransitionCounterPredicate(int transitionsCount)
         }
 
         public bool TransitionEnd(TransitionEndEvent e) => --_transitionsCount == 0;
-        public bool TransitionCancel(TransitionCancelEvent e) => true;
     }
 }
 

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/TransitionPredicates/TransitionNamePredicate.cs

@@ -15,7 +15,6 @@ public TransitionNamePredicate(StylePropertyName stylePropertyName)
         }
 
         public bool TransitionEnd(TransitionEndEvent e) => e.AffectsProperty(_stylePropertyName);
-        public bool TransitionCancel(TransitionCancelEvent e) => true;
     }
 }
 

src/UnityMvvmToolkit.UnityPackage/Assets/Plugins/UnityMvvmToolkit/Runtime/External/UniTask/TransitionPromise.cs

@@ -0,0 +1,191 @@
+#if UNITYMVVMTOOLKIT_UNITASK_SUPPORT
+
+namespace UnityMvvmToolkit.UniTask
+{
+    using System;
+    using Interfaces;
+    using System.Threading;
+    using Cysharp.Threading.Tasks;
+    using UnityEngine;
+    using UnityEngine.UIElements;
+
+    internal sealed class TransitionPromise<T> : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<TransitionPromise<T>>
+        where T : ITransitionPredicate
+    {
+        private static TaskPool<TransitionPromise<T>> _pool;
+
+        private TransitionPromise<T> _nextNode;
+        private UniTaskCompletionSourceCore<AsyncUnit> _core;
+
+        private int _initFrame;
+        private float _elapsed;
+        private float _delayTime;
+
+        private T _predicate;
+        private bool _completed;
+        private VisualElement _visualElement;
+        private CancellationToken _cancellationToken;
+
+        private readonly EventCallback<TransitionEndEvent> _endCallback;
+        private readonly EventCallback<TransitionCancelEvent> _cancelCallback;
+
+        private TransitionPromise()
+        {
+            _endCallback = OnTransitionEnd;
+            _cancelCallback = OnTransitionCancel;
+        }
+
+        public ref TransitionPromise<T> NextNode => ref _nextNode;
+
+        static TransitionPromise()
+        {
+            TaskPool.RegisterSizeGetter(typeof(TransitionPromise<T>), () => _pool.Size);
+        }
+
+        public static IUniTaskSource Create(VisualElement visualElement, T predicate, int timeoutMs,
+            PlayerLoopTiming timing, CancellationToken cancellationToken, out short token)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+            }
+
+            if (_pool.TryPop(out var transitionPromise) == false)
+            {
+                transitionPromise = new TransitionPromise<T>();
+            }
+
+            transitionPromise._elapsed = 0.0f;
+            transitionPromise._delayTime = (float) TimeSpan.FromMilliseconds(timeoutMs).TotalSeconds;
+            transitionPromise._initFrame = PlayerLoopHelper.IsMainThread ? Time.frameCount : -1;
+
+            transitionPromise._completed = false;
+            transitionPromise._predicate = predicate;
+            transitionPromise._visualElement = visualElement;
+            transitionPromise._cancellationToken = cancellationToken;
+
+            TaskTracker.TrackActiveTask(transitionPromise, 3);
+            PlayerLoopHelper.AddAction(timing, transitionPromise);
+
+            visualElement.RegisterCallback(transitionPromise._cancelCallback);
+            visualElement.RegisterCallback(transitionPromise._endCallback);
+
+            token = transitionPromise._core.Version;
+
+            return transitionPromise;
+        }
+
+        private void OnTransitionEnd(TransitionEndEvent e)
+        {
+            UnregisterCallbacks();
+
+            if (_completed)
+            {
+                Reset();
+                return;
+            }
+
+            _completed = true;
+
+            if (_cancellationToken.IsCancellationRequested)
+            {
+                _core.TrySetCanceled(_cancellationToken);
+            }
+            else if (_predicate.TransitionEnd(e))
+            {
+                _core.TrySetResult(AsyncUnit.Default);
+            }
+        }
+
+        private void OnTransitionCancel(TransitionCancelEvent e)
+        {
+            UnregisterCallbacks();
+
+            if (_completed)
+            {
+                Reset();
+                return;
+            }
+
+            _completed = true;
+            _core.TrySetCanceled(_cancellationToken.IsCancellationRequested ? _cancellationToken : default);
+        }
+
+        private void UnregisterCallbacks()
+        {
+            _visualElement.UnregisterCallback(_cancelCallback);
+            _visualElement.UnregisterCallback(_endCallback);
+        }
+
+        public void GetResult(short token)
+        {
+            _core.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 (_completed)
+            {
+                Reset();
+                return false;
+            }
+
+            if (_cancellationToken.IsCancellationRequested)
+            {
+                _completed = true;
+                _core.TrySetCanceled(_cancellationToken);
+
+                return false;
+            }
+
+            if (_elapsed == 0.0f && _initFrame == Time.frameCount)
+            {
+                return true;
+            }
+
+            _elapsed += Time.deltaTime;
+
+            if (_elapsed < _delayTime)
+            {
+                return true;
+            }
+
+            _core.TrySetResult(AsyncUnit.Default);
+
+            return false;
+        }
+
+        private void Reset()
+        {
+            TaskTracker.RemoveTracking(this);
+
+            _core.Reset();
+
+            _elapsed = default;
+            _delayTime = default;
+
+            _predicate = default;
+            _visualElement = default;
+            _cancellationToken = default;
+
+            _pool.TryPush(this);
+        }
+    }
+}
+
+#endif