Long-lived AsyncMethodVariant in MethodWithGCInfo/MethodWithToken (#121357)

Incorporates feedback from
#121218, creating a long-lived
async MethodDesc. There will be more warts that need to be worked out
once we actually compile these, in particular lots of
`(EcmaMethod)methodDesc.GetMethodDefinition()`, which doesn't work with
`AsyncMethodVariant`.

Author: jtschuster

src/coreclr/inc/corinfo.h

@@ -673,7 +673,8 @@ enum CorInfoCallConv
     CORINFO_CALLCONV_HASTHIS    = 0x20,
     CORINFO_CALLCONV_EXPLICITTHIS=0x40,
     CORINFO_CALLCONV_PARAMTYPE  = 0x80,     // Passed last. Same as CORINFO_GENERICS_CTXT_FROM_PARAMTYPEARG
-    CORINFO_CALLCONV_ASYNCCALL  = 0x100,    // Is this a call to an async function?
+    CORINFO_CALLCONV_ASYNCCALL  = 0x100,    // Is this a call with async calling convention?
+
 };
 
 // Represents the calling conventions supported with the extensible calling convention syntax

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

@@ -0,0 +1,90 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics;
+using Internal.TypeSystem;
+using Internal.TypeSystem.Ecma;
+
+namespace ILCompiler
+{
+    /// <summary>
+    /// MethodDesc that represents async calling convention entrypoint of a Task-returning method.
+    /// </summary>
+    public partial class AsyncMethodVariant : MethodDelegator
+    {
+        private MethodSignature _asyncSignature;
+
+        public AsyncMethodVariant(EcmaMethod wrappedMethod)
+            : base(wrappedMethod)
+        {
+            Debug.Assert(wrappedMethod.Signature.ReturnsTaskOrValueTask());
+        }
+
+        public EcmaMethod Target => (EcmaMethod)_wrappedMethod;
+
+        public override MethodSignature Signature
+        {
+            get
+            {
+                return _asyncSignature ?? InitializeSignature();
+            }
+        }
+
+        private MethodSignature InitializeSignature()
+        {
+            var signature = _wrappedMethod.Signature;
+            Debug.Assert(!signature.IsAsyncCall);
+            Debug.Assert(signature.ReturnsTaskOrValueTask());
+            TypeDesc md = signature.ReturnType;
+            MethodSignatureBuilder builder = new MethodSignatureBuilder(signature);
+            builder.ReturnType = md.HasInstantiation ? md.Instantiation[0] : this.Context.GetWellKnownType(WellKnownType.Void);
+            builder.Flags = signature.Flags | MethodSignatureFlags.AsyncCall;
+            return (_asyncSignature = builder.ToSignature());
+        }
+
+        public override MethodDesc GetCanonMethodTarget(CanonicalFormKind kind)
+        {
+            // We should not be calling GetCanonMethodTarget on generic definitions of anything
+            // and this MethodDesc is a generic definition.
+            Debug.Assert(!HasInstantiation && !OwningType.HasInstantiation);
+            return this;
+        }
+
+        public override MethodDesc GetMethodDefinition()
+        {
+            return this;
+        }
+
+        public override MethodDesc GetTypicalMethodDefinition()
+        {
+            return this;
+        }
+
+        public override MethodDesc InstantiateSignature(Instantiation typeInstantiation, Instantiation methodInstantiation)
+        {
+            throw new NotImplementedException();
+        }
+
+        public override string ToString() => $"Async variant: " + _wrappedMethod.ToString();
+
+        protected override int ClassCode => unchecked((int)0xd0fd1c1fu);
+
+        protected override int CompareToImpl(MethodDesc other, TypeSystemComparer comparer)
+        {
+            var asyncOther = (AsyncMethodVariant)other;
+            return comparer.Compare(_wrappedMethod, asyncOther._wrappedMethod);
+        }
+    }
+
+    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;
+        }
+    }
+}

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

@@ -0,0 +1,44 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Internal.IL;
+using Internal.TypeSystem;
+using Internal.TypeSystem.Ecma;
+
+using Debug = System.Diagnostics.Debug;
+
+namespace ILCompiler
+{
+    public partial class CompilerTypeSystemContext
+    {
+        public MethodDesc GetAsyncVariantMethod(MethodDesc taskReturningMethod)
+        {
+            Debug.Assert(taskReturningMethod.Signature.ReturnsTaskOrValueTask());
+            MethodDesc asyncMetadataMethodDef = taskReturningMethod.GetTypicalMethodDefinition();
+            MethodDesc result = _asyncVariantImplHashtable.GetOrCreateValue((EcmaMethod)asyncMetadataMethodDef);
+
+            if (asyncMetadataMethodDef != taskReturningMethod)
+            {
+                TypeDesc owningType = taskReturningMethod.OwningType;
+                if (owningType.HasInstantiation)
+                    result = GetMethodForInstantiatedType(result, (InstantiatedType)owningType);
+
+                if (taskReturningMethod.HasInstantiation)
+                    result = GetInstantiatedMethod(result, taskReturningMethod.Instantiation);
+            }
+
+            return result;
+        }
+
+        private sealed class AsyncVariantImplHashtable : LockFreeReaderHashtable<EcmaMethod, AsyncMethodVariant>
+        {
+            protected override int GetKeyHashCode(EcmaMethod key) => key.GetHashCode();
+            protected override int GetValueHashCode(AsyncMethodVariant value) => value.Target.GetHashCode();
+            protected override bool CompareKeyToValue(EcmaMethod key, AsyncMethodVariant value) => key == value.Target;
+            protected override bool CompareValueToValue(AsyncMethodVariant value1, AsyncMethodVariant value2)
+                => value1.Target == value2.Target;
+            protected override AsyncMethodVariant CreateValueFromKey(EcmaMethod key) => new AsyncMethodVariant(key);
+        }
+        private AsyncVariantImplHashtable _asyncVariantImplHashtable = new AsyncVariantImplHashtable();
+    }
+}

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

@@ -1,8 +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;
 using ILCompiler.DependencyAnalysis;
-
 using Internal.TypeSystem;
 using Internal.TypeSystem.Ecma;
 
@@ -135,5 +135,23 @@ public static bool NotCallableWithoutOwningEEType(this MethodDesc method)
                 (owningType is not MetadataType mdType || !mdType.IsModuleType) && /* Compiler parks some instance methods on the <Module> type */
                 !method.IsSharedByGenericInstantiations; /* Current impl limitation; can be lifted */
         }
+
+        public static bool ReturnsTaskOrValueTask(this MethodSignature method)
+        {
+            TypeDesc ret = method.ReturnType;
+
+            if (ret is MetadataType md
+                && md.Module == method.Context.SystemModule
+                && md.Namespace.SequenceEqual("System.Threading.Tasks"u8))
+            {
+                ReadOnlySpan<byte> name = md.Name;
+                if (name.SequenceEqual("Task"u8) || name.SequenceEqual("Task`1"u8)
+                    || name.SequenceEqual("ValueTask"u8) || name.SequenceEqual("ValueTask`1"u8))
+                {
+                    return true;
+                }
+            }
+            return false;
+        }
     }
 }

src/coreclr/tools/Common/JitInterface/AsyncMethodDesc.cs

@@ -877,6 +877,7 @@ private void Get_CORINFO_SIG_INFO(MethodSignature signature, CORINFO_SIG_INFO* s
 
             if (!signature.IsStatic) sig->callConv |= CorInfoCallConv.CORINFO_CALLCONV_HASTHIS;
             if (signature.IsExplicitThis) sig->callConv |= CorInfoCallConv.CORINFO_CALLCONV_EXPLICITTHIS;
+            if (signature.IsAsyncCall) sig->callConv |= CorInfoCallConv.CORINFO_CALLCONV_ASYNCCALL;
 
             TypeDesc returnType = signature.ReturnType;
 

src/coreclr/tools/Common/JitInterface/AsyncMethodDescFactory.cs

@@ -119,6 +119,7 @@ public unsafe struct CORINFO_SIG_INFO
         private uint totalILArgs() { return (uint)(numArgs + (hasImplicitThis() ? 1 : 0)); }
         private bool isVarArg() { return ((getCallConv() == CorInfoCallConv.CORINFO_CALLCONV_VARARG) || (getCallConv() == CorInfoCallConv.CORINFO_CALLCONV_NATIVEVARARG)); }
         internal bool hasTypeArg() { return ((callConv & CorInfoCallConv.CORINFO_CALLCONV_PARAMTYPE) != 0); }
+        private bool isAsyncCall() { return ((callConv & CorInfoCallConv.CORINFO_CALLCONV_ASYNCCALL) != 0); }
     };
 
     //----------------------------------------------------------------------------
@@ -377,6 +378,7 @@ public enum CorInfoCallConv
         CORINFO_CALLCONV_HASTHIS = 0x20,
         CORINFO_CALLCONV_EXPLICITTHIS = 0x40,
         CORINFO_CALLCONV_PARAMTYPE = 0x80,     // Passed last. Same as CORINFO_GENERICS_CTXT_FROM_PARAMTYPEARG
+        CORINFO_CALLCONV_ASYNCCALL = 0x100,    // Is this a call with async calling convention?
     }
 
     // Represents the calling conventions supported with the extensible calling convention syntax

src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs

@@ -23,7 +23,8 @@ public enum MethodSignatureFlags
 
         Static = 0x0010,
         ExplicitThis = 0x0020,
-        AsyncCallConv = 0x0040,
+
+        AsyncCall                            = 0x0100,
     }
 
     public enum EmbeddedSignatureDataKind
@@ -139,11 +140,11 @@ public bool IsExplicitThis
             }
         }
 
-        public bool IsAsyncCallConv
+        public bool IsAsyncCall
         {
             get
             {
-                return (_flags & MethodSignatureFlags.AsyncCallConv) != 0;
+                return (_flags & MethodSignatureFlags.AsyncCall) != 0;
             }
         }
 

src/coreclr/tools/Common/JitInterface/CorInfoTypes.cs

@@ -15,6 +15,7 @@ public sealed partial class EcmaMethod : MethodDesc, EcmaModule.IEntityHandleObj
     {
         private static class MethodFlags
         {
+#pragma warning disable IDE0055 // Disable formatting to keep aligned
             public const int BasicMetadataCache     = 0x00001;
             public const int Virtual                = 0x00002;
             public const int NewSlot                = 0x00004;
@@ -28,11 +29,12 @@ private static class MethodFlags
             public const int AggressiveOptimization = 0x00400;
             public const int NoOptimization         = 0x00800;
             public const int RequireSecObject       = 0x01000;
+            public const int Async                  = 0x02000;
 
-            public const int AttributeMetadataCache = 0x02000;
-            public const int Intrinsic              = 0x04000;
-            public const int UnmanagedCallersOnly   = 0x08000;
-            public const int Async                  = 0x10000;
+            public const int AttributeMetadataCache = 0x04000;
+            public const int Intrinsic              = 0x08000;
+            public const int UnmanagedCallersOnly   = 0x10000;
+#pragma warning restore IDE0055
         };
 
         private EcmaType _type;