Trace the bodies of async functions and async generators

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

@@ -681,11 +681,9 @@ private LightLambdaExpression CreateFunctionLambda() {
 
             List<MSAst.Expression> statements = new List<MSAst.Expression>();
 
-            // Add beginning sequence point. For async functions this body is deferred into the nested Async/AsyncEnumerable lambda,
-            // and the outer factory (which runs first, and is the only one traced today) re-emits its own entry point at StartIndex;
-            // anchor this one at HeaderIndex (end of the signature) so the two occupy distinct source locations should the deferred
-            // body ever become traceable. For single-line signatures both still land on the same line (differing only in column).
-            statements.Add(MakeFunctionEntrySequencePoint(IsAsync ? HeaderIndex : StartIndex));
+            // Add beginning sequence point: the location of the 'call' trace event. Async functions included, since this is their body too;
+            // their outer factory, which only creates the coroutine, is not traced (as in CPython).
+            statements.Add(MakeFunctionEntrySequencePoint(StartIndex));
 
 
             // For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close().
@@ -740,7 +738,6 @@ private LightLambdaExpression CreateFunctionLambda() {
                     var throwSlot = AsyncThrowSlot;
                     body = MSAst.Expression.Block(
                         [cts, excBox, sendSlot, throwSlot],
-                        MakeFunctionEntrySequencePoint(StartIndex),
                         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>))),
@@ -760,7 +757,6 @@ private LightLambdaExpression CreateFunctionLambda() {
                     // rather than whatever context happened to be current at construction.
                     body = MSAst.Expression.Block(
                         [cts, excBox],
-                        MakeFunctionEntrySequencePoint(StartIndex),
                         MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(CancellationTokenSource))),
                         MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(StrongBox<Exception>))),
                         Ast.Call(

src/core/IronPython/Compiler/GeneratorRewriter.cs

@@ -373,6 +373,11 @@ protected override Expression VisitTry(TryExpression node) {
             if (catchYields != tryYields) {
                 var block = new List<Expression>();
 
+                // A deferred handler runs if its variable is set, so the variables must be reset on each entry into the try: they are
+                // lifted, and would otherwise keep the exception of an earlier pass (e.g. of a loop) whose handler did not reach its end.
+                // A resumption into a deferred handler is routed past the reset, which keeps the exception for the rest of the handler.
+                var resetDeferredVars = new List<Expression>();
+
                 block.Add(MakeYieldRouter(node.Body.Type, tryYields, catchYields, tryStart));
                 block.Add(null); // empty slot to fill in later
 
@@ -395,7 +400,8 @@ protected override Expression VisitTry(TryExpression node) {
                     // the catch block had one. It needs to be hoisted because
                     // the catch might contain yields.
                     var deferredVar = c.Variable ?? Expression.Variable(c.Test, null);
-                    LiftVariable(deferredVar);                    
+                    LiftVariable(deferredVar);
+                    resetDeferredVars.Add(DelayedAssign(Visit(deferredVar), AstUtils.Constant(null, deferredVar.Type)));
 
                     // We need to ensure that filters can access the exception
                     // variable
@@ -431,7 +437,8 @@ protected override Expression VisitTry(TryExpression node) {
                     );
                 }
 
-                block[1] = Expression.MakeTry(null, @try, null, null, new ReadOnlyCollection<CatchBlock>(handlers));
+                resetDeferredVars.Add(Expression.MakeTry(null, @try, null, null, new ReadOnlyCollection<CatchBlock>(handlers)));
+                block[1] = Expression.Block(resetDeferredVars);
                 @try = Expression.Block(block);
                 Debug.Assert(@try.Type == node.Body.Type);
                 handlers = Array.Empty<CatchBlock>(); // so we don't reuse these

src/core/IronPython/Compiler/Parser.cs

@@ -1126,7 +1126,7 @@ private Statement ParseDecorated() {
                 res = cls;
             } else if (PeekToken() == Tokens.KeywordAsyncToken) {
                 Eat(TokenKind.KeywordAsync);
-                FunctionDefinition fnc = ParseFuncDef(true);
+                FunctionDefinition fnc = ParseFuncDef(asyncStart: GetStart());
                 fnc.Decorators = decorators.ToArray();
                 res = fnc;
             } else {
@@ -1139,10 +1139,11 @@ private Statement ParseDecorated() {
 
         // funcdef: [decorators] 'def' NAME parameters ':' suite
         // parameters: '(' [typedargslist] ')'
-        // this gets called with "def" as the look-ahead
-        private FunctionDefinition ParseFuncDef(bool isAsync = false) {
+        // this gets called with "def" as the look-ahead; asyncStart is the start of the preceding 'async' keyword, if any
+        private FunctionDefinition ParseFuncDef(int? asyncStart = null) {
+            bool isAsync = asyncStart.HasValue;
             Eat(TokenKind.KeywordDef);
-            var start = GetStart();
+            var start = asyncStart ?? GetStart();   // like in CPython, an async function starts at its 'async' keyword
             string name = ReadName();
 
             Eat(TokenKind.LeftParenthesis);
@@ -1469,7 +1470,7 @@ private Statement ParseAsyncStmt(bool onlyAllowDef = false) {
 
             switch (PeekToken().Kind) {
                 case TokenKind.KeywordDef:
-                    return ParseFuncDef(true);
+                    return ParseFuncDef(asyncStart: start);
                 case TokenKind.KeywordWith:
                     if (onlyAllowDef) goto default;
                     return ParseAsyncWithStmt(start);

src/core/IronPython/Runtime/Exceptions/TraceBack.cs

@@ -222,6 +222,12 @@ private void SetLineNumber(int newLineNum) {
             }
 
             FunctionCode funcCode = _debugProperties.Code;
+#if FEATURE_NET_ASYNC
+            // The frame of a traced async body reports its location but cannot change it.
+            if ((funcCode.Flags & FunctionAttributes.Coroutine) != 0) {
+                throw PythonOps.ValueError("can't jump in an async function or async generator");
+            }
+#endif
             Dictionary<int, Dictionary<int, bool>> loopAndFinallyLocations = _debugProperties.LoopAndFinallyLocations;
             Dictionary<int, bool> handlerLocations = _debugProperties.HandlerLocations;
 

src/core/IronPython/Runtime/FunctionCode.cs

@@ -754,11 +754,14 @@ private LambdaExpression GetGeneratorOrNormalLambdaTracing(PythonContext context
             // Under .NET-async, async bodies are lowered before they reach here and must NOT go through the generator rewriter:
             //  * plain async (Coroutine, no Generator) via AsyncExpression,
             //  * async generators (Generator+Coroutine) via AsyncEnumerableExpression.
-            //  Only a plain generator (Generator without Coroutine) is rewritten here.
-            if ((Flags & FunctionAttributes.Generator) == 0 || (Flags & FunctionAttributes.Coroutine) != 0) {
-#else
-            if ((Flags & (FunctionAttributes.Generator | FunctionAttributes.Coroutine)) == 0) {
+            // As in CPython, calling an async function only creates the coroutine (or async generator) object, so the outer factory lambda
+            // is not traced at all. The async body is traced instead.
+            if ((Flags & FunctionAttributes.Coroutine) != 0) {
+                var lambda = (LambdaExpression)Compiler.Ast.Node.RemoveFrame(_lambda.GetLambda());
+                return (LambdaExpression)new AsyncDebugHookInjector(context.DebugContext, debugInfo).Visit(lambda);
+            }
 #endif
+            if ((Flags & (FunctionAttributes.Generator | FunctionAttributes.Coroutine)) == 0) {
                 return context.DebugContext.TransformLambda((LambdaExpression)Compiler.Ast.Node.RemoveFrame(_lambda.GetLambda()), debugInfo);
             }
 
@@ -790,7 +793,7 @@ private LambdaExpression GetGeneratorOrNormalLambdaTracing(PythonContext context
         private LightLambdaExpression GetGeneratorOrNormalLambda() {
             LightLambdaExpression finalCode;
 #if FEATURE_NET_ASYNC
-            // See GetGeneratorOrNormalLambda's sibling in MakeDebuggableFunctionCode: only a plain generator
+            // See the sibling GetGeneratorOrNormalLambdaTracing: only a plain generator
             // (Generator without Coroutine) is rewritten; plain async and async generators are lowered earlier.
             if ((Flags & FunctionAttributes.Generator) == 0 || (Flags & FunctionAttributes.Coroutine) != 0) {
 #else
@@ -1204,4 +1207,35 @@ public Dictionary<int, bool> LoopOrFinallyIds {
             }
         }
     }
+
+#if FEATURE_NET_ASYNC
+    /// <summary>
+    ///   Swaps a lowered <see cref="AsyncExpression"/> / <see cref="AsyncEnumerableExpression"/> for a debug-enabled copy.
+    /// </summary>
+    /// <remarks>
+    ///   Used to build the tracing variant of an async function: the debug-enabled node instruments its own body when reduced, so the
+    ///   async body becomes traceable while the outer factory stays untraced. A function's own async node is the only one in its lambda
+    ///   — nested defs compile to their own lambdas/codes — so a single replacement is made, and the node is not descended into.
+    /// </remarks>
+    internal sealed class AsyncDebugHookInjector : ExpressionVisitor {
+        private readonly DebugContext _debugContext;
+        private readonly DebugLambdaInfo _debugLambdaInfo;
+
+        internal AsyncDebugHookInjector(DebugContext debugContext, DebugLambdaInfo debugLambdaInfo) {
+            _debugContext = debugContext;
+            _debugLambdaInfo = debugLambdaInfo;
+        }
+
+        protected override Expression VisitExtension(Expression node) {
+            switch (node) {
+                case AsyncExpression ae:
+                    return ae.WithDebugInfo(_debugContext, _debugLambdaInfo);
+                case AsyncEnumerableExpression aee:
+                    return aee.WithDebugInfo(_debugContext, _debugLambdaInfo);
+                default:
+                    return base.VisitExtension(node);
+            }
+        }
+    }
+#endif
 }

tests/IronPython.Tests/Cases/IronPythonCasesManifest.ini

@@ -5,7 +5,6 @@ Redirect=false
 Timeout=120000 # 2 minute timeout
 
 [IronPython.test_asyncgen]
-Ignore=true
 RunCondition='$(FRAMEWORK)' <> '.NETFramework,Version=v4.6.2'
 Reason=Requires IAsyncEnumerable & IAsyncDisposable
 

tests/suite/test_async.py

@@ -48,6 +48,12 @@ async def foo():
             return 42
         self.assertEqual(run_coro(foo()), 42)
 
+    def test_location_starts_at_async_keyword(self):
+        # As in CPython. It is also where a coroutine's 'call' trace event is reported.
+        import ast
+        node = ast.parse("x = 1\nasync def f():\n    pass\n").body[1]
+        self.assertEqual((node.lineno, node.col_offset), (2, 0))
+
     def test_coroutine_type(self):
         async def foo():
             return 1
@@ -171,6 +177,47 @@ async def test():
 
         self.assertEqual(run_coro(test()), 'caught: boom')
 
+    def test_await_in_handler_left_by_continue(self):
+        """A handler that awaits and then leaves its iteration does not run again in a later iteration that does not raise."""
+        async def noop():
+            pass
+
+        async def test():
+            log = []
+            for i in range(3):
+                try:
+                    if i == 0:
+                        raise ValueError()
+                    log.append(('try', i))
+                except ValueError:
+                    await noop()
+                    log.append(('except', i))
+                    continue
+            return log
+
+        self.assertEqual(run_coro(test()), [('except', 0), ('try', 1), ('try', 2)])
+
+    def test_await_in_handler_left_by_break(self):
+        """As above, with the handler breaking out of an inner loop that the outer loop enters again."""
+        async def noop():
+            pass
+
+        async def test():
+            log = []
+            for j in range(2):
+                for i in range(2):
+                    try:
+                        if j == 0:
+                            raise ValueError()
+                        log.append(('try', j, i))
+                    except ValueError:
+                        await noop()
+                        log.append(('except', j, i))
+                        break
+            return log
+
+        self.assertEqual(run_coro(test()), [('except', 0, 0), ('try', 1, 0), ('try', 1, 1)])
+
 
 class AsyncWithTest(unittest.TestCase):
     """Tests for async with statement."""