Implement async functions and generators (.NET) (#2046)

* Implement async support in coroutine handling

* Add async generator support

* Create async coroutine cold

* Simplify and clean up

* Add more async tests

* Add async generator tests (.NET)

* Update after review

Author: BCSharp

eng/net10.0.props

@@ -23,6 +23,7 @@
     <Features>$(Features);FEATURE_METADATA_READER</Features>
     <Features>$(Features);FEATURE_MMAP</Features>
     <Features>$(Features);FEATURE_NATIVE</Features>
+    <Features>$(Features);FEATURE_NET_ASYNC</Features>
     <Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
     <Features>$(Features);FEATURE_PIPES</Features>
     <Features>$(Features);FEATURE_PROCESS</Features>

eng/net8.0.props

@@ -23,6 +23,7 @@
     <Features>$(Features);FEATURE_METADATA_READER</Features>
     <Features>$(Features);FEATURE_MMAP</Features>
     <Features>$(Features);FEATURE_NATIVE</Features>
+    <Features>$(Features);FEATURE_NET_ASYNC</Features>
     <Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
     <Features>$(Features);FEATURE_PIPES</Features>
     <Features>$(Features);FEATURE_PROCESS</Features>

eng/net9.0.props

@@ -23,6 +23,7 @@
     <Features>$(Features);FEATURE_METADATA_READER</Features>
     <Features>$(Features);FEATURE_MMAP</Features>
     <Features>$(Features);FEATURE_NATIVE</Features>
+    <Features>$(Features);FEATURE_NET_ASYNC</Features>
     <Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
     <Features>$(Features);FEATURE_PIPES</Features>
     <Features>$(Features);FEATURE_PROCESS</Features>

src/core/IronPython/Compiler/Ast/AstMethods.cs

@@ -80,6 +80,11 @@ internal static class AstMethods {
         public static readonly MethodInfo FormatString = GetMethod((Func<CodeContext, string, object, string>)PythonOps.FormatString);
         public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func<object, object>)PythonOps.GeneratorCheckThrowableAndReturnSendValue);
         public static readonly MethodInfo MakeCoroutine = GetMethod((Func<PythonFunction, MutableTuple, object, PythonCoroutine>)PythonOps.MakeCoroutine);
+#if FEATURE_NET_ASYNC
+        public static readonly MethodInfo AsTaskForAwait = GetMethod((Func<object, System.Threading.Tasks.Task<object>>)PythonOps.AsTaskForAwait);
+        public static readonly MethodInfo MakeAsyncCoroutine = GetMethod((Func<PythonFunction, Func<System.Threading.Tasks.Task<object>>, System.Threading.CancellationTokenSource, System.Runtime.CompilerServices.StrongBox<System.Exception>, PythonCoroutine>)PythonOps.MakeAsyncCoroutine);
+        public static readonly MethodInfo MakeAsyncGenerator = GetMethod((Func<PythonFunction, System.Collections.Generic.IAsyncEnumerable<object>, System.Runtime.CompilerServices.StrongBox<object>, System.Runtime.CompilerServices.StrongBox<System.Exception>, System.Threading.CancellationTokenSource, PythonAsyncGenerator>)PythonOps.MakeAsyncGenerator);
+#endif
 
         // builtins
         public static readonly MethodInfo Format = GetMethod((Func<CodeContext, object, string, string>)PythonOps.Format);

src/core/IronPython/Compiler/Ast/AwaitExpression.cs

@@ -6,15 +6,41 @@
 
 using MSAst = System.Linq.Expressions;
 
+using IronPython.Runtime.Operations;
+
 using AstUtils = Microsoft.Scripting.Ast.Utils;
 
 namespace IronPython.Compiler.Ast {
     using Ast = MSAst.Expression;
 
     /// <summary>
-    /// Represents an await expression. Implemented as yield from expr.__await__().
+    /// Represents <c>await expr</c>. Under <c>FEATURE_NET_ASYNC</c> this compiles directly to a DLR async suspension point.
+    /// Otherwise it is desugared into <c>yield from expr.__await__()</c> against the enclosing generator-shaped coroutine state machine.
     /// </summary>
     public class AwaitExpression : Expression {
+#if FEATURE_NET_ASYNC
+        public AwaitExpression(Expression expression) {
+            Expression = expression;
+        }
+
+        public Expression Expression { get; }
+
+        public override MSAst.Expression Reduce() {
+            // await x -> AsyncHelpers-driven suspension on the Task produced by
+            // PythonOps.AsTaskForAwait(x).
+            return AstUtils.Await(
+                Ast.Call(
+                    AstMethods.AsTaskForAwait,
+                    AstUtils.Convert(Expression, typeof(object))));
+        }
+
+        public override void Walk(PythonWalker walker) {
+            if (walker.Walk(this)) {
+                Expression?.Walk(walker);
+            }
+            walker.PostWalk(this);
+        }
+#else
         private readonly Statement _statement;
         private readonly NameExpression _result;
 
@@ -60,6 +86,7 @@ public override void Walk(PythonWalker walker) {
             }
             walker.PostWalk(this);
         }
+#endif
 
         public override string NodeName => "await expression";
     }

src/core/IronPython/Compiler/Ast/FunctionDefinition.cs

@@ -8,6 +8,7 @@
 using System.Runtime.CompilerServices;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 
 using IronPython.Runtime;
 using IronPython.Runtime.Operations;
@@ -117,7 +118,25 @@ internal override int KwOnlyArgCount {
 
         public Expression ReturnAnnotation { get; internal set; }
 
+#if FEATURE_NET_ASYNC
+        // Under .NET-async, async functions are compiled directly to a Task<object?> via the DLR's AsyncExpression
+        // rather than reused through the generator state machine, so IsAsync does not imply generator-shaped emission.
+        internal override bool IsGeneratorMethod => IsGenerator;
+
+        // Async-generator (PEP 525) channels. The StrongBox *values* are per-async-generator instance,
+        // allocated per stack frame at runtime in the function body. 
+        // Declared/assigned in the body, captured by the generator (its yields read them through Parent),
+        // and handed to the PythonAsyncGenerator wrapper, which writes them before each resume:
+        //   AsyncSendSlot  — the value of `x = yield z` (asend(v); None for __anext__/async for).
+        //   AsyncThrowSlot — an exception to rethrow at the yield resume point (athrow/aclose).
+        private readonly MSAst.ParameterExpression _asyncSendSlot = MSAst.Expression.Variable(typeof(StrongBox<object>), "$asyncSend");
+        private readonly MSAst.ParameterExpression _asyncThrowSlot = MSAst.Expression.Variable(typeof(StrongBox<Exception>), "$asyncThrow");
+
+        internal MSAst.ParameterExpression AsyncSendSlot => _asyncSendSlot;
+        internal MSAst.ParameterExpression AsyncThrowSlot => _asyncThrowSlot;
+#else
         internal override bool IsGeneratorMethod => IsGenerator || IsAsync;
+#endif
 
         /// <summary>
         /// The function is a generator
@@ -182,9 +201,15 @@ internal override FunctionAttributes Flags {
                     fa |= FunctionAttributes.ContainsTryFinally;
                 }
 
+#if FEATURE_NET_ASYNC
+                if (IsGenerator) {
+                    fa |= FunctionAttributes.Generator;
+                }
+#else
                 if (IsGenerator || IsAsync) {
                     fa |= FunctionAttributes.Generator;
                 }
+#endif
 
                 if (IsAsync) {
                     fa |= FunctionAttributes.Coroutine;
@@ -357,7 +382,13 @@ internal MSAst.Expression MakeFunctionExpression() {
                                 annotations
                             )
                         ),
+#if FEATURE_NET_ASYNC
+                    // Async generators are lowered via AsyncEnumerableExpression in the body,
+                    // so they must not be wrapped as a PythonGenerator here — only plain (non-async) generators are.
+                    (IsGenerator && !IsAsync) ?
+#else
                     (IsGenerator || IsAsync) ?
+#endif
                         (MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold, IsAsync) :
                         (MSAst.Expression)code
                 );
@@ -659,7 +690,13 @@ private LightLambdaExpression CreateFunctionLambda() {
             // For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close().
             // The exception traceback needs to come from the generator's method body, and so we must do the check and throw
             // from inside the generator.
+#if FEATURE_NET_ASYNC
+            // Async generators have no backing PythonGenerator (they lower to IAsyncEnumerable via AsyncEnumerableExpression),
+            // so skip the $generator.CheckThrowable() prologue for them.
+            if (IsGenerator && !IsAsync) {
+#else
             if (IsGenerator || IsAsync) {
+#endif
                 MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None);
                 statements.Add(s1);
             }
@@ -681,6 +718,59 @@ private LightLambdaExpression CreateFunctionLambda() {
             body = Ast.Block(body, AstUtils.Empty());
             body = AddReturnTarget(body);
 
+#if FEATURE_NET_ASYNC
+            // Under .NET-async, an `async def` body returns a PythonCoroutine wrapping a Task<object?>.
+            // We pre-allocate a CancellationTokenSource and a StrongBox<Exception?> here
+            // so the same instances are shared with both AsyncExpression, which threads them into AsyncHelpers.DriveAsync
+            // and PythonCoroutine, which uses them to implement coro.throw(exc) on a running coroutine:
+            // write the exception to the box, cancel the CTS, and DriveAsync surfaces that exception in place of OperationCanceledException.
+            if (IsAsync) {
+                var cts = MSAst.Expression.Variable(typeof(CancellationTokenSource), "$cts");
+                var excBox = MSAst.Expression.Variable(typeof(StrongBox<Exception>), "$cancelExc");
+                var ctToken = MSAst.Expression.Property(cts, nameof(CancellationTokenSource.Token));
+                if (IsGenerator) {
+                    // Async generator: the body has both `await` and `yield`. Lower it to an
+                    // IAsyncEnumerable<object?> via AsyncEnumerableExpression, sharing the generator label so
+                    // the body's yields and the rewritten awaits land in one generator, then wrap it in a
+                    // PythonAsyncGenerator. The send/throw slots are per-generator StrongBoxes captured by the
+                    // generator (the body's yields read them) AND handed to the wrapper, which writes them
+                    // before each resume — see AsyncSendSlot / AsyncThrowSlot.
+                    var sendSlot = AsyncSendSlot;
+                    var throwSlot = AsyncThrowSlot;
+                    body = MSAst.Expression.Block(
+                        [cts, excBox, sendSlot, throwSlot],
+                        MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(CancellationTokenSource))),
+                        MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(StrongBox<Exception>))),
+                        MSAst.Expression.Assign(sendSlot, MSAst.Expression.New(typeof(StrongBox<object>))),
+                        MSAst.Expression.Assign(throwSlot, MSAst.Expression.New(typeof(StrongBox<Exception>))),
+                        Ast.Call(
+                            AstMethods.MakeAsyncGenerator,
+                            _functionParam,
+                            AstUtils.AsyncEnumerable(Name, body, GeneratorLabel, ctToken, excBox),
+                            sendSlot,
+                            throwSlot,
+                            cts));
+                } else {
+                    // Plain async def: the body returns a PythonCoroutine wrapping a Task<object?>.
+                    // Lazy start: hand MakeAsyncCoroutine a thunk (Func<Task<object?>>) instead of an already-running Task,
+                    // so the body doesn't execute until the coroutine is first driven (send/AsTask).
+                    // This makes calling an async def side-effect-free (PEP 492) and lets the body's first await capture the driver's SynchronizationContext
+                    // rather than whatever context happened to be current at construction.
+                    body = MSAst.Expression.Block(
+                        [cts, excBox],
+                        MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(CancellationTokenSource))),
+                        MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(StrongBox<Exception>))),
+                        Ast.Call(
+                            AstMethods.MakeAsyncCoroutine,
+                            _functionParam,
+                            MSAst.Expression.Lambda<Func<Task<object>>>(
+                                AstUtils.Async(Name, body, ctToken, excBox)),
+                            cts,
+                            excBox));
+                }
+            }
+#endif
+
             MSAst.Expression bodyStmt = body;
             if (localContext != null) {
                 var createLocal = CreateLocalContext(_parentContext);

src/core/IronPython/Compiler/Ast/ReturnStatement.cs

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the Apache 2.0 License.
 // See the LICENSE file in the project root for more information.
 
+using System.Diagnostics;
+
 using MSAst = System.Linq.Expressions;
 
 namespace IronPython.Compiler.Ast {
@@ -17,6 +19,16 @@ public ReturnStatement(Expression expression) {
 
         public override MSAst.Expression Reduce() {
             if (Parent.IsGeneratorMethod) {
+#if FEATURE_NET_ASYNC
+                // An async generator (`async def` with `yield`) lowers through the DLR generator via AsyncEnumerableExpression,
+                // which doesn't understand IronPython's -2 "generator-return" marker.
+                // `return` there is always bare (`return value` is a SyntaxError in async generators)
+                // and simply ends the async iteration — map it to a YieldBreak without a value.
+                if (Parent is FunctionDefinition { IsAsync: true }) {
+                    Debug.Assert(Expression == null, "async generators should not have a return value");
+                    return GlobalParent.AddDebugInfo(AstUtils.YieldBreak(GeneratorLabel), Span);
+                }
+#endif
                 // Reduce to a yield return with a marker of -2, this will be interpreted as a yield break with a return value
                 return GlobalParent.AddDebugInfo(AstUtils.YieldReturn(GeneratorLabel, TransformOrConstantNull(Expression, typeof(object)), -2), Span);
             }

src/core/IronPython/Compiler/Ast/YieldExpression.cs

@@ -7,6 +7,8 @@
 using MSAst = System.Linq.Expressions;
 
 using System;
+using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
 using System.Diagnostics;
 using Microsoft.Scripting;
 using Microsoft.Scripting.Runtime;
@@ -19,6 +21,13 @@ namespace IronPython.Compiler.Ast {
     //    x = yield z
     // The return value (x) is provided by calling Generator.Send()
     public class YieldExpression : Expression {
+#if FEATURE_NET_ASYNC
+        private static readonly System.Reflection.MethodInfo s_captureMethod
+            = typeof(ExceptionDispatchInfo).GetMethod(nameof(ExceptionDispatchInfo.Capture))!;
+        private static readonly System.Reflection.MethodInfo s_throwMethod
+            = typeof(ExceptionDispatchInfo).GetMethod(nameof(ExceptionDispatchInfo.Throw), Type.EmptyTypes)!;
+#endif
+
         public YieldExpression(Expression? expression) {
             Expression = expression;
         }
@@ -43,16 +52,40 @@ internal static MSAst.Expression CreateCheckThrowExpression(SourceSpan span) {
         }
 
         public override MSAst.Expression Reduce() {
+            MSAst.Expression yieldValue = Expression == null ? AstUtils.Constant(null) : AstUtils.Convert(Expression, typeof(object));
+
+#if FEATURE_NET_ASYNC
+            // An async generator (`async def` with `yield`) is lowered via AsyncEnumerableExpression
+            // and has no backing PythonGenerator, so there is no `$generator` to call CheckThrowable() on.
+            // Instead the resume reads two per-generator cells that PythonAsyncGenerator writes before advancing:
+            //   AsyncThrowSlot — if set (athrow/aclose), rethrow it here (preserving stack);
+            //                    cleared first so a body that catches it and yields again doesn't re-throw on the next resume.
+            //   AsyncSendSlot  — the value of the yield expression: the asend(v) value, or None.
+            if (Parent is FunctionDefinition { IsAsync: true } fd) {
+                MSAst.ParameterExpression sendSlot = fd.AsyncSendSlot;
+                MSAst.ParameterExpression throwSlot = fd.AsyncThrowSlot;
+                MSAst.ParameterExpression pending = Ast.Variable(typeof(Exception), "$athrow");
+                return Ast.Block(
+                    typeof(object),
+                    [pending],
+                    AstUtils.YieldReturn(GeneratorLabel, yieldValue),
+                    Ast.Assign(pending, Ast.Field(throwSlot, nameof(StrongBox<Exception>.Value))),
+                    Ast.Assign(Ast.Field(throwSlot, nameof(StrongBox<Exception>.Value)), Ast.Constant(null, typeof(Exception))),
+                    Ast.IfThen(
+                        Ast.ReferenceNotEqual(pending, Ast.Constant(null, typeof(Exception))),
+                        Ast.Call(Ast.Call(s_captureMethod, pending), s_throwMethod)),
+                    Ast.Field(sendSlot, nameof(StrongBox<object>.Value))
+                );
+            }
+#endif
+
             // (yield z) becomes:
             // .comma (1) {
             //    .void ( .yield_statement (_expression) ),
-            //    $gen.CheckThrowable() // <-- has return result from send            
+            //    $gen.CheckThrowable() // <-- has return result from send
             //  }
             return Ast.Block(
-                AstUtils.YieldReturn(
-                    GeneratorLabel,
-                    Expression == null ? AstUtils.Constant(null) : AstUtils.Convert(Expression, typeof(object))
-                ),
+                AstUtils.YieldReturn(GeneratorLabel, yieldValue),
                 CreateCheckThrowExpression(Span) // emits ($gen.CheckThrowable())
             );
         }

src/core/IronPython/Compiler/Parser.cs

@@ -1999,10 +1999,16 @@ private Expression ParseAtomExpr() {
                 if (current is null || !current.IsAsync) {
                     ReportSyntaxError("'await' outside async function");
                 }
+#if !FEATURE_NET_ASYNC
+                // Under the generator-based async path, `await` desugars to `yield from`,
+                // so the enclosing function must be marked a generator.
+                // Under FEATURE_NET_ASYNC, the body is lowered through AsyncExpression and is *not* a Python generator
+                // so this mark assignment is being skipped here.
                 if (current is not null) {
                     current.IsGenerator = true;
                     current.GeneratorStop = GeneratorStop;
                 }
+#endif
             }
 
             Expression ret = ParseAtom();