Insert a set of known types and methods references necessary for async thunks into the mutable module

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

@@ -78,5 +78,10 @@ public static bool IsAsyncVariant(this MethodDesc method)
         {
             return method.GetTypicalMethodDefinition() is AsyncMethodVariant;
         }
+
+        public static bool IsAsyncThunk(this MethodDesc method)
+        {
+            return method.IsAsyncVariant() ^ method.IsAsync;
+        }
     }
 }

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

@@ -0,0 +1,118 @@
+// 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.Collections.Generic;
+using System.Diagnostics;
+using ILCompiler;
+using Internal.TypeSystem;
+using Internal.TypeSystem.Ecma;
+
+namespace Internal.IL
+{
+    public class ILStubReferences
+    {
+        /// <summary>
+        /// Extracts all method and type references from the IL of a method.
+        /// This is used to find references in synthetic IL (like async thunks) that may need
+        /// to be added to the mutable module.
+        /// </summary>
+        public static List<TypeSystemEntity> GetNecessaryReferencesFromIL(MethodIL methodIL)
+        {
+            var references = new List<TypeSystemEntity>();
+            byte[] ilBytes = methodIL.GetILBytes();
+            ILReader reader = new ILReader(ilBytes);
+
+            while (reader.HasNext)
+            {
+                ILOpcode opcode = reader.ReadILOpcode();
+
+                switch (opcode)
+                {
+                    case ILOpcode.call:
+                    case ILOpcode.callvirt:
+                    case ILOpcode.newobj:
+                    case ILOpcode.ldftn:
+                    case ILOpcode.ldvirtftn:
+                    case ILOpcode.jmp:
+                    {
+                        int token = reader.ReadILToken();
+                        object obj = methodIL.GetObject(token, NotFoundBehavior.ReturnNull);
+                        if (obj is MethodDesc method)
+                        {
+                            references.Add(method);
+                        }
+                        break;
+                    }
+
+                    case ILOpcode.newarr:
+                    case ILOpcode.castclass:
+                    case ILOpcode.isinst:
+                    case ILOpcode.box:
+                    case ILOpcode.unbox:
+                    case ILOpcode.unbox_any:
+                    case ILOpcode.ldobj:
+                    case ILOpcode.stobj:
+                    case ILOpcode.cpobj:
+                    case ILOpcode.initobj:
+                    case ILOpcode.ldelem:
+                    case ILOpcode.stelem:
+                    case ILOpcode.ldelema:
+                    case ILOpcode.constrained:
+                    case ILOpcode.sizeof_:
+                    case ILOpcode.mkrefany:
+                    case ILOpcode.refanyval:
+                    {
+                        int token = reader.ReadILToken();
+                        object obj = methodIL.GetObject(token, NotFoundBehavior.ReturnNull);
+                        if (obj is TypeDesc type)
+                        {
+                            references.Add(type);
+                        }
+                        break;
+                    }
+
+                    case ILOpcode.ldtoken:
+                    {
+                        int token = reader.ReadILToken();
+                        object obj = methodIL.GetObject(token, NotFoundBehavior.ReturnNull);
+                        if (obj is MethodDesc method)
+                        {
+                            references.Add(method);
+                        }
+                        else if (obj is TypeDesc type)
+                        {
+                            references.Add(type);
+                        }
+                        break;
+                    }
+
+                    default:
+                        reader.Skip(opcode);
+                        break;
+                }
+            }
+
+            foreach (var region in methodIL.GetExceptionRegions())
+            {
+                if (region.Kind == ILExceptionRegionKind.Catch && region.ClassToken != 0)
+                {
+                    object obj = methodIL.GetObject(region.ClassToken, NotFoundBehavior.ReturnNull);
+                    if (obj is TypeDesc type)
+                        references.Add(type);
+                }
+            }
+
+            foreach (var local in methodIL.GetLocals())
+            {
+                TypeDesc type = local.Type;
+                // These types can be represented by ELEMENT_TYPE values and don't need references.
+                if (!type.IsPrimitive && !type.IsVoid && !type.IsObject && !type.IsString && !type.IsTypedReference)
+                    references.Add(type);
+            }
+
+            return references;
+        }
+    }
+}

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

@@ -108,6 +108,12 @@ public bool CanInline(MethodDesc caller, MethodDesc callee)
                 }
             }
 
+            if (callee.IsAsyncThunk())
+            {
+                // Async thunks require special handling in the compiler and should not be inlined
+                return false;
+            }
+
             _nodeFactory.DetectGenericCycles(caller, callee);
 
             return NodeFactory.CompilationModuleGroup.CanInline(caller, callee);
@@ -694,6 +700,50 @@ protected override void ComputeDependencyNodeDependencies(List<DependencyNodeCor
                                 CorInfoImpl.IsMethodCompilable(this, methodCodeNodeNeedingCode.Method))
                                 _methodsWhichNeedMutableILBodies.Add(ecmaMethod);
                         }
+                        if (method.IsAsyncThunk())
+                        {
+                            // The synthetic async thunks require references to methods/types that may not have existing methodRef entries in the version bubble.
+                            // These references need to be added to the mutable module if they don't exist.
+                            // Extract the required references by reading the IL of the thunk method.
+                            MethodIL methodIL = GetMethodIL(method);
+                            if (methodIL is null)
+                                continue;
+
+                            _nodeFactory.ManifestMetadataTable._mutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences = ((EcmaMethod)method.GetTypicalMethodDefinition()).Module;
+                            try
+                            {
+                                foreach (var entity in ILStubReferences.GetNecessaryReferencesFromIL(methodIL))
+                                {
+                                    switch (entity)
+                                    {
+                                        case MethodDesc md:
+                                            if (md.IsAsyncVariant())
+                                            {
+                                                Debug.Assert(((CompilerTypeSystemContext)md.Context).GetTargetOfAsyncVariantMethod(md.GetTypicalMethodDefinition()) == method);
+                                                Debug.Assert(!_nodeFactory.Resolver.GetModuleTokenForMethod(method, allowDynamicallyCreatedReference: false, throwIfNotFound: true).IsNull);
+                                            }
+                                            else if (_nodeFactory.Resolver.GetModuleTokenForMethod(md, allowDynamicallyCreatedReference: true, throwIfNotFound: false).IsNull)
+                                            {
+                                                _nodeFactory.ManifestMetadataTable._mutableModule.TryGetEntityHandle(md);
+                                                Debug.Assert(!_nodeFactory.Resolver.GetModuleTokenForMethod(md, allowDynamicallyCreatedReference: true, throwIfNotFound: true).IsNull);
+                                            }
+                                            break;
+                                        case EcmaType td:
+                                            if (_nodeFactory.Resolver.GetModuleTokenForType(td, allowDynamicallyCreatedReference: true, throwIfNotFound: false).IsNull)
+                                            {
+                                                _nodeFactory.ManifestMetadataTable._mutableModule.TryGetEntityHandle(td);
+                                            }
+                                            Debug.Assert(!_nodeFactory.Resolver.GetModuleTokenForType(td, allowDynamicallyCreatedReference: true, throwIfNotFound: true).IsNull);
+                                            break;
+                                    }
+                                }
+                            }
+                            finally
+                            {
+                                _nodeFactory.ManifestMetadataTable._mutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences = null;
+                            }
+                        }
+
                         if (!_nodeFactory.CompilationModuleGroup.VersionsWithMethodBody(method))
                         {
                             // Validate that the typedef tokens for all of the instantiation parameters of the method
@@ -977,5 +1027,6 @@ public string GetReproInstructions(MethodDesc method)
         {
             return _printReproInstructions(method);
         }
+
     }
 }

src/coreclr/tools/aot/ILCompiler.ReadyToRun/ILCompiler.ReadyToRun.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <OutputType>Library</OutputType>
     <AssemblyName>ILCompiler.ReadyToRun</AssemblyName>
@@ -45,6 +45,8 @@
     <Compile Include="..\..\Common\System\FormattingHelpers.cs" Link="Common\FormattingHelpers.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\ILProvider.cs" Link="IL\ILProvider.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\ILReader.cs" Link="IL\ILReader.cs" />
+    <Compile Include="..\..\Common\TypeSystem\IL\ILStubReferences.cs" Link="IL\ILStubReferences.cs" />
+    <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncThunks.cs" Link="IL\Stubs\AsyncThunks.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.cs" Link="IL\Stubs\AsyncResumptionStub.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.Sorting.cs" Link="IL\Stubs\AsyncResumptionStub.Sorting.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.Mangling.cs" Link="IL\Stubs\AsyncResumptionStub.Mangling.cs" />

src/coreclr/tools/aot/ILCompiler.ReadyToRun/JitInterface/CorInfoImpl.ReadyToRun.cs

@@ -1405,14 +1405,20 @@ private ModuleToken HandleToModuleToken(ref CORINFO_RESOLVED_TOKEN pResolvedToke
                 {
                     if (resultMethod is IL.Stubs.PInvokeTargetNativeMethod rawPinvoke)
                         resultMethod = rawPinvoke.Target;
+                    if (resultMethod is AsyncMethodVariant asyncVariant)
+                        resultMethod = asyncVariant.Target;
 
                     // It's okay to strip the instantiation away because we don't need a MethodSpec
                     // token - SignatureBuilder will generate the generic method signature
                     // using instantiation parameters from the MethodDesc entity.
                     resultMethod = resultMethod.GetTypicalMethodDefinition();
 
                     Debug.Assert(resultMethod is EcmaMethod);
-                    Debug.Assert(_compilation.NodeFactory.CompilationModuleGroup.VersionsWithType(((EcmaMethod)resultMethod).OwningType));
+                    if (!_compilation.NodeFactory.CompilationModuleGroup.VersionsWithType(((EcmaMethod)resultMethod).OwningType))
+                    {
+                        ModuleToken result = _compilation.NodeFactory.Resolver.GetModuleTokenForMethod(resultMethod, allowDynamicallyCreatedReference: true, throwIfNotFound: true);
+                        return result;
+                    }
                     token = (mdToken)MetadataTokens.GetToken(((EcmaMethod)resultMethod).Handle);
                     module = ((EcmaMethod)resultMethod).Module;
                 }
@@ -1432,7 +1438,11 @@ private ModuleToken HandleToModuleToken(ref CORINFO_RESOLVED_TOKEN pResolvedToke
                 {
                     if (resultDef is EcmaType ecmaType)
                     {
-                        Debug.Assert(_compilation.NodeFactory.CompilationModuleGroup.VersionsWithType(ecmaType));
+                        if (!_compilation.NodeFactory.CompilationModuleGroup.VersionsWithType(ecmaType))
+                        {
+                            ModuleToken result = _compilation.NodeFactory.Resolver.GetModuleTokenForType(ecmaType, allowDynamicallyCreatedReference: true, throwIfNotFound: true);
+                            return result;
+                        }
                         token = (mdToken)MetadataTokens.GetToken(ecmaType.Handle);
                         module = ecmaType.Module;
                     }