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

Async thunk IL has references to methods and types that may not have MethodRef
or TypeRef tokens in the assembly being crossgen'ed. There are some
assertions to validate that any references in "faux il" version with the
assembly, which isn't true for these types and methods. To work around
this, we can inject these references into the mutable module, a
synthetic module added to the assembly for things like this. When we are about
to compile an async thunk method, we check to see if all the required
references are already in the mutable module, and add them if not.

Author: jtschuster

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/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRun/ModuleTokenResolver.cs

@@ -52,17 +52,20 @@ public void InitManifestMutableModule(MutableModule mutableModule)
             _manifestMutableModule = mutableModule;
         }
 
-        public ModuleToken GetModuleTokenForType(EcmaType type, bool allowDynamicallyCreatedReference, bool throwIfNotFound = true)
+        public ModuleToken GetModuleTokenForType(TypeDesc type, bool allowDynamicallyCreatedReference, bool throwIfNotFound = true)
         {
-            if (_compilationModuleGroup.VersionsWithType(type))
-            {
-                return new ModuleToken(type.Module, (mdToken)MetadataTokens.GetToken(type.Handle));
-            }
-
             ModuleToken token;
-            if (_typeToRefTokens.TryGetValue(type, out token))
+            if (type is EcmaType ecmaType)
             {
-                return token;
+                if (_compilationModuleGroup.VersionsWithType(ecmaType))
+                {
+                    return new ModuleToken(ecmaType.Module, (mdToken)MetadataTokens.GetToken(ecmaType.Handle));
+                }
+
+                if (_typeToRefTokens.TryGetValue(ecmaType, out token))
+                {
+                    return token;
+                }
             }
 
             // If the token was not lazily mapped, search the input compilation set for a type reference token
@@ -178,7 +181,7 @@ private void DecodeMethodSignatureToDiscoverUsedTypeTokens(BlobHandle signatureH
         {
             MetadataReader metadataReader = token.MetadataReader;
             TokenResolverProvider rentedProvider = TokenResolverProvider.Rent(this, token.Module);
-            SignatureDecoder<DummyTypeInfo, ModuleTokenResolver> sigDecoder = new (rentedProvider, metadataReader, this);
+            SignatureDecoder<DummyTypeInfo, ModuleTokenResolver> sigDecoder = new(rentedProvider, metadataReader, this);
             BlobReader signature = metadataReader.GetBlobReader(signatureHandle);
 
             SignatureHeader header = signature.ReadSignatureHeader();

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,52 @@ 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 TypeDesc 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;
+                                        default:
+                                            throw new NotImplementedException("Unexpected entity type in Async thunk.");
+                                    }
+                                }
+                            }
+                            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 +1029,6 @@ public string GetReproInstructions(MethodDesc method)
         {
             return _printReproInstructions(method);
         }
+
     }
 }