给PrimeTween新增小小的异步代码支持

最近做的新项目没有使用 DOTween 了，转而使用性能比较高的 PrimeTween。这是一个比较新的插件，看了下大概是前两年上架的。
比起历史悠久的 DOTween 相比， PrimeTween 虽然性能很高，但是不像 DOTween 有丰富的特性以及一些额外的其它插件的联动支持。

具体的就说 UniTask 吧，他有一个专门的拓展类给予 DOTwee 异步相关的支持。这里面有我现在马上就需要用到的 “使用 CancellationToken 取消对 Tween 的异步等待” 的功能。

PrimeTween 没有这个支持，没关系，我只要照着这个拓展类抄一个就好，虽然这俩Tween的实现不太一样，但是核心的行为还是非常相似的。

思路

稍微点F12跳两下就发现，其实这个功能的实现思路就是通过方法 CancellationToken.Register 去给C ancellationToken 收到取消请求的时机注册一个回调。在这个回调中去处理关于Tween的各种处理。

DOTween 的拓展处理中会把原本 Tween 的正常播放完成相关回调拿出来然后缓存在自己的结构中并换上自己的响应动画的正常完成。PrimeTween 和 DOTween 不一样的地方在于 PrimeTween 的回调不能重复绑定，所以照着抄的时候为了保证代码的一致性，我选择用 PrimeTween 插件中的 Sequence 把原本的 Tween/Sequence 包装起来，似乎这个应该是他插件本身就推荐的做法。当然包起来以后就不可以再编辑原 Tween 的一些参数了。

PrimeTween 高性能的原因似乎是他把实际执行 Tween 行为的实例都用对象池管理起来了，不像 DOTween 需要用户自己管理动画实例对象。而且 DOTween 的用于存放相关动画参数的数据结构 TweenSetting 好像也是类，这点其实让我有点讨厌。

代码

namespace Cysharp.Threading.Tasks
{
    public enum TweenCancelBehaviour
    {
        KillAndCancelAwait,
        // TODO @Hiko finish those cancel behaviour if really need :)
        // Kill,
        //KillWithCompleteCallback,
    }

    public  static class PrimeTweenAsyncExtensions
    {
        public struct TweenAwaiter : ICriticalNotifyCompletion
        {
            private readonly Tween tween;
            private readonly Sequence warppedSequence;

            // killed(non active) as completed.
            public bool IsCompleted => !tween.isAlive;

            public TweenAwaiter(Tween tween)
            {
                this.tween = tween;
                warppedSequence = Sequence.Create(this.tween);
            }

            public TweenAwaiter GetAwaiter()
            {
                return this;
            }

            public void GetResult() { }

            public void OnCompleted(Action continuation)
            {
                UnsafeOnCompleted(continuation);
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                // onKill is called after OnCompleted, both Complete(false/true) and Kill(false/true).
                warppedSequence.OnComplete(PooledPrimeTweenCallback.Create(continuation));
            }
        }

        public static TweenAwaiter GetAwaiter(this Tween tween)
        {
            return new TweenAwaiter(tween);
        }

        public static UniTask WithCancellation(in this Tween tween, CancellationToken cancellationToken)
        {
            if (!tween.isAlive) return UniTask.CompletedTask;
            return new UniTask(TweenConfiguredSource.Create(tween, TweenCancelBehaviour.KillAndCancelAwait, cancellationToken, out var token, out _), token);
        }

        public static UniTask WithCancellation(in this Tween tween, CancellationToken cancellationToken, out Sequence warppedTween)
        {
            warppedTween = default;
            if (!tween.isAlive) return UniTask.CompletedTask;
            return new UniTask(TweenConfiguredSource.Create(tween, TweenCancelBehaviour.KillAndCancelAwait, cancellationToken, out var token, out warppedTween), token);
        }

        private sealed class TweenConfiguredSource : IUniTaskSource, ITaskPoolNode<TweenConfiguredSource>
        {
            private static TaskPool<TweenConfiguredSource> s_pool;

            private TweenConfiguredSource m_nextNode;

            private Tween m_tween;
            // use sequence to warp the original tween cuz original tween may already bind a completion callback
            private Sequence m_warppedSequence;

            private TweenCancelBehaviour m_cancelBehaviour;
            private CancellationToken m_cancellationToken;
            private CancellationTokenRegistration m_cancellationRegistration;
            bool m_isCanceled;

            private UniTaskCompletionSourceCore<AsyncUnit> m_core;

            /// <summary>
            /// event will be triggered when the tween completes normally
            /// </summary>
            private readonly Action m_tweenCompleteNotifyCallback;

            public ref TweenConfiguredSource NextNode => ref m_nextNode;

            static TweenConfiguredSource()
            {
                TaskPool.RegisterSizeGetter(typeof(TweenConfiguredSource), () => s_pool.Size);
            }

            private TweenConfiguredSource()
            {
                m_tweenCompleteNotifyCallback = OnTweenCompleteCallbackDelegate;
            }

            public static IUniTaskSource Create(Tween tween, TweenCancelBehaviour cancelBehaviour, CancellationToken cancellationToken, out short token)
            {
                return Create(tween, cancelBehaviour, cancellationToken, out token, out _);
            }

            public static IUniTaskSource Create(Tween tween, TweenCancelBehaviour cancelBehaviour, CancellationToken cancellationToken, out short token, out Sequence warppedTween)
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    DoCancelBeforeCreate(tween, cancelBehaviour);
                    warppedTween = default;
                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
                }

                if (!s_pool.TryPop(out var result))
                {
                    result = new TweenConfiguredSource();
                }

                result.m_tween = tween;
                result.m_cancelBehaviour = cancelBehaviour;
                result.m_cancellationToken = cancellationToken;
                result.m_isCanceled = false;
                // warp the original tween and bind our comlete callback
                result.m_warppedSequence = Sequence.Create(result.m_tween);
                result.m_warppedSequence.OnComplete(result.m_tweenCompleteNotifyCallback);

                if (cancellationToken.CanBeCanceled)
                {
                    result.m_cancellationRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(x =>
                    {
                        TweenConfiguredSource source = (TweenConfiguredSource)x;
                        switch (source.m_cancelBehaviour)
                        {
                            case TweenCancelBehaviour.KillAndCancelAwait:
                            default:
                                // token canceled stop tween
                                source.m_isCanceled = true;
                                source.m_warppedSequence.Stop();
                                source.OnTweenCompleteCallbackDelegate();
                                break;
                        }
                    }, result);
                }

                TaskTracker.TrackActiveTask(result, 3);
                token = result.m_core.Version;
                warppedTween = result.m_warppedSequence;
                return result;
            }

            private static void DoCancelBeforeCreate(Tween tween, TweenCancelBehaviour tweenCancelBehaviour)
            {
                switch (tweenCancelBehaviour)
                {
                    case TweenCancelBehaviour.KillAndCancelAwait:
                    default:
                        tween.Stop();
                        break;
                }
            }

            public UniTaskStatus UnsafeGetStatus()
            {
                return m_core.UnsafeGetStatus();
            }

            public void OnCompleted(Action<object> continuation, object state, short token)
            {
                m_core.OnCompleted(continuation, state, token);
            }

            public void GetResult(short token)
            {
                try
                {
                    m_core.GetResult(token);
                }
                finally
                {
                    TryReturnSelfToPool();
                }
            }

            public UniTaskStatus GetStatus(short token)
            {
                return m_core.GetStatus(token);
            }

            private bool TryReturnSelfToPool()
            {
                TaskTracker.RemoveTracking(this);
                m_core.Reset();
                m_cancellationRegistration.Dispose();

                m_tween = default;
                m_warppedSequence = default;
                m_cancellationToken = default;
                return s_pool.TryPush(this);
            }

            /// <summary>
            /// Tween completes normally
            /// </summary>
            private void OnTweenCompleteCallbackDelegate()
            {
                if (m_isCanceled)
                {
                    m_core.TrySetCanceled(m_cancellationToken);
                }
                else
                {
                    m_core.TrySetResult(AsyncUnit.Default);
                }
            }
        }

        //----------------------------------

        public struct SequenceAwaiter : ICriticalNotifyCompletion
        {
            private readonly Sequence originSequence;
            private readonly Sequence warppedSequence;

            // killed(non active) as completed.
            public bool IsCompleted => !originSequence.isAlive;

            public SequenceAwaiter(Sequence sequence)
            {
                originSequence = sequence;
                warppedSequence = Sequence.Create();
                warppedSequence.Insert(0f, sequence);
            }

            public SequenceAwaiter GetAwaiter()
            {
                return this;
            }

            public void GetResult() { }

            public void OnCompleted(Action continuation)
            {
                UnsafeOnCompleted(continuation);
            }

            public void UnsafeOnCompleted(Action continuation)
            {
                // onKill is called after OnCompleted, both Complete(false/true) and Kill(false/true).
                warppedSequence.OnComplete(PooledPrimeTweenCallback.Create(continuation));
            }
        }

        public static SequenceAwaiter GetAwaiter(this Sequence tween)
        {
            return new SequenceAwaiter(tween);
        }

        public static UniTask WithCancellation(in this Sequence sequence, CancellationToken cancellationToken)
        {
            if (!sequence.isAlive) return UniTask.CompletedTask;
            return new UniTask(SequenceConfiguredSource.Create(sequence, TweenCancelBehaviour.KillAndCancelAwait, cancellationToken, out var token, out _), token);
        }

        public static UniTask WithCancellation(in this Sequence sequence, CancellationToken cancellationToken, out Sequence warppedTween)
        {
            warppedTween = default;
            if (!sequence.isAlive) return UniTask.CompletedTask;
            return new UniTask(SequenceConfiguredSource.Create(sequence, TweenCancelBehaviour.KillAndCancelAwait, cancellationToken, out var token, out warppedTween), token);
        }

        private sealed class SequenceConfiguredSource : IUniTaskSource, ITaskPoolNode<SequenceConfiguredSource>
        {
            private static TaskPool<SequenceConfiguredSource> s_pool;

            private SequenceConfiguredSource m_nextNode;

            private Sequence m_originSequence;
            // use sequence to warp the original sequence cuz original tween may already bind a completion callback
            private Sequence m_warppedSequence;

            private TweenCancelBehaviour m_cancelBehaviour;
            private CancellationToken m_cancellationToken;
            private CancellationTokenRegistration m_cancellationRegistration;
            bool m_isCanceled;

            private UniTaskCompletionSourceCore<AsyncUnit> m_core;

            private readonly Action m_SequenceCompleteNotifyCallback;

            public ref SequenceConfiguredSource NextNode => ref m_nextNode;

            static SequenceConfiguredSource()
            {
                TaskPool.RegisterSizeGetter(typeof(SequenceConfiguredSource), () => s_pool.Size);
            }

            private SequenceConfiguredSource()
            {
                m_SequenceCompleteNotifyCallback = OnSequenceCompleteCallbackDelegate;
            }

            public static IUniTaskSource Create(Sequence Sequence, TweenCancelBehaviour cancelBehaviour, CancellationToken cancellationToken, out short token)
            {
                return Create(Sequence, cancelBehaviour, cancellationToken, out token, out _);
            }

            public static IUniTaskSource Create(Sequence Sequence, TweenCancelBehaviour cancelBehaviour, CancellationToken cancellationToken, out short token, out Sequence warppedSequence)
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    DoCancelBeforeCreate(Sequence, cancelBehaviour);
                    warppedSequence = default;
                    return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
                }

                if (!s_pool.TryPop(out var result))
                {
                    result = new SequenceConfiguredSource();
                }

                result.m_originSequence = Sequence;
                result.m_cancelBehaviour = cancelBehaviour;
                result.m_cancellationToken = cancellationToken;
                result.m_isCanceled = false;

                result.m_warppedSequence = Sequence.Create();
                result.m_warppedSequence.Insert(0f, result.m_originSequence);
                result.m_warppedSequence.OnComplete(result.m_SequenceCompleteNotifyCallback);

                if (cancellationToken.CanBeCanceled)
                {
                    result.m_cancellationRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(x =>
                    {
                        SequenceConfiguredSource source = (SequenceConfiguredSource)x;
                        switch (source.m_cancelBehaviour)
                        {
                            case TweenCancelBehaviour.KillAndCancelAwait:
                                source.m_isCanceled = true;
                                source.m_warppedSequence.Stop();
                                source.OnSequenceCompleteCallbackDelegate();
                                break;
                            default:
                                break;
                        }
                    }, result);
                }

                TaskTracker.TrackActiveTask(result, 3);
                token = result.m_core.Version;
                warppedSequence = result.m_warppedSequence;
                return result;
            }

            private static void DoCancelBeforeCreate(Sequence Sequence, TweenCancelBehaviour tweenCancelBehaviour)
            {
                switch (tweenCancelBehaviour)
                {
                    case TweenCancelBehaviour.KillAndCancelAwait:
                    default:
                        Sequence.Stop();
                        break;
                }
            }

            public UniTaskStatus UnsafeGetStatus()
            {
                return m_core.UnsafeGetStatus();
            }

            public void OnCompleted(Action<object> continuation, object state, short token)
            {
                m_core.OnCompleted(continuation, state, token);
            }

            public void GetResult(short token)
            {
                try
                {
                    m_core.GetResult(token);
                }
                finally
                {
                    TryReturnSelfToPool();
                }
            }

            public UniTaskStatus GetStatus(short token)
            {
                return m_core.GetStatus(token);
            }

            private bool TryReturnSelfToPool()
            {
                TaskTracker.RemoveTracking(this);
                m_core.Reset();
                m_cancellationRegistration.Dispose();

                m_originSequence = default;
                m_warppedSequence = default;
                m_cancellationToken = default;
                return s_pool.TryPush(this);
            }

            private void OnSequenceCompleteCallbackDelegate()
            {
                if (m_isCanceled)
                {
                    m_core.TrySetCanceled(m_cancellationToken);
                }
                else
                {
                    m_core.TrySetResult(AsyncUnit.Default);
                }
            }
        }

    }

    public sealed class PooledPrimeTweenCallback
    {
        private static readonly ConcurrentQueue<PooledPrimeTweenCallback> s_pool = new ConcurrentQueue<PooledPrimeTweenCallback>();

        private readonly Action m_runDelegate;
        private Action m_continuation;

        private PooledPrimeTweenCallback()
        {
            m_runDelegate = Run;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static Action Create(Action continuation)
        {
            if (!s_pool.TryDequeue(out PooledPrimeTweenCallback item))
            {
                item = new PooledPrimeTweenCallback();
            }

            item.m_continuation = continuation;
            return item.m_runDelegate;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private void Run()
        {
            Action call = m_continuation;
            m_continuation = null;
            if (call != null)
            {
                s_pool.Enqueue(this);
                call.Invoke();
            }
        }

    }
}