Virtual methods with runtime async

src/coreclr/tools/Common/Compiler/AsyncMethodVariant.cs

@@ -79,12 +79,9 @@ protected override int CompareToImpl(MethodDesc other, TypeSystemComparer compar
 
     public static class AsyncMethodVariantExtensions
     {
-        /// <summary>
-        /// Returns true if this MethodDesc is an AsyncMethodVariant, which should not escape the jit interface.
-        /// </summary>
         public static bool IsAsyncVariant(this MethodDesc method)
         {
-            return method is AsyncMethodVariant;
+            return method.GetMethodDefinition() is AsyncMethodVariant;
         }
     }
 }

src/coreclr/tools/Common/Compiler/CompilerTypeSystemContext.Async.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Generic;
+
 using Internal.IL;
 using Internal.TypeSystem;
 using Internal.TypeSystem.Ecma;
@@ -11,6 +13,141 @@ namespace ILCompiler
 {
     public partial class CompilerTypeSystemContext
     {
+        private sealed class AsyncAwareVirtualMethodResolutionAlgorithm : MetadataVirtualMethodAlgorithm
+        {
+            private readonly CompilerTypeSystemContext _context;
+
+            public AsyncAwareVirtualMethodResolutionAlgorithm(CompilerTypeSystemContext context)
+                => _context = context;
+
+            private MethodDesc DecomposeAsyncVariant(MethodDesc method, out bool isAsyncVariant)
+            {
+                isAsyncVariant = method.IsAsyncVariant();
+                return isAsyncVariant ? _context.GetTargetOfAsyncVariantMethod(method) : method;
+            }
+
+            public override MethodDesc FindVirtualFunctionTargetMethodOnObjectType(MethodDesc targetMethod, TypeDesc objectType)
+            {
+                targetMethod = DecomposeAsyncVariant(targetMethod, out bool isAsyncSlot);
+                MethodDesc result = base.FindVirtualFunctionTargetMethodOnObjectType(targetMethod, objectType);
+                if (result != null && isAsyncSlot)
+                    result = _context.GetAsyncVariantMethod(result);
+
+                return result;
+            }
+
+            public override DefaultInterfaceMethodResolution ResolveInterfaceMethodToDefaultImplementationOnType(MethodDesc interfaceMethod, TypeDesc currentType, out MethodDesc impl)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                DefaultInterfaceMethodResolution result = base.ResolveInterfaceMethodToDefaultImplementationOnType(interfaceMethod, currentType, out impl);
+                if (impl != null && isAsyncSlot)
+                    impl = _context.GetAsyncVariantMethod(impl);
+
+                return result;
+            }
+
+            public override MethodDesc ResolveInterfaceMethodToStaticVirtualMethodOnType(MethodDesc interfaceMethod, TypeDesc currentType)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                MethodDesc result = base.ResolveInterfaceMethodToStaticVirtualMethodOnType(interfaceMethod, currentType);
+                if (result != null && isAsyncSlot)
+                    result = _context.GetAsyncVariantMethod(result);
+
+                return result;
+            }
+            public override MethodDesc ResolveInterfaceMethodToVirtualMethodOnType(MethodDesc interfaceMethod, TypeDesc currentType)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                MethodDesc result = base.ResolveInterfaceMethodToVirtualMethodOnType(interfaceMethod, currentType);
+                if (result != null && isAsyncSlot)
+                    result = _context.GetAsyncVariantMethod(result);
+
+                return result;
+            }
+            public override DefaultInterfaceMethodResolution ResolveVariantInterfaceMethodToDefaultImplementationOnType(MethodDesc interfaceMethod, TypeDesc currentType, out MethodDesc impl)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                DefaultInterfaceMethodResolution result = base.ResolveVariantInterfaceMethodToDefaultImplementationOnType(interfaceMethod, currentType, out impl);
+                if (impl != null && isAsyncSlot)
+                    impl = _context.GetAsyncVariantMethod(impl);
+
+                return result;
+            }
+            public override MethodDesc ResolveVariantInterfaceMethodToStaticVirtualMethodOnType(MethodDesc interfaceMethod, TypeDesc currentType)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                MethodDesc result = base.ResolveVariantInterfaceMethodToStaticVirtualMethodOnType(interfaceMethod, currentType);
+                if (result != null && isAsyncSlot)
+                    result = _context.GetAsyncVariantMethod(result);
+
+                return result;
+            }
+            public override MethodDesc ResolveVariantInterfaceMethodToVirtualMethodOnType(MethodDesc interfaceMethod, TypeDesc currentType)
+            {
+                interfaceMethod = DecomposeAsyncVariant(interfaceMethod, out bool isAsyncSlot);
+                MethodDesc result = base.ResolveVariantInterfaceMethodToVirtualMethodOnType(interfaceMethod, currentType);
+                if (result != null && isAsyncSlot)
+                    result = _context.GetAsyncVariantMethod(result);
+
+                return result;
+            }
+
+            public override IEnumerable<MethodDesc> ComputeAllVirtualSlots(TypeDesc type)
+            {
+                foreach (MethodDesc method in base.ComputeAllVirtualSlots(type))
+                {
+                    yield return method;
+
+                    // We create an async variant slot for any Task-returning method, not just runtime-async.
+                    // This is not a problem in practice because the slot is still subject to dependency
+                    // analysis and if not used, will not be generated.
+                    //
+                    // The reason why we need it is this:
+                    //
+                    // interface IFoo
+                    // {
+                    //     [RuntimeAsyncMethodGeneration(true)]
+                    //     Task Method();
+                    // }
+                    //
+                    // class Base
+                    // {
+                    //     [RuntimeAsyncMethodGeneration(false)]
+                    //     public virtual Task Method();
+                    // }
+                    //
+                    // class Derived : Base, IFoo
+                    // {
+                    //     // Q: The runtime-async implementation for IFoo.Method
+                    //     //    comes from Base. However Base was not runtime-async and we
+                    //     //    didn't know about IFoo in Base either. Who has the slot?
+                    //     // A: Base has the runtime-async slot, despite the method not being runtime-async.
+                    // }
+                    if (method.GetTypicalMethodDefinition().Signature.ReturnsTaskOrValueTask())
+                        yield return _context.GetAsyncVariantMethod(method);
+                }
+            }
+        }
+
+        public MethodDesc GetTargetOfAsyncVariantMethod(MethodDesc asyncVariantMethod)
+        {
+            var asyncMethodVariantDefinition = (AsyncMethodVariant)asyncVariantMethod.GetTypicalMethodDefinition();
+            MethodDesc result = asyncMethodVariantDefinition.Target;
+
+            // If there are generics involved, we need to specialize
+            if (asyncVariantMethod != asyncMethodVariantDefinition)
+            {
+                TypeDesc owningType = asyncVariantMethod.OwningType;
+                if (owningType != asyncMethodVariantDefinition.OwningType)
+                    result = GetMethodForInstantiatedType(result, (InstantiatedType)owningType);
+
+                if (asyncVariantMethod.HasInstantiation && !asyncVariantMethod.IsMethodDefinition)
+                    result = GetInstantiatedMethod(result, asyncVariantMethod.Instantiation);
+            }
+
+            return result;
+        }
+
         public MethodDesc GetAsyncVariantMethod(MethodDesc taskReturningMethod)
         {
             Debug.Assert(taskReturningMethod.Signature.ReturnsTaskOrValueTask());

src/coreclr/tools/Common/Compiler/CompilerTypeSystemContext.cs

@@ -17,7 +17,6 @@ namespace ILCompiler
     public partial class CompilerTypeSystemContext : MetadataTypeSystemContext, IMetadataStringDecoderProvider
     {
         private readonly MetadataRuntimeInterfacesAlgorithm _metadataRuntimeInterfacesAlgorithm = new MetadataRuntimeInterfacesAlgorithm();
-        private readonly MetadataVirtualMethodAlgorithm _virtualMethodAlgorithm = new MetadataVirtualMethodAlgorithm();
 
         private MetadataStringDecoder _metadataStringDecoder;
 

src/coreclr/tools/Common/TypeSystem/IL/NativeAotILProvider.cs

@@ -378,8 +378,7 @@ public override MethodIL GetMethodIL(MethodDesc method)
                 }
                 else
                 {
-                    // TODO: Emit thunk with async calling convention
-                    throw new NotImplementedException();
+                    return AsyncThunkILEmitter.EmitAsyncMethodThunk(asyncVariantImpl, asyncVariantImpl.Target);
                 }
             }
             else

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncThunks.cs

@@ -36,5 +36,19 @@ public static MethodIL EmitTaskReturningThunk(MethodDesc taskReturningMethod, Me
 
             return emitter.Link(taskReturningMethod);
         }
+
+        public static MethodIL EmitAsyncMethodThunk(MethodDesc asyncMethod, MethodDesc taskReturningMethod)
+        {
+            TypeSystemContext context = asyncMethod.Context;
+
+            var emitter = new ILEmitter();
+            var codestream = emitter.NewCodeStream();
+
+            // TODO: match EmitAsyncMethodThunk in CoreCLR VM
+
+            codestream.EmitCallThrowHelper(emitter, context.GetHelperEntryPoint("ThrowHelpers"u8, "ThrowNotSupportedException"u8));
+
+            return emitter.Link(asyncMethod);
+        }
     }
 }

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/CompilerTypeSystemContext.Aot.cs

@@ -39,6 +39,8 @@ public SharedGenericsConfiguration GenericsConfig
         private readonly Int128FieldLayoutAlgorithm _int128FieldLayoutAlgorithm;
         private readonly TypeWithRepeatedFieldsFieldLayoutAlgorithm _typeWithRepeatedFieldsFieldLayoutAlgorithm;
 
+        private readonly AsyncAwareVirtualMethodResolutionAlgorithm _virtualMethodAlgorithm;
+
         private TypeDesc[] _arrayOfTInterfaces;
         private TypeDesc[] _arrayEnumeratorOfTInterfaces;
         private ArrayOfTRuntimeInterfacesAlgorithm _arrayOfTRuntimeInterfacesAlgorithm;
@@ -53,6 +55,8 @@ public CompilerTypeSystemContext(TargetDetails details, SharedGenericsMode gener
         {
             _genericsMode = genericsMode;
 
+            _virtualMethodAlgorithm = new AsyncAwareVirtualMethodResolutionAlgorithm(this);
+
             _vectorOfTFieldLayoutAlgorithm = new VectorOfTFieldLayoutAlgorithm(_metadataFieldLayoutAlgorithm);
             _vectorFieldLayoutAlgorithm = new VectorFieldLayoutAlgorithm(_metadataFieldLayoutAlgorithm);
             _int128FieldLayoutAlgorithm = new Int128FieldLayoutAlgorithm(_metadataFieldLayoutAlgorithm);

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/ReadyToRunCompilerContext.cs

@@ -11,6 +11,8 @@ namespace ILCompiler
 {
     partial class CompilerTypeSystemContext
     {
+        private readonly MetadataVirtualMethodAlgorithm _virtualMethodAlgorithm = new MetadataVirtualMethodAlgorithm();
+
         public CompilerTypeSystemContext(TargetDetails details, SharedGenericsMode genericsMode)
             : base(details)
         {