Merge pull request #89 from AnnulusGames/feature-tovaluetask

Add: ToValueTask

src/LitMotion/Assets/LitMotion/Runtime/Internal/AwaitableMotionConfiguredSource.cs → src/LitMotion/Assets/LitMotion/Runtime/Internal/Tasks/AwaitableMotionConfiguredSource.cs

@@ -1,72 +1,14 @@
+#if UNITY_2023_1_OR_NEWER
 using System;
-using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using System.Threading;
 using UnityEngine;
 
 namespace LitMotion
 {
-#if UNITY_2023_1_OR_NEWER
-
-    internal sealed class AwaitableMotionConfiguredSource
+    internal sealed class AwaitableMotionConfiguredSource : IMotionTaskSourcePoolNode<AwaitableMotionConfiguredSource>
     {
-        [StructLayout(LayoutKind.Auto)]
-        public struct Pool
-        {
-            static readonly int MaxPoolSize = int.MaxValue;
-
-            int gate;
-            int size;
-            AwaitableMotionConfiguredSource root;
-
-            public readonly int Size => size;
-
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            public bool TryPop(out AwaitableMotionConfiguredSource result)
-            {
-                if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
-                {
-                    var v = root;
-                    if (v != null)
-                    {
-                        ref var nextNode = ref v.NextNode;
-                        root = nextNode;
-                        nextNode = null;
-                        size--;
-                        result = v;
-                        Volatile.Write(ref gate, 0);
-                        return true;
-                    }
-
-                    Volatile.Write(ref gate, 0);
-                }
-                result = default;
-                return false;
-            }
+        static MotionTaskSourcePool<AwaitableMotionConfiguredSource> pool;
 
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            public bool TryPush(AwaitableMotionConfiguredSource item)
-            {
-                if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
-                {
-                    if (size < MaxPoolSize)
-                    {
-                        item.NextNode = root;
-                        root = item;
-                        size++;
-                        Volatile.Write(ref gate, 0);
-                        return true;
-                    }
-                    else
-                    {
-                        Volatile.Write(ref gate, 0);
-                    }
-                }
-                return false;
-            }
-        }
-
-        static Pool pool;
         AwaitableMotionConfiguredSource nextNode;
         public ref AwaitableMotionConfiguredSource NextNode => ref nextNode;
 
@@ -215,6 +157,6 @@ void RestoreOriginalCallback()
             MotionStorageManager.SetMotionCallbacks(motionHandle, callbacks);
         }
     }
+}
 
 #endif
-}

src/LitMotion/Assets/LitMotion/Runtime/Internal/Tasks/IMotionTaskSourcePoolNode.cs

@@ -0,0 +1,7 @@
+namespace LitMotion
+{
+    internal interface IMotionTaskSourcePoolNode<T> where T : class
+    {
+        ref T NextNode { get; }
+    }
+}

src/LitMotion/Assets/LitMotion/Runtime/Internal/Tasks/TaskSourcePool.cs

@@ -0,0 +1,65 @@
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using System.Threading;
+
+namespace LitMotion
+{
+    // This implementation is based on UniTask's TaskPool<T>
+    // Reference: https://github.com/Cysharp/UniTask/blob/64792b672d35e43b3412fc74861f8bdbf41e3a6f/src/UniTask/Assets/Plugins/UniTask/Runtime/TaskPool.cs
+
+    [StructLayout(LayoutKind.Auto)]
+    internal struct MotionTaskSourcePool<T> where T : class, IMotionTaskSourcePoolNode<T>
+    {
+        static readonly int MaxPoolSize = int.MaxValue;
+
+        int gate;
+        int size;
+        T root;
+
+        public readonly int Size => size;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool TryPop(out T result)
+        {
+            if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
+            {
+                var v = root;
+                if (v != null)
+                {
+                    ref var nextNode = ref v.NextNode;
+                    root = nextNode;
+                    nextNode = null;
+                    size--;
+                    result = v;
+                    Volatile.Write(ref gate, 0);
+                    return true;
+                }
+
+                Volatile.Write(ref gate, 0);
+            }
+            result = default;
+            return false;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool TryPush(T item)
+        {
+            if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
+            {
+                if (size < MaxPoolSize)
+                {
+                    item.NextNode = root;
+                    root = item;
+                    size++;
+                    Volatile.Write(ref gate, 0);
+                    return true;
+                }
+                else
+                {
+                    Volatile.Write(ref gate, 0);
+                }
+            }
+            return false;
+        }
+    }
+}

src/LitMotion/Assets/LitMotion/Runtime/Internal/Tasks/ValueTaskMotionConfiguredSource.cs

@@ -0,0 +1,164 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks.Sources;
+
+namespace LitMotion
+{
+    internal sealed class ValueTaskMotionConfiguredSource : IValueTaskSource, IMotionTaskSourcePoolNode<ValueTaskMotionConfiguredSource>
+    {
+        static MotionTaskSourcePool<ValueTaskMotionConfiguredSource> pool;
+
+        ValueTaskMotionConfiguredSource nextNode;
+        public ref ValueTaskMotionConfiguredSource NextNode => ref nextNode;
+
+        readonly Action onCancelCallbackDelegate;
+        readonly Action onCompleteCallbackDelegate;
+
+        MotionHandle motionHandle;
+        CancellationToken cancellationToken;
+        CancellationTokenRegistration cancellationRegistration;
+
+        Action originalCompleteAction;
+        Action originalCancelAction;
+        ManualResetValueTaskSourceCore<object> core;
+
+        static ValueTaskMotionConfiguredSource FromCanceled(out short token)
+        {
+            if (canceledSource == null)
+            {
+                canceledSource = new();
+                canceledSource.core.SetException(new OperationCanceledException());
+            }
+
+            token = canceledSource.Version;
+            return canceledSource;
+        }
+        static ValueTaskMotionConfiguredSource canceledSource;
+
+        ValueTaskMotionConfiguredSource()
+        {
+            onCancelCallbackDelegate = new(OnCancelCallbackDelegate);
+            onCompleteCallbackDelegate = new(OnCompleteCallbackDelegate);
+        }
+
+        public static IValueTaskSource Create(MotionHandle motionHandle, CancellationToken cancellationToken, out short token)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                motionHandle.Cancel();
+                return FromCanceled(out token);
+            }
+
+            if (!pool.TryPop(out var result))
+            {
+                result = new ValueTaskMotionConfiguredSource();
+            }
+
+            result.motionHandle = motionHandle;
+            result.cancellationToken = cancellationToken;
+
+            var callbacks = MotionStorageManager.GetMotionCallbacks(motionHandle);
+            result.originalCancelAction = callbacks.OnCancelAction;
+            result.originalCompleteAction = callbacks.OnCompleteAction;
+            callbacks.OnCancelAction = result.onCancelCallbackDelegate;
+            callbacks.OnCompleteAction = result.onCompleteCallbackDelegate;
+            MotionStorageManager.SetMotionCallbacks(motionHandle, callbacks);
+
+            if (result.originalCancelAction == result.onCancelCallbackDelegate)
+            {
+                result.originalCancelAction = null;
+            }
+            if (result.originalCompleteAction == result.onCompleteCallbackDelegate)
+            {
+                result.originalCompleteAction = null;
+            }
+
+            if (cancellationToken.CanBeCanceled)
+            {
+                result.cancellationRegistration = cancellationToken.Register(static x =>
+                {
+                    var source = (ValueTaskMotionConfiguredSource)x;
+                    var motionHandle = source.motionHandle;
+                    if (motionHandle.IsActive())
+                    {
+                        motionHandle.Cancel();
+                    }
+                    else
+                    {
+                        source.core.SetException(new OperationCanceledException());
+                    }
+                }, result);
+            }
+
+            token = result.core.Version;
+            return result;
+        }
+
+        public short Version => core.Version;
+
+        void OnCompleteCallbackDelegate()
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                core.SetException(new OperationCanceledException());
+            }
+            else
+            {
+                originalCompleteAction?.Invoke();
+                core.SetResult(null);
+            }
+        }
+
+        void OnCancelCallbackDelegate()
+        {
+            originalCancelAction?.Invoke();
+            core.SetException(new OperationCanceledException());
+        }
+
+        public void GetResult(short token)
+        {
+            try
+            {
+                core.GetResult(token);
+            }
+            finally
+            {
+                TryReturn();
+            }
+        }
+
+        public ValueTaskSourceStatus GetStatus(short token)
+        {
+            return core.GetStatus(token);
+        }
+
+        public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
+        {
+            core.OnCompleted(continuation, state, token, flags);
+        }
+
+        bool TryReturn()
+        {
+            core.Reset();
+            cancellationRegistration.Dispose();
+
+            RestoreOriginalCallback();
+
+            motionHandle = default;
+            cancellationToken = default;
+            originalCompleteAction = default;
+            originalCancelAction = default;
+
+            return pool.TryPush(this);
+        }
+
+        void RestoreOriginalCallback()
+        {
+            if (!motionHandle.IsActive()) return;
+            var callbacks = MotionStorageManager.GetMotionCallbacks(motionHandle);
+            callbacks.OnCancelAction = originalCancelAction;
+            callbacks.OnCompleteAction = originalCompleteAction;
+            MotionStorageManager.SetMotionCallbacks(motionHandle, callbacks);
+        }
+    }
+}

src/LitMotion/Assets/LitMotion/Runtime/MotionHandleExtensions.cs

@@ -1,8 +1,7 @@
 using System;
 using System.Collections;
-using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using System.Threading;
+using System.Threading.Tasks;
 using UnityEngine;
 
 namespace LitMotion
@@ -120,6 +119,13 @@ public static IEnumerator ToYieldInteraction(this MotionHandle handle)
                 yield return null;
             }
         }
+
+        public static ValueTask ToValueTask(this MotionHandle handle, CancellationToken cancellationToken = default)
+        {
+            if (!handle.IsActive()) return default;
+            var source = ValueTaskMotionConfiguredSource.Create(handle, cancellationToken, out var token);
+            return new ValueTask(source, token);
+        }
 
 #if UNITY_2023_1_OR_NEWER
         public static Awaitable ToAwaitable(this MotionHandle handle, CancellationToken cancellationToken = default)