reduce AsyncBuilder heap allocation

src/UniTask.NetCoreSandbox/Program.cs

@@ -205,7 +205,7 @@ namespace NetCoreSandbox
 #endif
 
 
-            AsyncTest().Forget();
+            // AsyncTest().Forget();
 
 
             //AsyncTest().Forget();

src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs

@@ -13,8 +13,8 @@ namespace Cysharp.Threading.Tasks.CompilerServices
     public struct AsyncUniTaskMethodBuilder
     {
         // cache items.
-        AutoResetUniTaskCompletionSource promise;
+        internal IMoveNextRunnerPromise runnerPromise;
-        internal IMoveNextRunner runner;
+        Exception ex;
 
         // 1. Static Create method.
         [DebuggerHidden]
@@ -31,18 +31,18 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
             get
             {
-                if (promise != null)
+                if (runnerPromise != null)
                 {
-                    return promise.Task;
+                    return runnerPromise.Task;
                 }
-
+                else if (ex != null)
-                if (runner == null)
+                {
+                    return UniTask.FromException(ex);
+                }
+                else
                 {
                     return UniTask.CompletedTask;
                 }
-
-                promise = AutoResetUniTaskCompletionSource.Create();
-                return promise.Task;
             }
         }
 
@@ -50,22 +50,13 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         [DebuggerHidden]
         public void SetException(Exception exception)
         {
-            var p = promise; // after return, promise will clear so require to store local.
+            if (runnerPromise == null)
-
-            // runner is finished, return first.
-            if (runner != null)
             {
-                runner.Return();
+                ex = exception;
-                runner = null;
-            }
-
-            if (p != null)
-            {
-                p.TrySetException(exception);
             }
             else
             {
-                promise = AutoResetUniTaskCompletionSource.CreateFromException(exception, out _);
+                runnerPromise.SetException(exception);
             }
         }
 
@@ -73,18 +64,9 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         [DebuggerHidden]
         public void SetResult()
         {
-            var p = promise; // after return, promise will clear so require to store local.
+            if (runnerPromise != null)
-
-            // runner is finished, return first.
-            if (runner != null)
             {
-                runner.Return();
+                runnerPromise.SetResult();
-                runner = null;
-            }
-
-            if (p != null)
-            {
-                p.TrySetResult();
             }
         }
 
@@ -94,16 +76,12 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             where TAwaiter : INotifyCompletion
             where TStateMachine : IAsyncStateMachine
         {
-            if (promise == null)
+            if (runnerPromise == null)
             {
-                promise = AutoResetUniTaskCompletionSource.Create();
+                MoveNextRunnerPromise<TStateMachine>.SetStateMachine(ref this, ref stateMachine);
-            }
-            if (runner == null)
-            {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
             }
 
-            awaiter.OnCompleted(runner.CallMoveNext);
+            awaiter.OnCompleted(runnerPromise.MoveNext);
         }
 
         // 6. AwaitUnsafeOnCompleted
@@ -113,16 +91,12 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             where TAwaiter : ICriticalNotifyCompletion
             where TStateMachine : IAsyncStateMachine
         {
-            if (promise == null)
+            if (runnerPromise == null)
             {
-                promise = AutoResetUniTaskCompletionSource.Create();
+                MoveNextRunnerPromise<TStateMachine>.SetStateMachine(ref this, ref stateMachine);
-            }
-            if (runner == null)
-            {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
             }
 
-            awaiter.UnsafeOnCompleted(runner.CallMoveNext);
+            awaiter.UnsafeOnCompleted(runnerPromise.MoveNext);
         }
 
         // 7. Start
@@ -161,8 +135,8 @@ namespace Cysharp.Threading.Tasks.CompilerServices
     public struct AsyncUniTaskMethodBuilder<T>
     {
         // cache items.
-        AutoResetUniTaskCompletionSource<T> promise;
+        internal IMoveNextRunnerPromise<T> runnerPromise;
-        internal IMoveNextRunner runner;
+        Exception ex;
         T result;
 
         // 1. Static Create method.
@@ -179,18 +153,18 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         {
             get
             {
-                if (promise != null)
+                if (runnerPromise != null)
                 {
-                    return promise.Task;
+                    return runnerPromise.Task;
                 }
-
+                else if (ex != null)
-                if (runner == null)
+                {
+                    return UniTask.FromException<T>(ex);
+                }
+                else
                 {
                     return UniTask.FromResult(result);
                 }
-
-                promise = AutoResetUniTaskCompletionSource<T>.Create();
-                return promise.Task;
             }
         }
 
@@ -198,22 +172,13 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         [DebuggerHidden]
         public void SetException(Exception exception)
         {
-            var p = promise; // after return, promise will clear so require to store local.
+            if (runnerPromise == null)
-
-            // runner is finished, return first.
-            if (runner != null)
             {
-                runner.Return();
+                ex = exception;
-                runner = null;
-            }
-
-            if (p == null)
-            {
-                promise = AutoResetUniTaskCompletionSource<T>.CreateFromException(exception, out _);
             }
             else
             {
-                p.TrySetException(exception);
+                runnerPromise.SetException(exception);
             }
         }
 
@@ -221,22 +186,13 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         [DebuggerHidden]
         public void SetResult(T result)
         {
-            var p = promise; // after return, promise will clear so require to store local.
+            if (runnerPromise == null)
-
-            // runner is finished, return first.
-            if (runner != null)
-            {
-                runner.Return();
-                runner = null;
-            }
-
-            if (p == null)
             {
                 this.result = result;
             }
             else
             {
-                p.TrySetResult(result);
+                runnerPromise.SetResult(result);
             }
         }
 
@@ -246,16 +202,12 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             where TAwaiter : INotifyCompletion
             where TStateMachine : IAsyncStateMachine
         {
-            if (promise == null)
+            if (runnerPromise == null)
             {
-                promise = AutoResetUniTaskCompletionSource<T>.Create();
+                MoveNextRunnerPromise<TStateMachine, T>.SetStateMachine(ref this, ref stateMachine);
-            }
-            if (runner == null)
-            {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
             }
 
-            awaiter.OnCompleted(runner.CallMoveNext);
+            awaiter.OnCompleted(runnerPromise.MoveNext);
         }
 
         // 6. AwaitUnsafeOnCompleted
@@ -265,16 +217,12 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             where TAwaiter : ICriticalNotifyCompletion
             where TStateMachine : IAsyncStateMachine
         {
-            if (promise == null)
+            if (runnerPromise == null)
             {
-                promise = AutoResetUniTaskCompletionSource<T>.Create();
+                MoveNextRunnerPromise<TStateMachine, T>.SetStateMachine(ref this, ref stateMachine);
-            }
-            if (runner == null)
-            {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
             }
 
-            awaiter.UnsafeOnCompleted(runner.CallMoveNext);
+            awaiter.UnsafeOnCompleted(runnerPromise.MoveNext);
         }
 
         // 7. Start

src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs

@@ -65,7 +65,7 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         {
             if (runner == null)
             {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
+                MoveNextRunner<TStateMachine>.SetStateMachine(ref this, ref stateMachine);
             }
 
             awaiter.OnCompleted(runner.CallMoveNext);
@@ -80,7 +80,7 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         {
             if (runner == null)
             {
-                MoveNextRunner<TStateMachine>.SetRunner(ref this, ref stateMachine);
+                MoveNextRunner<TStateMachine>.SetStateMachine(ref this, ref stateMachine);
             }
 
             awaiter.UnsafeOnCompleted(runner.CallMoveNext);

src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/MoveNextRunner.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using Cysharp.Threading.Tasks.Internal;
 
 namespace Cysharp.Threading.Tasks.CompilerServices
@@ -14,13 +15,29 @@ namespace Cysharp.Threading.Tasks.CompilerServices
         void Return();
     }
 
+    internal interface IMoveNextRunnerPromise : IUniTaskSource
+    {
+        Action MoveNext { get; }
+        UniTask Task { get; }
+        void SetResult();
+        void SetException(Exception exception);
+    }
+
+    internal interface IMoveNextRunnerPromise<T> : IUniTaskSource<T>
+    {
+        Action MoveNext { get; }
+        UniTask<T> Task { get; }
+        void SetResult(T result);
+        void SetException(Exception exception);
+    }
+
     internal sealed class MoveNextRunner<TStateMachine> : IMoveNextRunner, IPromisePoolItem
         where TStateMachine : IAsyncStateMachine
     {
         static PromisePool<MoveNextRunner<TStateMachine>> pool = new PromisePool<MoveNextRunner<TStateMachine>>();
 
         TStateMachine stateMachine;
-        internal readonly Action callMoveNext;
+        readonly Action callMoveNext;
 
         public Action CallMoveNext => callMoveNext;
 
@@ -29,31 +46,7 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             callMoveNext = Run;
         }
 
-        public static void SetRunner(ref AsyncUniTaskMethodBuilder builder, ref TStateMachine stateMachine)
+        public static void SetStateMachine(ref AsyncUniTaskVoidMethodBuilder builder, ref TStateMachine stateMachine)
-        {
-            var result = pool.TryRent();
-            if (result == null)
-            {
-                result = new MoveNextRunner<TStateMachine>();
-            }
-
-            builder.runner = result; // set runner before copied.
-            result.stateMachine = stateMachine; // copy struct StateMachine(in release build).
-        }
-
-        public static void SetRunner<T>(ref AsyncUniTaskMethodBuilder<T> builder, ref TStateMachine stateMachine)
-        {
-            var result = pool.TryRent();
-            if (result == null)
-            {
-                result = new MoveNextRunner<TStateMachine>();
-            }
-
-            builder.runner = result; // set runner before copied.
-            result.stateMachine = stateMachine; // copy struct StateMachine(in release build).
-        }
-
-        public static void SetRunner(ref AsyncUniTaskVoidMethodBuilder builder, ref TStateMachine stateMachine)
         {
             var result = pool.TryRent();
             if (result == null)
@@ -82,5 +75,223 @@ namespace Cysharp.Threading.Tasks.CompilerServices
             stateMachine = default;
         }
     }
+
+    internal class MoveNextRunnerPromise<TStateMachine> : IMoveNextRunnerPromise, IUniTaskSource, IPromisePoolItem
+        where TStateMachine : IAsyncStateMachine
+    {
+        static readonly PromisePool<MoveNextRunnerPromise<TStateMachine>> pool = new PromisePool<MoveNextRunnerPromise<TStateMachine>>();
+
+        TStateMachine stateMachine;
+
+        public Action MoveNext { get; }
+
+        UniTaskCompletionSourceCore<AsyncUnit> core;
+
+        MoveNextRunnerPromise()
+        {
+            MoveNext = Run;
+        }
+
+        public static void SetStateMachine(ref AsyncUniTaskMethodBuilder builder, ref TStateMachine stateMachine)
+        {
+            var result = pool.TryRent();
+            if (result == null)
+            {
+                result = new MoveNextRunnerPromise<TStateMachine>();
+            }
+
+            TaskTracker.TrackActiveTask(result, 2);
+
+            builder.runnerPromise = result; // set runner before copied.
+            result.stateMachine = stateMachine; // copy struct StateMachine(in release build).
+        }
+
+        [DebuggerHidden]
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        void Run()
+        {
+            stateMachine.MoveNext();
+        }
+
+        public UniTask Task
+        {
+            [DebuggerHidden]
+            get
+            {
+                return new UniTask(this, core.Version);
+            }
+        }
+
+        [DebuggerHidden]
+        public void SetResult()
+        {
+            core.TrySetResult(AsyncUnit.Default);
+        }
+
+        [DebuggerHidden]
+        public void SetException(Exception exception)
+        {
+            core.TrySetException(exception);
+        }
+
+        [DebuggerHidden]
+        public void GetResult(short token)
+        {
+            try
+            {
+                TaskTracker.RemoveTracking(this);
+                core.GetResult(token);
+            }
+            finally
+            {
+                pool.TryReturn(this);
+            }
+        }
+
+        [DebuggerHidden]
+        public UniTaskStatus GetStatus(short token)
+        {
+            return core.GetStatus(token);
+        }
+
+        [DebuggerHidden]
+        public UniTaskStatus UnsafeGetStatus()
+        {
+            return core.UnsafeGetStatus();
+        }
+
+        [DebuggerHidden]
+        public void OnCompleted(Action<object> continuation, object state, short token)
+        {
+            core.OnCompleted(continuation, state, token);
+        }
+
+        [DebuggerHidden]
+        void IPromisePoolItem.Reset()
+        {
+            stateMachine = default;
+            core.Reset();
+        }
+
+        ~MoveNextRunnerPromise()
+        {
+            if (pool.TryReturn(this))
+            {
+                GC.ReRegisterForFinalize(this);
+            }
+        }
+    }
+
+    internal class MoveNextRunnerPromise<TStateMachine, T> : IMoveNextRunnerPromise<T>, IUniTaskSource<T>, IPromisePoolItem
+        where TStateMachine : IAsyncStateMachine
+    {
+        static readonly PromisePool<MoveNextRunnerPromise<TStateMachine, T>> pool = new PromisePool<MoveNextRunnerPromise<TStateMachine, T>>();
+
+        TStateMachine stateMachine;
+
+        public Action MoveNext { get; }
+
+        UniTaskCompletionSourceCore<T> core;
+
+        MoveNextRunnerPromise()
+        {
+            MoveNext = Run;
+        }
+
+        public static void SetStateMachine(ref AsyncUniTaskMethodBuilder<T> builder, ref TStateMachine stateMachine)
+        {
+            var result = pool.TryRent();
+            if (result == null)
+            {
+                result = new MoveNextRunnerPromise<TStateMachine, T>();
+            }
+
+            TaskTracker.TrackActiveTask(result, 2);
+
+            builder.runnerPromise = result; // set runner before copied.
+            result.stateMachine = stateMachine; // copy struct StateMachine(in release build).
+        }
+
+        [DebuggerHidden]
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        void Run()
+        {
+            stateMachine.MoveNext();
+        }
+
+        public UniTask<T> Task
+        {
+            [DebuggerHidden]
+            get
+            {
+                return new UniTask<T>(this, core.Version);
+            }
+        }
+
+        [DebuggerHidden]
+        public void SetResult(T result)
+        {
+            core.TrySetResult(result);
+        }
+
+        [DebuggerHidden]
+        public void SetException(Exception exception)
+        {
+            core.TrySetException(exception);
+        }
+
+        [DebuggerHidden]
+        public T GetResult(short token)
+        {
+            try
+            {
+                TaskTracker.RemoveTracking(this);
+                return core.GetResult(token);
+            }
+            finally
+            {
+                pool.TryReturn(this);
+            }
+        }
+
+        [DebuggerHidden]
+        void IUniTaskSource.GetResult(short token)
+        {
+            GetResult(token);
+        }
+
+        [DebuggerHidden]
+        public UniTaskStatus GetStatus(short token)
+        {
+            return core.GetStatus(token);
+        }
+
+        [DebuggerHidden]
+        public UniTaskStatus UnsafeGetStatus()
+        {
+            return core.UnsafeGetStatus();
+        }
+
+        [DebuggerHidden]
+        public void OnCompleted(Action<object> continuation, object state, short token)
+        {
+            core.OnCompleted(continuation, state, token);
+        }
+
+        [DebuggerHidden]
+        void IPromisePoolItem.Reset()
+        {
+            stateMachine = default;
+            core.Reset();
+        }
+
+        ~MoveNextRunnerPromise()
+        {
+            if (pool.TryReturn(this))
+            {
+                GC.ReRegisterForFinalize(this);
+            }
+        }
+    }
 }
 

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

@@ -12,10 +12,11 @@ namespace Cysharp.Threading.Tasks
 {
     internal static class AwaiterActions
     {
-        internal static readonly Action<object> InvokeActionDelegate = InvokeAction;
+        internal static readonly Action<object> InvokeContinuationDelegate = Continuation;
 
         [DebuggerHidden]
-        static void InvokeAction(object state)
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        static void Continuation(object state)
         {
             ((Action)state).Invoke();
         }
@@ -312,7 +313,7 @@ namespace Cysharp.Threading.Tasks
                 }
                 else
                 {
-                    task.source.OnCompleted(AwaiterActions.InvokeActionDelegate, continuation, task.token);
+                    task.source.OnCompleted(AwaiterActions.InvokeContinuationDelegate, continuation, task.token);
                 }
             }
 
@@ -326,7 +327,7 @@ namespace Cysharp.Threading.Tasks
                 }
                 else
                 {
-                    task.source.OnCompleted(AwaiterActions.InvokeActionDelegate, continuation, task.token);
+                    task.source.OnCompleted(AwaiterActions.InvokeContinuationDelegate, continuation, task.token);
                 }
             }
 
@@ -650,7 +651,7 @@ namespace Cysharp.Threading.Tasks
                 }
                 else
                 {
-                    s.OnCompleted(AwaiterActions.InvokeActionDelegate, continuation, task.token);
+                    s.OnCompleted(AwaiterActions.InvokeContinuationDelegate, continuation, task.token);
                 }
             }
 
@@ -665,7 +666,7 @@ namespace Cysharp.Threading.Tasks
                 }
                 else
                 {
-                    s.OnCompleted(AwaiterActions.InvokeActionDelegate, continuation, task.token);
+                    s.OnCompleted(AwaiterActions.InvokeContinuationDelegate, continuation, task.token);
                 }
             }