Async continuations for native AOT

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -12,6 +12,10 @@
 using System.Threading;
 using System.Threading.Tasks;
 
+#if NATIVEAOT
+using Internal.Runtime;
+#endif
+
 namespace System.Runtime.CompilerServices
 {
     internal struct ExecutionAndSyncBlockStore
@@ -160,11 +164,16 @@ private struct RuntimeAsyncAwaitState
 
         private static unsafe Continuation AllocContinuation(Continuation prevContinuation, MethodTable* contMT)
         {
+#if NATIVEAOT
+            Continuation newContinuation = (Continuation)RuntimeImports.RhNewObject(contMT);
+#else
             Continuation newContinuation = (Continuation)RuntimeTypeHandle.InternalAllocNoChecks(contMT);
+#endif
             prevContinuation.Next = newContinuation;
             return newContinuation;
         }
 
+#if !NATIVEAOT
         private static unsafe Continuation AllocContinuationMethod(Continuation prevContinuation, MethodTable* contMT, int keepAliveOffset, MethodDesc* method)
         {
             LoaderAllocator loaderAllocator = RuntimeMethodHandle.GetLoaderAllocator(new RuntimeMethodHandleInternal((IntPtr)method));
@@ -186,6 +195,7 @@ private static unsafe Continuation AllocContinuationClass(Continuation prevConti
             }
             return newContinuation;
         }
+#endif
 
         [BypassReadyToRun]
         [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)]

src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj

@@ -55,6 +55,10 @@
       <Link>ExceptionStringID.cs</Link>
     </Compile>
   </ItemGroup>
+  <ItemGroup>
+    <!-- TODO: (async) once we know which helpers can actually be shared, move those to libraries partition -->
+    <Compile Include="$(CoreClrProjectRoot)System.Private.CoreLib\src\System\Runtime\CompilerServices\AsyncHelpers.CoreCLR.cs" />
+  </ItemGroup>
   <ItemGroup>
     <Compile Include="Internal\Runtime\CompilerHelpers\DelegateHelpers.cs" />
     <Compile Include="Internal\Runtime\CompilerHelpers\LibraryInitializer.cs" />

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

@@ -0,0 +1,85 @@
+// 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 System.Text;
+
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    /// <summary>
+    /// An async continuation type. The code generator will request this to store local state
+    /// through an async suspension/resumption. We only identify these using a <see cref="GCPointerMap"/>
+    /// since that's all the code generator cares about - size of the type, and where the GC pointers are.
+    /// </summary>
+    public sealed partial class AsyncContinuationType : MetadataType
+    {
+        private readonly MetadataType _continuationBaseType;
+        public GCPointerMap PointerMap { get; }
+
+        public override DefType[] ExplicitlyImplementedInterfaces => [];
+        public override ReadOnlySpan<byte> Name => Encoding.UTF8.GetBytes(DiagnosticName);
+        public override ReadOnlySpan<byte> Namespace => [];
+
+        // We don't lay these out using MetadataType metadata.
+        // Autolayout (which we'd get due to GC pointers) would likely not match what codegen expects.
+        public override bool IsExplicitLayout => throw new NotImplementedException();
+        public override bool IsSequentialLayout => throw new NotImplementedException();
+        public override bool IsExtendedLayout => throw new NotImplementedException();
+        public override bool IsAutoLayout => throw new NotImplementedException();
+        public override ClassLayoutMetadata GetClassLayout() => throw new NotImplementedException();
+
+        public override bool IsBeforeFieldInit => false;
+        public override ModuleDesc Module => _continuationBaseType.Module;
+        public override MetadataType BaseType => _continuationBaseType;
+        public override bool IsSealed => true;
+        public override bool IsAbstract => false;
+        public override MetadataType ContainingType => null;
+        public override PInvokeStringFormat PInvokeStringFormat => default;
+        public override string DiagnosticName => $"ContinuationType_{PointerMap}";
+        public override string DiagnosticNamespace => "";
+        protected override int ClassCode => 0x528741a;
+        public override TypeSystemContext Context => _continuationBaseType.Context;
+
+        public AsyncContinuationType(MetadataType continuationBaseType, GCPointerMap pointerMap)
+            => (_continuationBaseType, PointerMap) = (continuationBaseType, pointerMap);
+
+        public override bool HasCustomAttribute(string attributeNamespace, string attributeName) => false;
+        public override IEnumerable<MetadataType> GetNestedTypes() => [];
+        public override MetadataType GetNestedType(string name) => null;
+        protected override MethodImplRecord[] ComputeVirtualMethodImplsForType() => [];
+        public override MethodImplRecord[] FindMethodsImplWithMatchingDeclName(ReadOnlySpan<byte> name) => [];
+
+        protected override int CompareToImpl(TypeDesc other, TypeSystemComparer comparer)
+        {
+            Debug.Assert(_continuationBaseType == ((AsyncContinuationType)other)._continuationBaseType);
+            GCPointerMap otherPointerMap = ((AsyncContinuationType)other).PointerMap;
+            return PointerMap.CompareTo(otherPointerMap);
+        }
+
+        public override int GetHashCode() => PointerMap.GetHashCode();
+
+        protected override TypeFlags ComputeTypeFlags(TypeFlags mask)
+        {
+            TypeFlags flags = 0;
+
+            if ((mask & TypeFlags.HasGenericVarianceComputed) != 0)
+            {
+                flags |= TypeFlags.HasGenericVarianceComputed;
+            }
+
+            if ((mask & TypeFlags.CategoryMask) != 0)
+            {
+                flags |= TypeFlags.Class;
+            }
+
+            flags |= TypeFlags.HasFinalizerComputed;
+            flags |= TypeFlags.AttributeCacheComputed;
+
+            return flags;
+        }
+    }
+}

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

@@ -15,7 +15,7 @@ public MethodDesc GetAsyncVariantMethod(MethodDesc taskReturningMethod)
         {
             Debug.Assert(taskReturningMethod.Signature.ReturnsTaskOrValueTask());
             MethodDesc asyncMetadataMethodDef = taskReturningMethod.GetTypicalMethodDefinition();
-            MethodDesc result = _asyncVariantImplHashtable.GetOrCreateValue((EcmaMethod)asyncMetadataMethodDef);
+            MethodDesc result = _asyncVariantHashtable.GetOrCreateValue((EcmaMethod)asyncMetadataMethodDef);
 
             if (asyncMetadataMethodDef != taskReturningMethod)
             {
@@ -30,7 +30,7 @@ public MethodDesc GetAsyncVariantMethod(MethodDesc taskReturningMethod)
             return result;
         }
 
-        private sealed class AsyncVariantImplHashtable : LockFreeReaderHashtable<EcmaMethod, AsyncMethodVariant>
+        private sealed class AsyncVariantHashtable : LockFreeReaderHashtable<EcmaMethod, AsyncMethodVariant>
         {
             protected override int GetKeyHashCode(EcmaMethod key) => key.GetHashCode();
             protected override int GetValueHashCode(AsyncMethodVariant value) => value.Target.GetHashCode();
@@ -39,6 +39,32 @@ protected override bool CompareValueToValue(AsyncMethodVariant value1, AsyncMeth
                 => value1.Target == value2.Target;
             protected override AsyncMethodVariant CreateValueFromKey(EcmaMethod key) => new AsyncMethodVariant(key);
         }
-        private AsyncVariantImplHashtable _asyncVariantImplHashtable = new AsyncVariantImplHashtable();
+        private AsyncVariantHashtable _asyncVariantHashtable = new AsyncVariantHashtable();
+
+        public MetadataType GetContinuationType(GCPointerMap pointerMap)
+        {
+            return _continuationTypeHashtable.GetOrCreateValue(pointerMap);
+        }
+
+        private sealed class ContinuationTypeHashtable : LockFreeReaderHashtable<GCPointerMap, AsyncContinuationType>
+        {
+            private readonly CompilerTypeSystemContext _parent;
+            private MetadataType _continuationType;
+
+            public ContinuationTypeHashtable(CompilerTypeSystemContext parent)
+                => _parent = parent;
+
+            protected override int GetKeyHashCode(GCPointerMap key) => key.GetHashCode();
+            protected override int GetValueHashCode(AsyncContinuationType value) => value.PointerMap.GetHashCode();
+            protected override bool CompareKeyToValue(GCPointerMap key, AsyncContinuationType value) => key.Equals(value.PointerMap);
+            protected override bool CompareValueToValue(AsyncContinuationType value1, AsyncContinuationType value2)
+                => value1.PointerMap.Equals(value2.PointerMap);
+            protected override AsyncContinuationType CreateValueFromKey(GCPointerMap key)
+            {
+                _continuationType ??= _parent.SystemModule.GetKnownType("System.Runtime.CompilerServices"u8, "Continuation"u8);
+                return new AsyncContinuationType(_continuationType, key);
+            }
+        }
+        private ContinuationTypeHashtable _continuationTypeHashtable;
     }
 }

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

@@ -685,6 +685,7 @@ private void CompileMethodCleanup()
 
 #if !READYTORUN
             _debugInfo = null;
+            _asyncResumptionStub = null;
 #endif
 
             _debugLocInfos = null;
@@ -3401,7 +3402,19 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut)
         private CORINFO_CLASS_STRUCT_* getContinuationType(nuint dataSize, ref bool objRefs, nuint objRefsSize)
         {
             Debug.Assert(objRefsSize == (dataSize + (nuint)(PointerSize - 1)) / (nuint)PointerSize);
+#if READYTORUN
             throw new NotImplementedException("getContinuationType");
+#else
+            GCPointerMapBuilder gcMapBuilder = new GCPointerMapBuilder((int)dataSize, PointerSize);
+            ReadOnlySpan<bool> bools = MemoryMarshal.CreateReadOnlySpan(ref objRefs, (int)objRefsSize);
+            for (int i = 0; i < bools.Length; i++)
+            {
+                if (bools[i])
+                    gcMapBuilder.MarkGCPointer(i * PointerSize);
+            }
+
+            return ObjectToHandle(_compilation.TypeSystemContext.GetContinuationType(gcMapBuilder.ToGCMap()));
+#endif
         }
 
         private mdToken getMethodDefFromMethod(CORINFO_METHOD_STRUCT_* hMethod)
@@ -3756,7 +3769,14 @@ private bool getTailCallHelpers(ref CORINFO_RESOLVED_TOKEN callToken, CORINFO_SI
         private CORINFO_METHOD_STRUCT_* getAsyncResumptionStub(ref void* entryPoint)
 #pragma warning restore CA1822 // Mark members as static
         {
+#if READYTORUN
             throw new NotImplementedException("Crossgen2 does not support runtime-async yet");
+#else
+            _asyncResumptionStub ??= new AsyncResumptionStub(MethodBeingCompiled);
+
+            entryPoint = (void*)ObjectToHandle(_compilation.NodeFactory.MethodEntrypoint(_asyncResumptionStub));
+            return ObjectToHandle(_asyncResumptionStub);
+#endif
         }
 
         private byte[] _code;

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncResumptionStub.Mangling.cs

@@ -0,0 +1,15 @@
+// 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.Stubs;
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    public partial class AsyncResumptionStub : ILStubMethod, IPrefixMangledMethod
+    {
+        MethodDesc IPrefixMangledMethod.BaseMethod => _owningMethod;
+
+        string IPrefixMangledMethod.Prefix => "Resume";
+    }
+}

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncResumptionStub.Sorting.cs

@@ -0,0 +1,18 @@
+// 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.Stubs;
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    public partial class AsyncResumptionStub : ILStubMethod
+    {
+        protected override int ClassCode => 0x773ab1;
+
+        protected override int CompareToImpl(MethodDesc other, TypeSystemComparer comparer)
+        {
+            return comparer.Compare(_owningMethod, ((AsyncResumptionStub)other)._owningMethod);
+        }
+    }
+}

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

@@ -0,0 +1,53 @@
+// 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 Internal.IL;
+using Internal.IL.Stubs;
+using Internal.TypeSystem;
+
+using Debug = System.Diagnostics.Debug;
+
+namespace ILCompiler
+{
+    public partial class AsyncResumptionStub : ILStubMethod
+    {
+        private readonly MethodDesc _owningMethod;
+        private MethodSignature _signature;
+
+        public AsyncResumptionStub(MethodDesc owningMethod)
+        {
+            Debug.Assert(owningMethod.IsAsyncVariant()
+                || (owningMethod.IsAsync && !owningMethod.Signature.ReturnsTaskOrValueTask()));
+            _owningMethod = owningMethod;
+        }
+
+        public override ReadOnlySpan<byte> Name => _owningMethod.Name;
+        public override string DiagnosticName => _owningMethod.DiagnosticName;
+
+        public override TypeDesc OwningType => _owningMethod.OwningType;
+
+        public override MethodSignature Signature => _signature ??= InitializeSignature();
+
+        public override TypeSystemContext Context => _owningMethod.Context;
+
+        private MethodSignature InitializeSignature()
+        {
+            TypeDesc objectType = Context.GetWellKnownType(WellKnownType.Object);
+            TypeDesc byrefByte = Context.GetWellKnownType(WellKnownType.Byte).MakeByRefType();
+            return _signature = new MethodSignature(0, 0, objectType, [objectType, byrefByte]);
+        }
+
+        public override MethodIL EmitIL()
+        {
+            var emitter = new ILEmitter();
+            ILCodeStream codeStream = emitter.NewCodeStream();
+
+            // TODO: match getAsyncResumptionStub from CoreCLR VM
+            codeStream.EmitCallThrowHelper(emitter, Context.GetHelperEntryPoint("ThrowHelpers"u8, "ThrowNotSupportedException"u8));
+
+            return emitter.Link(this);
+        }
+    }
+}

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

@@ -59,6 +59,7 @@ public CompilerTypeSystemContext(TargetDetails details, SharedGenericsMode gener
             _typeWithRepeatedFieldsFieldLayoutAlgorithm = new TypeWithRepeatedFieldsFieldLayoutAlgorithm(_metadataFieldLayoutAlgorithm);
 
             _delegateInfoHashtable = new DelegateInfoHashtable(delegateFeatures);
+            _continuationTypeHashtable = new ContinuationTypeHashtable(this);
 
             _genericCycleDetector = new LazyGenericsSupport.GenericCycleDetector(genericCycleDepthCutoff, genericCycleBreadthCutoff);