Improved async value retrieval

[timcassell]

Similar to #75196, but this time focusing on awaiters rather than async method builders.

Consider ValueTask<TResult> that stores the TResult _result in the struct, so that synchronous value retrieval can occur without allocating on the heap. ValueTaskAwaiter<TResult> has a ValueTask<TResult> field, so it can do the same thing. Currently, async state machines store the entire awaiter, whether it completes synchronously or asynchronously. If the ValueTask<TResult> was pending by the time it was awaited, that's wasted space being stored in the state machine. When the async continuation occurs, the result is retrieved from the backing heap object, not from the result stored in the awaiter.

The compiler might recognize a new pattern: public AsyncResultGetterType GetAsyncResultGetter(). When this method exists on an awaiter, the compiler may use that to store on the state machine instead of the awaiter itself.

The async result getter type must have a GetResult method, which can have the shape of void GetResult() for void values, T GetResult() for T values if the type is generic, or T GetResult<T>(). That last one is new and not available for existing awaiters. The idea is that the async value will be retrieved from a heap object, and only 1 awaiter is used at a time, so the state machine can be optimized by re-using a single async result getter type for multiple awaiter types. For example, await FuncAsync<object>(); await FuncAsync<int>();.

Here are benchmark results comparing the current async state machines with the new optimized version:

Method	Mean	Error	StdDev	Code Size	Allocated
Current	290.5 ns	1.11 ns	1.04 ns	3,068 B	248 B
New	240.4 ns	2.12 ns	1.77 ns	2,911 B	168 B

Code

[MemoryDiagnoser(false)]
[DisassemblyDiagnoser(100)]
public class AsyncStateMachineBenchmarks
{
    private readonly MyValueTaskSource<object> _source1 = new MyValueTaskSource<object>();
    private readonly MyValueTaskSource<Matrix4x4> _source2 = new MyValueTaskSource<Matrix4x4>();

    [Benchmark]
    public void Current()
    {
        _ = AsyncFuncCurrent();

        _source1._continuation.Invoke();
        _source1._continuation = null;
        _source2._continuation.Invoke();
        _source2._continuation = null;
    }

    [Benchmark]
    public void New()
    {
        _ = AsyncFuncNew();

        _source1._continuation.Invoke();
        _source1._continuation = null;
        _source2._continuation.Invoke();
        _source2._continuation = null;
    }

    //private async Task AsyncFunc()
    //{
    //    _ = await new MyValueTask<object>(_source1);
    //    _ = await new MyValueTask<Matrix4x4>(_source2);
    //}

    [StructLayout(LayoutKind.Auto)]
    private struct AsyncFuncCurrentStateMachine : IAsyncStateMachine
    {
        public int __state;

        public AsyncTaskMethodBuilder t__builder;

        public AsyncStateMachineBenchmarks __this;

        private MyValueTaskAwaiter<object> u__1;

        private MyValueTaskAwaiter<Matrix4x4> u__2;

        private void MoveNext()
        {
            int num = __state;
            try
            {
                MyValueTaskAwaiter<Matrix4x4> awaiter;
                MyValueTaskAwaiter<object> awaiter2;
                if (num != 0)
                {
                    if (num == 1)
                    {
                        awaiter = u__2;
                        u__2 = default;
                        num = (__state = -1);
                        goto IL_00be;
                    }
                    awaiter2 = new MyValueTask<object>(__this._source1).GetAwaiter();
                    if (!awaiter2.IsCompleted)
                    {
                        num = (__state = 0);
                        u__1 = awaiter2;
                        t__builder.AwaitUnsafeOnCompleted(ref awaiter2, ref this);
                        return;
                    }
                }
                else
                {
                    awaiter2 = u__1;
                    u__1 = default;
                    num = (__state = -1);
                }
                awaiter2.GetResult();
                awaiter = new MyValueTask<Matrix4x4>(__this._source2).GetAwaiter();
                if (!awaiter.IsCompleted)
                {
                    num = (__state = 1);
                    u__2 = awaiter;
                    t__builder.AwaitUnsafeOnCompleted(ref awaiter, ref this);
                    return;
                }
                goto IL_00be;
                IL_00be:
                awaiter.GetResult();
            }
            catch (Exception exception)
            {
                __state = -2;
                t__builder.SetException(exception);
                return;
            }
            __state = -2;
            t__builder.SetResult();
        }

        void IAsyncStateMachine.MoveNext()
        {
            //ILSpy generated this explicit interface implementation from .override directive in MoveNext
            this.MoveNext();
        }

        private void SetStateMachine(IAsyncStateMachine stateMachine)
        {
            t__builder.SetStateMachine(stateMachine);
        }

        void IAsyncStateMachine.SetStateMachine(IAsyncStateMachine stateMachine)
        {
            //ILSpy generated this explicit interface implementation from .override directive in SetStateMachine
            this.SetStateMachine(stateMachine);
        }
    }

    private Task AsyncFuncCurrent()
    {
        AsyncFuncCurrentStateMachine stateMachine = default;
        stateMachine.t__builder = AsyncTaskMethodBuilder.Create();
        stateMachine.__this = this;
        stateMachine.__state = -1;
        stateMachine.t__builder.Start(ref stateMachine);
        return stateMachine.t__builder.Task;
    }

    [StructLayout(LayoutKind.Auto)]
    private struct AsyncFuncNewStateMachine : IAsyncStateMachine
    {
        public int __state;

        public AsyncTaskMethodBuilder t__builder;

        public AsyncStateMachineBenchmarks __this;

        private MyValueTaskResultGetter u__1;

        private void MoveNext()
        {
            int num = __state;
            try
            {
                MyValueTaskResultGetter resultGetter;
                if (num != 0)
                {
                    if (num == 1)
                    {
                        resultGetter = u__1;
                        u__1 = default;
                        num = (__state = -1);
                        resultGetter.GetResult<Matrix4x4>();
                        goto IL_00be;
                    }
                    var awaiter1 = new MyValueTask<object>(__this._source1).GetAwaiter();
                    if (!awaiter1.IsCompleted)
                    {
                        num = (__state = 0);
                        u__1 = awaiter1.GetAsyncResultGetter();
                        t__builder.AwaitUnsafeOnCompleted(ref awaiter1, ref this);
                        return;
                    }
                    awaiter1.GetResult();
                }
                else
                {
                    resultGetter = u__1;
                    u__1 = default;
                    num = (__state = -1);
                    resultGetter.GetResult<object>();
                }
                var awaiter2 = new MyValueTask<Matrix4x4>(__this._source2).GetAwaiter();
                if (!awaiter2.IsCompleted)
                {
                    num = (__state = 1);
                    u__1 = awaiter2.GetAsyncResultGetter();
                    t__builder.AwaitUnsafeOnCompleted(ref awaiter2, ref this);
                    return;
                }
                awaiter2.GetResult();
                goto IL_00be;
            IL_00be:;
            }
            catch (Exception exception)
            {
                __state = -2;
                t__builder.SetException(exception);
                return;
            }
            __state = -2;
            t__builder.SetResult();
        }

        void IAsyncStateMachine.MoveNext()
        {
            //ILSpy generated this explicit interface implementation from .override directive in MoveNext
            this.MoveNext();
        }

        private void SetStateMachine(IAsyncStateMachine stateMachine)
        {
            t__builder.SetStateMachine(stateMachine);
        }

        void IAsyncStateMachine.SetStateMachine(IAsyncStateMachine stateMachine)
        {
            //ILSpy generated this explicit interface implementation from .override directive in SetStateMachine
            this.SetStateMachine(stateMachine);
        }
    }

    private Task AsyncFuncNew()
    {
        AsyncFuncNewStateMachine stateMachine = default;
        stateMachine.t__builder = AsyncTaskMethodBuilder.Create();
        stateMachine.__this = this;
        stateMachine.__state = -1;
        stateMachine.t__builder.Start(ref stateMachine);
        return stateMachine.t__builder.Task;
    }
}

public readonly struct MyValueTask<T>
{
    internal readonly MyValueTaskSource<T> _source;
    internal readonly T _completedResult;

    public MyValueTask(T completedResult)
    {
        _source = null;
        _completedResult = completedResult;
    }

    internal MyValueTask(MyValueTaskSource<T> source)
    {
        _source = source;
        _completedResult = default;
    }

    public MyValueTaskAwaiter<T> GetAwaiter()
        => new MyValueTaskAwaiter<T>(this);
}

internal abstract class MyValueTaskSourceBase { }

internal sealed class MyValueTaskSource<T> : MyValueTaskSourceBase
{
    internal Action _continuation;
    internal T _result;
}

public readonly struct MyValueTaskAwaiter<T> : ICriticalNotifyCompletion
{
    private readonly MyValueTask<T> _task;

    internal MyValueTaskAwaiter(MyValueTask<T> task)
    {
        _task = task;
    }

    public bool IsCompleted
        => _task._source is null;

    public T GetResult()
        => _task._source is null
            ? _task._completedResult
            : _task._source._result;

    public void OnCompleted(Action continuation)
        => _task._source._continuation = continuation;

    public void UnsafeOnCompleted(Action continuation)
        => OnCompleted(continuation);

    public MyValueTaskResultGetter GetAsyncResultGetter()
        => new MyValueTaskResultGetter(_task._source);
}

public readonly struct MyValueTaskResultGetter
{
    private readonly MyValueTaskSourceBase _source;

    internal MyValueTaskResultGetter(MyValueTaskSourceBase source)
    {
        _source = source;
    }

    public void GetResult() { }

    public TResult GetResult<TResult>()
        => Unsafe.As<MyValueTaskSource<TResult>>(_source)._result;
}