Custom AsyncMethodBuilder method overloads

[timcassell]

I tried adding method overloads for my custom builder to add a feature for my library in async functions, using the awaiter type.

public struct PromiseMethodBuilder<T>
{
    ...

    [SecuritySafeCritical]
    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : ICriticalNotifyCompletion
        where TStateMachine : IAsyncStateMachine
    {
        Console.WriteLine("AwaitUnsafeOnCompleted default");
        ...
    }

    [SecuritySafeCritical]
    public void AwaitUnsafeOnCompleted<TStateMachine>(ref PromiseAwaiterVoid awaiter, ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine
    {
        Console.WriteLine("AwaitUnsafeOnCompleted PromiseAwaiterVoid");
        ...
    }

    [SecuritySafeCritical]
    public void AwaitUnsafeOnCompleted<T, TStateMachine>(ref PromiseAwaiter<T> awaiter, ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine
    {
        Console.WriteLine("AwaitUnsafeOnCompleted PromiseAwaiter<T>");
        ...
    }
    
    ...
}

Using the builder with those overloads directly works as I expect:

var deferred = Promise.NewDeferred<int>();
PromiseAwaiter<int> awaiter = deferred.Promise.GetAwaiter();
IAsyncStateMachine stateMachine = null;
new PromiseMethodBuilder<int>().AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine); // Prints "AwaitUnsafeOnCompleted PromiseAwaiter<T>"

But when I tested it with an actual async Promise function, it printed AwaitUnsafeOnCompleted default.

var deferred = Promise.NewDeferred<int>();

async Promise<int> Func()
{
    return await deferred.Promise; // Prints "AwaitUnsafeOnCompleted default"
}

So the compiler chose to use the default overload despite the existence of the better overload candidate that the direct call uses.

Is this possible to work this way? Does this need to be a new feature request to allow this?

I'm not sure if this should be in runtime or csharplang.

@stephentoub

[Edit] I understand why this is the case (the compiler always includes the <TAwaiter, TStateMachine> type arguments to the method call explicitly), but is it possible to override that for specific types?

[stephentoub]

You would need an addition to the language/ compiler. But why do you need overloads rather than just dynamically type testing the TAwaiter? The latter is what the built-in builders do for special-casing the awaiters they know about, no overloads required.

[DaZombieKiller]

@timcassell You can use https://github.com/sschoener/unity-asm-explorer-package to view the generated assembly in Unity, though I'm not sure if it still works with the latest versions of the engine (I have not personally tested). Of course, this doesn't apply to IL2CPP, only the Mono JIT.

[timcassell]

[Edit] Sorry, I misread.

[timcassell]

If I used a sequential layout instead of auto, could I rely on that for an unsafe cast? Since I don't actually care what the T is for this purpose.

public readonly struct Promise
{
    internal readonly Promise<Internal.VoidResult> _target;
}

[StructLayout(LayoutKind.Sequential)]
public readonly struct Promise<T>
{
    internal readonly Internal.PromiseRef _ref;
    internal readonly int _depth;
    internal readonly short _id;
    internal readonly T _result;
}

public readonly struct PromiseAwaiter : IPromiseAwaiter
{
    internal readonly PromiseAwaiter<Internal.VoidResult> _target;
}

public readonly struct PromiseAwaiter<T> : IPromiseAwaiter
{
    internal readonly Promise<T> _target;
}

...

Unsafe.As<TAwaiter, PromiseAwaiter>(ref awaiter).AwaitUnsafeOnCompletedInternal(ref this);

[timcassell]

For anyone also wondering, I just profiled builds in Unity, and it looks like neither the Mono nor IL2CPP runtimes optimize away the ((IMyInterface)awaiter).MyMethod(); or awaiter is IMyInterface boxes, but they do both optimize away the null != default(T) box.

[timcassell]

This is my workaround to get it to work without boxing in Unity:

internal abstract class AwaitOverrider<T> where T : INotifyCompletion
{
    internal static AwaitOverrider<T> awaitOverrider;

    internal abstract void AwaitOnCompletedInternal(ref T awaiter, Internal.AsyncPromiseRef asyncPromiseRef);
}

internal sealed class AwaitOverriderImpl<TAwaiter> : AwaitOverrider<TAwaiter> where TAwaiter : struct, ICriticalNotifyCompletion, IPromiseAwaiter
{
    internal static void Create()
    {
        awaitOverrider = new AwaitOverriderImpl<TAwaiter>();
    }

    internal override void AwaitOnCompletedInternal(ref TAwaiter awaiter, Internal.AsyncPromiseRef asyncPromiseRef)
    {
        awaiter.AwaitOnCompletedInternal(asyncPromiseRef);
    }
}

internal interface IPromiseAwaiter
{
    void AwaitOnCompletedInternal(Internal.AsyncPromiseRef asyncPromiseRef);
}

public struct PromiseAwaiter<T> : ICriticalNotifyCompletion, IPromiseAwaiter
{
    static PromiseAwaiter()
    {
        AwaitOverriderImpl<PromiseAwaiter<T>>.Create();
    }

    void IPromiseAwaiter.AwaitOnCompletedInternal(Internal.AsyncPromiseRef asyncPromiseRef)
    {
        // Internal implementation...
    }
}

if (null != default(TAwaiter) && typeof(IPromiseAwaiter).IsAssignableFrom(typeof(TAwaiter)))
{
    AwaitOverrider<TAwaiter>.awaitOverrider.AwaitOnCompletedInternal(ref awaiter, _promiseRef);
}
else
{
    // Normal path...
}

This could also be optimized more with a function pointer instead of an object instance, but I'm still supporting older Unity versions with C#7. And the IsAssignableFrom could be changed to a null check, but I haven't benchmarked that to see if it would be faster.