#nullable enable

namespace BrokenAsyncMethodBuilder

{

    using System;

    using System.Diagnostics;

    using System.Runtime.CompilerServices;

    using System.Threading;

    using System.Threading.Tasks;

​

    [AsyncMethodBuilder(typeof(CustomAwaitableAsyncMethodBuilder<>))]

    public readonly struct CustomAwaitable<T>

    {

        readonly ValueTask<T> _valueTask;

​

        public CustomAwaitable(ValueTask<T> valueTask)

        {

            _valueTask = valueTask;

        }

​

        public ValueTaskAwaiter<T> GetAwaiter() => _valueTask.GetAwaiter();

    }

​

    public struct CustomAwaitableAsyncMethodBuilder<T>

    //     ^^^^^^ Only works if you change this to `class` or run in Debug mode

    {

        #region fields

​

        Exception? _exception;

        bool _hasResult;

        SpinLock _lock;

        T? _result;

        TaskCompletionSource<T>? _source;

​

        #endregion

​

        #region properties

​

        public CustomAwaitable<T> Task

        {

            get

            {

                var lockTaken = false;

                try

                {

                    _lock.Enter(ref lockTaken);

                    if (_exception is not null)

                        return new CustomAwaitable<T>(ValueTask.FromException<T>(_exception));

                    if (_hasResult)

                        return new CustomAwaitable<T>(ValueTask.FromResult(_result!));

                    return new CustomAwaitable<T>(

                        new ValueTask<T>(

                            (_source ??= new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously))

                            .Task

                        )

                    );

                }

                finally

                {

                    if (lockTaken)

                        _lock.Exit();

                }

            }

        }

​

        public void AwaitOnCompleted<TAwaiter, TStateMachine>(

            ref TAwaiter awaiter,

            ref TStateMachine stateMachine)

            where TAwaiter : INotifyCompletion

            where TStateMachine : IAsyncStateMachine =>

            awaiter.OnCompleted(stateMachine.MoveNext);

​

        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(

            ref TAwaiter awaiter,

            ref TStateMachine stateMachine)

            where TAwaiter : ICriticalNotifyCompletion

            where TStateMachine : IAsyncStateMachine =>

            awaiter.UnsafeOnCompleted(stateMachine.MoveNext);

​

        #endregion

​

        #region methods

​

        public static CustomAwaitableAsyncMethodBuilder<T> Create() => new()

        {

            _lock = new SpinLock(Debugger.IsAttached)

        };

​

        public void SetException(Exception exception)

        {

            var lockTaken = false;

            try

            {

                _lock.Enter(ref lockTaken);

                if (Volatile.Read(ref _source) is {} source)

                {

                    source.TrySetException(exception);

                }

                else

                {

                    _exception = exception;

                }

            }

            finally

            {

                if (lockTaken)

                    _lock.Exit();

            }

        }

​

        public void SetResult(T result)

        {

            var lockTaken = false;

            try

            {

                _lock.Enter(ref lockTaken);

                if (Volatile.Read(ref _source) is {} source)

                {

                    source.TrySetResult(result);

                }

                else

                {

                    _result = result;

                    _hasResult = true;

                }

            }

            finally

            {

                if (lockTaken)

                    _lock.Exit();

            }

        }

​

        public void SetStateMachine(IAsyncStateMachine stateMachine) {}

​

        public void Start<TStateMachine>(ref TStateMachine stateMachine)

            where TStateMachine : IAsyncStateMachine => stateMachine.MoveNext();

​

        #endregion

    }

​

    class Program

    {

        static async Task Main()

        {

            var expected = Guid.NewGuid().ToString();

            async CustomAwaitable<string> GetValueAsync()

            {

                await Task.Yield();

                return expected;

            }

​

            var actual = await GetValueAsync();

​

            if (!ReferenceEquals(expected, actual))

                throw new Exception();

​

            Console.WriteLine("Done!");

        }

    }

}