Create API to await multiple PlayerLoopTiming without any allocation

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/ContinuationQueue.cs

@@ -154,6 +154,7 @@ namespace Cysharp.Threading.Tasks.Internal
         [System.Diagnostics.DebuggerHidden]
         void RunCore()
         {
+            PlayerLoopHelper.SetCurrentPlayerLoopTiming(timing);
             {
                 bool lockTaken = false;
                 try

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/PlayerLoopRunner.cs

@@ -134,6 +134,7 @@ namespace Cysharp.Threading.Tasks.Internal
         [System.Diagnostics.DebuggerHidden]
         void RunCore()
         {
+            PlayerLoopHelper.SetCurrentPlayerLoopTiming(timing);
             lock (runningAndQueueLock)
             {
                 running = true;

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/UniTaskPlayerLoopSubSystem.cs

@@ -0,0 +1,325 @@
+using System;
+using System.Threading;
+using UnityEngine;
+
+#if UNITY_EDITOR
+using UnityEditor;
+#endif
+
+namespace Cysharp.Threading.Tasks.Internal
+{
+    internal sealed class UniTaskPlayerLoopSubSystem : IPlayerLoopItem
+    {
+        
+        public UniTaskPlayerLoopSubSystem()
+        {
+            this.unhandledExceptionCallback = ex => Debug.LogException(ex);
+#if UNITY_EDITOR
+            EditorApplication.playModeStateChanged += OnPlayModeStateChanged;
+#endif
+        }
+
+        const int MaxArrayLength = 0X7FEFFFFF;
+        const int InitialSize = 16;
+
+        void Run() 
+        {
+            RunContinuations();
+            RunLoopItems();
+        }
+        void Clear()
+        {
+            ClearContinuations();
+            ClearLoopItems();
+        }
+
+        public bool MoveNext()
+        {
+            Run();
+            return true;
+        }
+#if UNITY_EDITOR
+        void OnPlayModeStateChanged(PlayModeStateChange state)
+        {
+            if (state == PlayModeStateChange.EnteredEditMode || state == PlayModeStateChange.EnteredPlayMode)
+            {
+                return;
+            }
+            Clear();
+        }
+
+#endif
+
+        #region ContinuationQueue
+        SpinLock gate = new SpinLock();
+        bool dequing = false;
+
+        int actionListCount = 0;
+        Action[] actionList = new Action[InitialSize];
+
+        int waitingListCount = 0;
+        Action[] waitingList = new Action[InitialSize];
+
+
+        public void Enqueue(Action continuation)
+        {
+            bool lockTaken = false;
+            try
+            {
+                gate.Enter(ref lockTaken);
+
+                if (dequing)
+                {
+                    // Ensure Capacity
+                    if (waitingList.Length == waitingListCount)
+                    {
+                        var newLength = waitingListCount * 2;
+                        if ((uint)newLength > MaxArrayLength) newLength = MaxArrayLength;
+
+                        var newArray = new Action[newLength];
+                        Array.Copy(waitingList, newArray, waitingListCount);
+                        waitingList = newArray;
+                    }
+                    waitingList[waitingListCount] = continuation;
+                    waitingListCount++;
+                }
+                else
+                {
+                    // Ensure Capacity
+                    if (actionList.Length == actionListCount)
+                    {
+                        var newLength = actionListCount * 2;
+                        if ((uint)newLength > MaxArrayLength) newLength = MaxArrayLength;
+
+                        var newArray = new Action[newLength];
+                        Array.Copy(actionList, newArray, actionListCount);
+                        actionList = newArray;
+                    }
+                    actionList[actionListCount] = continuation;
+                    actionListCount++;
+                }
+            }
+            finally
+            {
+                if (lockTaken) gate.Exit(false);
+            }
+        }
+
+
+        void RunContinuations()
+        {
+            {
+                bool lockTaken = false;
+                try
+                {
+                    gate.Enter(ref lockTaken);
+                    if (actionListCount == 0) return;
+                    dequing = true;
+                }
+                finally
+                {
+                    if (lockTaken) gate.Exit(false);
+                }
+            }
+
+            for (int i = 0; i < actionListCount; i++)
+            {
+                var action = actionList[i];
+                actionList[i] = null;
+
+                try
+                {
+                    action();
+                }
+                catch (Exception ex)
+                {
+                    UnityEngine.Debug.LogException(ex);
+                }
+            }
+
+            {
+                bool lockTaken = false;
+                try
+                {
+                    gate.Enter(ref lockTaken);
+                    dequing = false;
+
+                    var swapTempActionList = actionList;
+
+                    actionListCount = waitingListCount;
+                    actionList = waitingList;
+
+                    waitingListCount = 0;
+                    waitingList = swapTempActionList;
+                }
+                finally
+                {
+                    if (lockTaken) gate.Exit(false);
+                }
+            }
+        }
+        void ClearContinuations()
+        {
+            actionListCount = 0;
+            actionList = new Action[InitialSize];
+
+            waitingListCount = 0;
+            waitingList = new Action[InitialSize];
+        }
+        #endregion
+
+        #region PlayerLoopRunner
+        readonly object runningAndQueueLock = new object();
+        readonly object arrayLock = new object();
+        readonly Action<Exception> unhandledExceptionCallback;
+
+        int tail = 0;
+        bool running = false;
+        IPlayerLoopItem[] loopItems = new IPlayerLoopItem[InitialSize];
+        MinimumQueue<IPlayerLoopItem> waitQueue = new MinimumQueue<IPlayerLoopItem>(InitialSize);
+        public void AddAction(IPlayerLoopItem item)
+        {
+            lock (runningAndQueueLock)
+            {
+                if (running)
+                {
+                    waitQueue.Enqueue(item);
+                    return;
+                }
+            }
+
+            lock (arrayLock)
+            {
+                // Ensure Capacity
+                if (loopItems.Length == tail)
+                {
+                    Array.Resize(ref loopItems, checked(tail * 2));
+                }
+                loopItems[tail++] = item;
+            }
+        }
+        void RunLoopItems()
+        {
+
+            lock (runningAndQueueLock)
+            {
+                running = true;
+            }
+
+            lock (arrayLock)
+            {
+                var j = tail - 1;
+                var loopItems = this.loopItems;
+                // eliminate array-bound check for i
+                for (int i = 0; i < loopItems.Length; i++)
+                {
+                    var action = loopItems[i];
+                    if (action != null)
+                    {
+                        try
+                        {
+                            if (!action.MoveNext())
+                            {
+                                loopItems[i] = null;
+                            }
+                            else
+                            {
+                                continue; // next i 
+                            }
+                        }
+                        catch (Exception ex)
+                        {
+                            loopItems[i] = null;
+                            try
+                            {
+                                unhandledExceptionCallback(ex);
+                            }
+                            catch { }
+                        }
+                    }
+
+                    // find null, loop from tail
+                    while (i < j)
+                    {
+                        var fromTail = loopItems[j];
+                        if (fromTail != null)
+                        {
+                            try
+                            {
+                                if (!fromTail.MoveNext())
+                                {
+                                    loopItems[j] = null;
+                                    j--;
+                                    continue; // next j
+                                }
+                                else
+                                {
+                                    // swap
+                                    loopItems[i] = fromTail;
+                                    loopItems[j] = null;
+                                    j--;
+                                    goto NEXT_LOOP; // next i
+                                }
+                            }
+                            catch (Exception ex)
+                            {
+                                loopItems[j] = null;
+                                j--;
+                                try
+                                {
+                                    unhandledExceptionCallback(ex);
+                                }
+                                catch { }
+                                continue; // next j
+                            }
+                        }
+                        else
+                        {
+                            j--;
+                        }
+                    }
+
+                    tail = i; // loop end
+                    break; // LOOP END
+
+                NEXT_LOOP:
+                    continue;
+                }
+
+
+                lock (runningAndQueueLock)
+                {
+                    running = false;
+                    while (waitQueue.Count != 0)
+                    {
+                        if (loopItems.Length == tail)
+                        {
+                            Array.Resize(ref loopItems, checked(tail * 2));
+                        }
+                        loopItems[tail++] = waitQueue.Dequeue();
+                    }
+                }
+            }
+
+        }
+
+
+
+
+        void ClearLoopItems()
+        {
+            lock (arrayLock)
+            {
+                for (var index = 0; index < loopItems.Length; index++)
+                {
+                    loopItems[index] = null;
+                }
+            }
+        }
+        #endregion
+
+    }
+
+    
+
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/Internal/UniTaskPlayerLoopSubSystem.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: 21e9b6c323ae3574cab6ff0bf6906a99
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

src/UniTask/Assets/Plugins/UniTask/Runtime/PlayerLoopHelper.cs

@@ -8,6 +8,7 @@ using System.Threading;
 
 #if UNITY_2019_3_OR_NEWER
 using UnityEngine.LowLevel;
+using System.Collections.Generic;
 #else
 using UnityEngine.Experimental.LowLevel;
 #endif
@@ -97,11 +98,34 @@ namespace Cysharp.Threading.Tasks
         public static bool IsMainThread => Thread.CurrentThread.ManagedThreadId == mainThreadId;
 
         static int mainThreadId;
+
+        [ThreadStatic]
+        static int syncState;
+        /// <summary>
+        /// In main thread,(<see cref="PlayerLoopTiming"/>)(<see cref="SyncState"/> - 1) is current <see cref="PlayerLoopTiming"/>.Otherwise, it is 0.
+        /// </summary>
+        internal static int SyncState => syncState;
+
+        internal static void SetCurrentPlayerLoopTiming(PlayerLoopTiming timing) => syncState = (int)timing + 1;
+
+        /// <summary>
+        /// If we are in main thread,returns current <see cref="PlayerLoopTiming"/>. Otherwise,returns <see langword="null"/>.
+        /// </summary>
+        /// <returns></returns>
+        public static PlayerLoopTiming? TryGetCurrentPlayerLoopTiming()
+        {
+            var value = SyncState - 1;
+            return value == -1 ? (PlayerLoopTiming?)null : (PlayerLoopTiming)value;
+        }
+
         static string applicationDataPath;
         static SynchronizationContext unitySynchronizationContetext;
         static ContinuationQueue[] yielders;
         static PlayerLoopRunner[] runners;
 
+        static readonly Dictionary<SyncParams, UniTaskPlayerLoopSubSystem> subSystems = new Dictionary<SyncParams, UniTaskPlayerLoopSubSystem>();//TODO:Replace this with much faster dictionary.
+        static SpinLock subSystemsLock;
+
         static PlayerLoopSystem[] InsertRunner(PlayerLoopSystem loopSystem,
             Type loopRunnerYieldType, ContinuationQueue cq, Type lastLoopRunnerYieldType, ContinuationQueue lastCq,
             Type loopRunnerType, PlayerLoopRunner runner, Type lastLoopRunnerType, PlayerLoopRunner lastRunner)
@@ -231,6 +255,7 @@ namespace Cysharp.Threading.Tasks
 #endif
             
             Initialize(ref playerLoop);
+            SetCurrentPlayerLoopTiming(PlayerLoopTiming.Initialization);
         }
 
 
@@ -278,6 +303,32 @@ namespace Cysharp.Threading.Tasks
 
 #endif
 
+        
+
+        internal static UniTaskPlayerLoopSubSystem GetOrCreateSubSystem(SyncParams syncParams)
+        {
+            var lockTaken = false;
+            try
+            {
+                subSystemsLock.Enter(ref lockTaken);
+                if (!subSystems.TryGetValue(syncParams, out var subSystem))
+                {
+                    subSystem = new UniTaskPlayerLoopSubSystem();
+                    foreach (var timing in syncParams.EnumeratePlayerLoopTimings())
+                    {
+                        AddAction(timing, subSystem);
+                    }
+                    subSystems.Add(syncParams, subSystem);
+                }
+                return subSystem;
+            }
+            finally
+            {
+                if (lockTaken) subSystemsLock.Exit(false);
+            }
+        }
+
+
         public static void Initialize(ref PlayerLoopSystem playerLoop)
         {
             yielders = new ContinuationQueue[14];
@@ -333,11 +384,33 @@ namespace Cysharp.Threading.Tasks
             runners[(int)timing].AddAction(action);
         }
 
+        internal static void AddAction(SyncParams syncParams, IPlayerLoopItem action)
+        {
+            if(syncParams== SyncParams.ThreadPool)
+            {
+                throw new ArgumentException();
+            }
+            else
+            {
+                GetOrCreateSubSystem(syncParams).AddAction(action);
+            }
+
+        }
+
         public static void AddContinuation(PlayerLoopTiming timing, Action continuation)
         {
             yielders[(int)timing].Enqueue(continuation);
         }
 
+        internal static void AddContinuation(SyncParams syncParams, Action continuation)
+        {
+            if(syncParams== SyncParams.ThreadPool)
+            {
+                new SwitchToThreadPoolAwaitable.Awaiter().OnCompleted(continuation);//TODO:Should we use UnsafeOnCompleted at here? 
+            }
+            GetOrCreateSubSystem(syncParams).Enqueue(continuation);
+        }
+
         // Diagnostics helper
 
 #if UNITY_2019_3_OR_NEWER

src/UniTask/Assets/Plugins/UniTask/Runtime/SyncParams.cs

@@ -0,0 +1,90 @@
+using Cysharp.Threading.Tasks;
+using System;
+using System.Runtime.CompilerServices;
+
+namespace Cysharp.Threading.Tasks
+{
+    [Flags]
+    public enum SyncParams
+    {
+        ThreadPool = 0,
+
+        #region PlayerLoopTiming
+        Initialization = 1 << PlayerLoopTiming.Initialization,
+        LastInitialization = 1 << PlayerLoopTiming.LastInitialization,
+
+        EarlyUpdate = 1 << PlayerLoopTiming.EarlyUpdate,
+        LastEarlyUpdate = 1 << PlayerLoopTiming.LastEarlyUpdate,
+
+        FixedUpdate = 1 << PlayerLoopTiming.FixedUpdate,
+        LastFixedUpdate = 1 << PlayerLoopTiming.LastFixedUpdate,
+
+        PreUpdate = 1 << PlayerLoopTiming.PreUpdate,
+        LastPreUpdate = 1 << PlayerLoopTiming.LastPreUpdate,
+
+        Update = 1 << PlayerLoopTiming.Update,
+        LastUpdate = 1 << PlayerLoopTiming.LastUpdate,
+
+        PreLateUpdate = 1 << PlayerLoopTiming.PreLateUpdate,
+        LastPreLateUpdate = 1 << PlayerLoopTiming.LastPreLateUpdate,
+
+        PostLateUpdate = 1 << PlayerLoopTiming.PostLateUpdate,
+        LastPostLateUpdate = 1 << PlayerLoopTiming.LastPostLateUpdate,
+        #endregion
+
+        MainThread = EarlyUpdate | PreUpdate | Update | PreLateUpdate | PostLateUpdate,
+    }
+    public struct SyncParamsPlayerLoopTimingEnumerable
+    {
+        internal SyncParams syncParams;
+        public SyncParamsPlayerLoopTimingEnumerator GetEnumerator() => new SyncParamsPlayerLoopTimingEnumerator() { syncParams = syncParams };
+
+    }
+
+    public static class SyncParamsHelpers
+    {
+        static readonly byte[] TrailingZeroCountDeBruijn =
+        {
+            00, 01, 28, 02, 29, 14, 24, 03,
+            30, 22, 20, 15, 25, 17, 04, 08,
+            31, 27, 13, 23, 21, 19, 16, 07,
+            26, 12, 18, 06, 11, 05, 10, 09
+        };
+        /// <summary>
+        /// From System.Numerics.BitOperations
+        /// </summary>
+        /// <param name="value"></param>
+        /// <returns></returns>
+        /// <remarks>License:https://github.com/dotnet/runtime/blob/6072e4d3a7a2a1493f514cdf4be75a3d56580e84/LICENSE.TXT </remarks>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static int TrailingZeroCount(uint value)
+        {
+            // Using deBruijn sequence, k=2, n=5 (2^5=32) : 0b_0000_0111_0111_1100_1011_0101_0011_0001u
+            //TODO:Avoid array bounds check
+            return TrailingZeroCountDeBruijn[(int)(((value & (uint)-(int)value) * 0x077CB531u) >> 27)];
+        }
+        public static SyncParamsPlayerLoopTimingEnumerable EnumeratePlayerLoopTimings(this SyncParams syncParams) => new SyncParamsPlayerLoopTimingEnumerable() { syncParams = syncParams };
+
+        public static SyncParams FromPlayerLoopTiming(PlayerLoopTiming playerLoopTiming) => (SyncParams)(1 << (int)playerLoopTiming);
+    }
+
+    public struct SyncParamsPlayerLoopTimingEnumerator
+    {
+        internal SyncParams syncParams;
+        public bool MoveNext()
+        {
+            if (syncParams == SyncParams.ThreadPool)
+            {
+                return false;
+            }
+            var value = (int)syncParams;
+            var tzcnt = SyncParamsHelpers.TrailingZeroCount((uint)value);
+            Current = (PlayerLoopTiming)tzcnt;
+            syncParams = (SyncParams)(value ^ (1 << tzcnt));
+            return true;
+        }
+        public PlayerLoopTiming Current { get; private set; }
+    }
+
+
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/SyncParams.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: ebb14e764a35bbf42add491f1035c8a0
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Threading.cs

@@ -28,6 +28,12 @@ namespace Cysharp.Threading.Tasks
             return new SwitchToMainThreadAwaitable(timing, cancellationToken);
         }
 
+
+        public static SwitchToSyncParamsAwaitable SwitchToSyncParams(SyncParams syncParams,CancellationToken cancellationToken = default)
+        {
+            return new SwitchToSyncParamsAwaitable(syncParams, cancellationToken);
+        }
+
         /// <summary>
         /// Return to mainthread(same as await SwitchToMainThread) after using scope is closed.
         /// </summary>
@@ -140,6 +146,67 @@ namespace Cysharp.Threading.Tasks
         }
     }
 
+    public struct SwitchToSyncParamsAwaitable
+    {
+        readonly SyncParams syncParams;
+        readonly CancellationToken cancellationToken;
+
+        internal SwitchToSyncParamsAwaitable(SyncParams syncParams, CancellationToken cancellationToken)
+        {
+            this.syncParams = syncParams;
+            this.cancellationToken = cancellationToken;
+        }
+
+        public Awaiter GetAwaiter() => new Awaiter(syncParams, cancellationToken);
+
+        public struct Awaiter : ICriticalNotifyCompletion
+        {
+            readonly SyncParams syncParams;
+            readonly CancellationToken cancellationToken;
+
+            internal Awaiter(SyncParams syncParams, CancellationToken cancellationToken)
+            {
+                this.syncParams = syncParams;
+                this.cancellationToken = cancellationToken;
+            }
+
+            public bool IsCompleted
+            {
+                get
+                {
+                    if(PlayerLoopHelper.TryGetCurrentPlayerLoopTiming() is PlayerLoopTiming timing)
+                    {
+                        foreach (var syncTiming in syncParams.EnumeratePlayerLoopTimings())
+                        {
+                            if (syncTiming == timing)
+                            {
+                                return true;
+                            }
+                        }
+                    }
+                    else if (syncParams == SyncParams.ThreadPool)
+                    {
+                        return true;
+                    }
+
+                    return false;
+                }
+            }
+
+            public void GetResult() { cancellationToken.ThrowIfCancellationRequested(); }
+
+            public void OnCompleted(Action continuation)
+            {
+                PlayerLoopHelper.AddContinuation(syncParams, continuation);
+            }
+
+            public void UnsafeOnCompleted(Action continuation)
+            {
+                PlayerLoopHelper.AddContinuation(syncParams, continuation);
+            }
+        }
+    }
+
     public struct ReturnToMainThread
     {
         readonly PlayerLoopTiming playerLoopTiming;

src/UniTask/Assets/Tests/SwitchToSyncParamsTest.cs

@@ -0,0 +1,53 @@
+using Cysharp.Threading.Tasks;
+using System.Collections;
+using System.Collections.Generic;
+using UnityEngine;
+using UnityEngine.TestTools;
+using FluentAssertions;
+namespace Cysharp.Threading.TasksTests
+{
+    public class SwitchToSyncParamsTest
+    {
+        [UnityTest]
+        public IEnumerator SyncParamsTest()
+        {
+            return UniTask.ToCoroutine(async () =>
+            {
+                var switchToMainThread = UniTask.SwitchToSyncParams(SyncParams.MainThread);
+                switchToMainThread.GetAwaiter().IsCompleted.Should().BeTrue();
+                await switchToMainThread;
+                var switchToThreadPool = UniTask.SwitchToSyncParams(SyncParams.ThreadPool);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToThreadPool;
+                PlayerLoopHelper.IsMainThread.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.IsMainThread.Should().BeTrue();
+
+                await UniTask.Yield(PlayerLoopTiming.Initialization);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.EarlyUpdate);
+                await UniTask.Yield(PlayerLoopTiming.LastEarlyUpdate);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.PreUpdate);
+                await UniTask.Yield(PlayerLoopTiming.LastPreUpdate);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.Update);
+                await UniTask.Yield(PlayerLoopTiming.LastUpdate);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.PreLateUpdate);
+                await UniTask.Yield(PlayerLoopTiming.LastPreLateUpdate);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.PostLateUpdate);
+                await UniTask.Yield(PlayerLoopTiming.LastPostLateUpdate);
+                switchToThreadPool.GetAwaiter().IsCompleted.Should().BeFalse();
+                await switchToMainThread;
+                PlayerLoopHelper.TryGetCurrentPlayerLoopTiming().Should().Be(PlayerLoopTiming.EarlyUpdate);
+            });
+        }
+    }
+}

src/UniTask/Assets/Tests/SwitchToSyncParamsTest.cs.meta

@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: d3694a1d095bc654c8fdcfe5d79cb78b
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: