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

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

@@ -51,7 +51,9 @@ public static MethodIL EmitTaskReturningThunk(MethodDesc taskReturningMethod, Me
             {
                 codestream.BeginTry(tryFinallyRegion);
                 codestream.Emit(ILOpcode.nop);
-                ILExceptionRegionBuilder tryCatchRegion = emitter.NewCatchRegion(context.GetWellKnownType(WellKnownType.Object));
+
+                TypeDesc exceptionType = context.GetWellKnownType(WellKnownType.Exception);
+                ILExceptionRegionBuilder tryCatchRegion = emitter.NewCatchRegion(exceptionType);
                 {
                     codestream.BeginTry(tryCatchRegion);
 
@@ -148,8 +150,6 @@ public static MethodIL EmitTaskReturningThunk(MethodDesc taskReturningMethod, Me
                 {
                     codestream.BeginHandler(tryCatchRegion);
 
-                    TypeDesc exceptionType = context.GetWellKnownType(WellKnownType.Exception);
-
                     MethodDesc fromExceptionMd;
                     if (logicalReturnType != null)
                     {

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRun/ModuleTokenResolver.cs

@@ -7,6 +7,7 @@
 using System.Reflection.Metadata;
 using System.Reflection.Metadata.Ecma335;
 
+using Internal.IL;
 using Internal.JitInterface;
 using Internal.TypeSystem;
 using Internal.TypeSystem.Ecma;
@@ -50,6 +51,7 @@ public void SetModuleIndexLookup(Func<IEcmaModule, int> moduleIndexLookup)
         public void InitManifestMutableModule(MutableModule mutableModule)
         {
             _manifestMutableModule = mutableModule;
+            InitializeKnownMethodsAndTypes();
         }
 
         public ModuleToken GetModuleTokenForType(EcmaType type, bool allowDynamicallyCreatedReference, bool throwIfNotFound = true)
@@ -436,5 +438,123 @@ public DummyTypeInfo GetTypeFromSpecification(MetadataReader reader, ModuleToken
                 return DummyTypeInfo.Instance;
             }
         }
+
+        private ImmutableHashSet<MethodDesc> KnownMethods;
+        private ImmutableHashSet<TypeDesc> KnownTypes;
+
+        private void InitializeKnownMethodsAndTypes()
+        {
+            var knownTypes = ImmutableHashSet.CreateBuilder<TypeDesc>();
+            var knownMethods = ImmutableHashSet.CreateBuilder<MethodDesc>();
+
+            var exception = CompilerContext.SystemModule.GetKnownType("System"u8, "Exception"u8);
+            knownTypes.Add(exception);
+
+            // AsyncHelpers type and methods
+            var asyncHelpers = CompilerContext.SystemModule.GetKnownType("System.Runtime.CompilerServices"u8, "AsyncHelpers"u8);
+            knownTypes.Add(asyncHelpers);
+
+            knownMethods.Add(asyncHelpers.GetKnownMethod("AsyncCallContinuation"u8, null));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("TaskFromException"u8, new MethodSignature(MethodSignatureFlags.Static, 0, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task"u8), new TypeDesc[] { exception })));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("ValueTaskFromException"u8, new MethodSignature(MethodSignatureFlags.Static, 0, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask"u8), new TypeDesc[] { exception })));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("TaskFromException"u8, new MethodSignature(MethodSignatureFlags.Static, 1, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task`1"u8).MakeInstantiatedType(CompilerContext.GetSignatureVariable(0, method: true)), new TypeDesc[] { exception })));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("ValueTaskFromException"u8, new MethodSignature(MethodSignatureFlags.Static, 1, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask`1"u8).MakeInstantiatedType(CompilerContext.GetSignatureVariable(0, method: true)), new TypeDesc[] { exception })));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("FinalizeTaskReturningThunk"u8, new MethodSignature(MethodSignatureFlags.Static, 0, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task"u8), Array.Empty<TypeDesc>())));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("FinalizeValueTaskReturningThunk"u8, new MethodSignature(MethodSignatureFlags.Static, 0, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask"u8), Array.Empty<TypeDesc>())));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("FinalizeTaskReturningThunk"u8, new MethodSignature(MethodSignatureFlags.Static, 1, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task`1"u8).MakeInstantiatedType(CompilerContext.GetSignatureVariable(0, method: true)), Array.Empty<TypeDesc>())));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("FinalizeValueTaskReturningThunk"u8, new MethodSignature(MethodSignatureFlags.Static, 1, CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask`1"u8).MakeInstantiatedType(CompilerContext.GetSignatureVariable(0, method: true)), Array.Empty<TypeDesc>())));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("TransparentAwait"u8, null));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("CompletedTask"u8, null));
+            knownMethods.Add(asyncHelpers.GetKnownMethod("CompletedTaskResult"u8, null));
+
+            // ExecutionAndSyncBlockStore type and methods
+            var executionSyncBlockStore = CompilerContext.SystemModule.GetKnownType("System.Runtime.CompilerServices"u8, "ExecutionAndSyncBlockStore"u8);
+            knownTypes.Add(executionSyncBlockStore);
+            knownMethods.Add(executionSyncBlockStore.GetKnownMethod("Push"u8, null));
+            knownMethods.Add(executionSyncBlockStore.GetKnownMethod("Pop"u8, null));
+
+            // Task types and methods
+            var task = CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task"u8);
+            knownTypes.Add(task);
+            knownMethods.Add(task.GetKnownMethod("FromResult"u8, null));
+            knownMethods.Add(task.GetKnownMethod("get_CompletedTask"u8, null));
+            knownMethods.Add(task.GetKnownMethod("get_IsCompleted"u8, null));
+
+            var taskGeneric = CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "Task`1"u8);
+            knownTypes.Add(taskGeneric);
+            knownMethods.Add(taskGeneric.GetKnownMethod("get_Result"u8, null));
+
+            // ValueTask types and methods
+            var valueTask = CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask"u8);
+            knownTypes.Add(valueTask);
+            knownMethods.Add(valueTask.GetKnownMethod("FromResult"u8, null));
+            knownMethods.Add(valueTask.GetKnownMethod("get_CompletedTask"u8, null));
+            knownMethods.Add(valueTask.GetKnownMethod("get_IsCompleted"u8, null));
+            knownMethods.Add(valueTask.GetKnownMethod("ThrowIfCompletedUnsuccessfully"u8, null));
+            knownMethods.Add(valueTask.GetKnownMethod("AsTaskOrNotifier"u8, null));
+
+            var valueTaskGeneric = CompilerContext.SystemModule.GetKnownType("System.Threading.Tasks"u8, "ValueTask`1"u8);
+            knownTypes.Add(valueTaskGeneric);
+            knownMethods.Add(valueTaskGeneric.GetKnownMethod("get_Result"u8, null));
+            knownMethods.Add(valueTaskGeneric.GetKnownMethod("AsTaskOrNotifier"u8, null));
+
+            KnownTypes = knownTypes.ToImmutable();
+            KnownMethods = knownMethods.ToImmutable();
+
+            try
+            {
+                Debug.Assert(_manifestMutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences == null);
+                _manifestMutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences = _manifestMutableModule.Context.SystemModule;
+                _manifestMutableModule.AddingReferencesToR2RKnownTypesAndMethods = true;
+                foreach (var type in KnownTypes)
+                {
+                    // We can skip types that are already in the version bubble
+                    if (_compilationModuleGroup.VersionsWithType(type))
+                        continue;
+
+                    EntityHandle? handle = _manifestMutableModule.TryGetEntityHandle(type);
+                    if (!handle.HasValue)
+                        throw new NotImplementedException($"Entity handle for known type ({type}) not found in manifest module");
+
+                    var token = new ModuleToken(_manifestMutableModule, handle.Value);
+                    this.AddModuleTokenForType(type, token);
+                    Debug.Assert(!default(ModuleToken).Equals(this.GetModuleTokenForType((EcmaType)type, allowDynamicallyCreatedReference: true, throwIfNotFound: true)));
+                }
+
+                foreach (var method in KnownMethods)
+                {
+                    // We can skip methods that are already in the version bubble
+                    if (_compilationModuleGroup.VersionsWithMethodBody((EcmaMethod)method))
+                        continue;
+
+                    EntityHandle ? handle = _manifestMutableModule.TryGetEntityHandle(method);
+                    if (!handle.HasValue)
+                        throw new NotImplementedException($"Entity handle for known method ({method}) not found in manifest module");
+
+                    var token = new ModuleToken(_manifestMutableModule, handle.Value);
+                    this.AddModuleTokenForMethod(method, token);
+                    Debug.Assert(!default(ModuleToken).Equals(this.GetModuleTokenForMethod(method, allowDynamicallyCreatedReference: true, throwIfNotFound: true)));
+                }
+            }
+            finally
+            {
+                _manifestMutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences = null;
+                _manifestMutableModule.AddingReferencesToR2RKnownTypesAndMethods = false;
+            }
+        }
+
+        public bool IsKnownMutableModuleMethod(EcmaMethod method)
+        {
+            if (KnownMethods is null)
+                InitializeKnownMethodsAndTypes();
+            return KnownMethods.Contains(method);
+        }
+
+        public bool IsKnownMutableModuleType(EcmaType type)
+        {
+            if (KnownTypes is null)
+                InitializeKnownMethodsAndTypes();
+            return KnownTypes.Contains(type);
+        }
     }
 }

src/coreclr/tools/aot/ILCompiler.ReadyToRun/Compiler/DependencyAnalysis/ReadyToRunCodegenNodeFactory.cs

@@ -713,7 +713,7 @@ public void AttachToDependencyGraph(DependencyAnalyzerBase<NodeFactory> graph, I
             ManifestMetadataTable = new ManifestMetadataTableNode(this);
             Header.Add(Internal.Runtime.ReadyToRunSectionType.ManifestMetadata, ManifestMetadataTable, ManifestMetadataTable);
             Resolver.SetModuleIndexLookup(ManifestMetadataTable.ModuleToIndex);
-            ((ReadyToRunILProvider)ilProvider).InitManifestMutableModule(ManifestMetadataTable._mutableModule);
+            ((ReadyToRunILProvider)ilProvider).InitManifestMutableModule(ManifestMetadataTable._mutableModule, Resolver);
             Resolver.InitManifestMutableModule(ManifestMetadataTable._mutableModule);
 
             ManifestAssemblyMvidHeaderNode mvidTableNode = new ManifestAssemblyMvidHeaderNode(ManifestMetadataTable);

src/coreclr/tools/aot/ILCompiler.ReadyToRun/IL/ReadyToRunILProvider.cs

@@ -14,6 +14,7 @@
 using System.Buffers.Binary;
 using System.Reflection.Metadata;
 using System.Reflection.Metadata.Ecma335;
+using ILCompiler.DependencyAnalysis.ReadyToRun;
 
 namespace Internal.IL
 {
@@ -34,7 +35,7 @@ public ReadyToRunILProvider(CompilationModuleGroup compilationModuleGroup)
             _compilationModuleGroup = compilationModuleGroup;
         }
 
-        public void InitManifestMutableModule(MutableModule module)
+        public void InitManifestMutableModule(MutableModule module, ModuleTokenResolver resolver)
         {
             _manifestMutableModule = module;
         }
@@ -169,7 +170,6 @@ public override MethodIL GetMethodIL(MethodDesc method)
                 {
                     // We should not be creating any AsyncMethodVariants yet.
                     // This hasn't been implemented.
-                    Debug.Assert(!method.IsAsyncVariant());
                     return null;
                 }
 

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

@@ -45,6 +45,7 @@
     <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\Stubs\AsyncThunks.cs" Link="IL\Stubs\AsyncThunks.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\ComparerIntrinsics.cs" Link="IL\Stubs\ComparerIntrinsics.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\InterlockedIntrinsics.cs" Link="IL\Stubs\InterlockedIntrinsics.cs" />
     <Compile Include="..\..\Common\TypeSystem\IL\Stubs\RuntimeHelpersIntrinsics.cs" Link="IL\Stubs\RuntimeHelpersIntrinsics.cs" />
@@ -327,6 +328,9 @@
     <Compile Include="..\..\Common\TypeSystem\IL\HelperExtensions.cs">
       <Link>IL\HelperExtensions.cs</Link>
     </Compile>
+    <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.cs" Link="IL\Stubs\AsyncResumptionStub.cs" />
+    <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.Mangling.cs" Link="IL\Stubs\AsyncResumptionStub.Mangling.cs" />
+    <Compile Include="..\..\Common\TypeSystem\IL\Stubs\AsyncResumptionStub.Sorting.cs" Link="IL\Stubs\AsyncResumptionStub.Sorting.cs" />
     <Compile Include="..\..\Common\JitInterface\JitTypeNameFormatter.cs">
       <Link>JitInterface\JitTypeNameFormatter.cs</Link>
     </Compile>

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

@@ -756,6 +756,12 @@ public void CompileMethod(MethodWithGCInfo methodCodeNodeNeedingCode, Logger log
                         logger.Writer.WriteLine($"Info: Method `{MethodBeingCompiled}` was not compiled because it has an non version resilient signature.");
                     return;
                 }
+                if (MethodBeingCompiled.IsAsyncVariant())
+                {
+                    if (logger.IsVerbose)
+                        logger.Writer.WriteLine($"Info: Method `{MethodBeingCompiled}` was not compiled because async variants are not supported yet.");
+                    return;
+                }
                 MethodIL methodIL = _compilation.GetMethodIL(MethodBeingCompiled);
                 if (methodIL == null)
                 {
@@ -1412,7 +1418,12 @@ private ModuleToken HandleToModuleToken(ref CORINFO_RESOLVED_TOKEN pResolvedToke
                     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))
+                    {
+                        Debug.Assert(_compilation.NodeFactory.Resolver.IsKnownMutableModuleMethod((EcmaMethod)resultMethod));
+                        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 +1443,12 @@ 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))
+                        {
+                            Debug.Assert(_compilation.NodeFactory.Resolver.IsKnownMutableModuleType(ecmaType));
+                            ModuleToken result = _compilation.NodeFactory.Resolver.GetModuleTokenForType(ecmaType, allowDynamicallyCreatedReference: true, throwIfNotFound: true);
+                            return result;
+                        }
                         token = (mdToken)MetadataTokens.GetToken(ecmaType.Handle);
                         module = ecmaType.Module;
                     }

src/coreclr/tools/aot/ILCompiler.ReadyToRun/TypeSystem/Mutable/MutableModule.cs

@@ -39,7 +39,8 @@ protected override EntityHandle GetNonNestedResolutionScope(MetadataType metadat
                 if (!_mutableModule._moduleToModuleRefString.TryGetValue(module, out moduleRefString))
                 {
                     Debug.Assert(_mutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences != null &&
-                        _mutableModule._compilationGroup.CrossModuleInlineableModule(_mutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences));
+                        _mutableModule._compilationGroup.CrossModuleInlineableModule(_mutableModule.ModuleThatIsCurrentlyTheSourceOfNewReferences)
+                        || _mutableModule.AddingReferencesToR2RKnownTypesAndMethods);
 
                     if (module == _typeSystemContext.SystemModule)
                     {
@@ -303,6 +304,7 @@ class DisableNewTokensException : Exception { }
 
         public bool DisableNewTokens;
         public ModuleDesc ModuleThatIsCurrentlyTheSourceOfNewReferences;
+        public bool AddingReferencesToR2RKnownTypesAndMethods;
         private ReadyToRunCompilationModuleGroupBase _compilationGroup;
         private Dictionary<ModuleDesc, string> _moduleToModuleRefString = new Dictionary<ModuleDesc, string>();