Trace async bodies with a debug hook

src/core/Microsoft.Dynamic/Ast/AsyncDebugSupport.cs

@@ -0,0 +1,213 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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.
+
+#nullable enable
+
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Linq.Expressions;
+using System.Reflection;
+using System.Runtime.CompilerServices;
+
+using Microsoft.Scripting.Debugging;
+using Microsoft.Scripting.Debugging.CompilerServices;
+
+namespace Microsoft.Scripting.Ast {
+    /// <summary>
+    ///   Compiles an async body for tracing.
+    /// </summary>
+    /// <remarks>
+    ///   Used by <see cref="AsyncRewriter"/> and <c>AsyncEnumerableRewriter</c> when the async node carries a debug hook (set by the host
+    ///   only when it builds the tracing variant of the code, e.g. when <c>sys.settrace</c> is active). In two steps:
+    ///   <list type="number">
+    ///     <item><see cref="InjectDebugPoints"/> instruments the body before it is sliced into a state machine, so that the debug points
+    ///     end up in the same state machine as the awaits and yields.</item>
+    ///     <item><see cref="TraceStateMachine"/> wraps the resulting state machine in a <see cref="TracingStateMachine"/>, which dispatches
+    ///     the trace events while the driver pumps it.</item>
+    ///   </list>
+    ///
+    ///   <para>The body's locals are not lifted, so its frame reports its location but no locals; the scope map is therefore uniform
+    ///   (every sequence point maps to no locals).</para>
+    /// </remarks>
+    internal sealed class AsyncDebugSupport {
+        private static readonly ConstructorInfo s_tracingStateMachineCtor
+            = typeof(TracingStateMachine).GetConstructor(
+                BindingFlags.Instance | BindingFlags.NonPublic, binder: null,
+                [typeof(IEnumerator<object?>), typeof(DebugContext), typeof(FunctionInfo), typeof(StrongBox<object?>), typeof(StrongBox<Exception?>), typeof(bool)],
+                modifiers: null)!;
+        private static readonly MethodInfo s_onExceptionCaughtMethod
+            = typeof(TracingStateMachine).GetMethod(nameof(TracingStateMachine.OnExceptionCaught), BindingFlags.Instance | BindingFlags.NonPublic)!;
+
+        private readonly DebugContext _debugContext;
+        private readonly DebugLambdaInfo _lambdaInfo;
+        private readonly string? _name;
+        private readonly ParameterExpression _exceptionSlot;
+
+        // Holds the TracingStateMachine, so that the body's exception handlers can report to it.
+        private readonly ParameterExpression _tracer = Expression.Variable(typeof(TracingStateMachine), "$tracer");
+
+        // Null until the body is instrumented, and also if it has no sequence points (then there is nothing to trace).
+        private FunctionInfo? _functionInfo;
+
+        /// <param name="exceptionSlot">The body's exception slot, which the body rethrows from at each resume point.</param>
+        internal AsyncDebugSupport(DebugContext debugContext, DebugLambdaInfo lambdaInfo, string? name, ParameterExpression exceptionSlot) {
+            _debugContext = debugContext;
+            _lambdaInfo = lambdaInfo;
+            _name = name;
+            _exceptionSlot = exceptionSlot;
+        }
+
+
+        /// <summary>
+        ///   Instruments <paramref name="body"/> to report its sequence points and the exceptions it handles.
+        /// </summary>
+        /// <remarks>
+        ///   Each <see cref="DebugInfoExpression"/> becomes <c>{ yield DebugPoint(cookie); rethrow-if-pending }</c> on
+        ///   <paramref name="yieldLabel"/>, the label the body's awaits and yields are sliced on, and each exception handler first reports
+        ///   the exception it handles. The rethrow lets an exception raised by the trace callback surface inside the body.
+        /// </remarks>
+        internal Expression InjectDebugPoints(Expression body, LabelTarget yieldLabel) {
+            var injector = new DebugPointInjector(_debugContext, yieldLabel, _exceptionSlot, _tracer);
+            Expression instrumentedBody = injector.Visit(body);
+            if (injector.SequencePoints.Length == 0) {
+                return body;
+            }
+
+            _functionInfo = DebugContext.CreateFunctionInfo(
+                generatorFactory: null,     // the frame cannot be remapped to a generator, so it cannot change its location either
+                _name ?? "<async>",
+                injector.SequencePoints,
+                injector.VariableScopeMap,
+                Array.Empty<VariableInfo>(),
+                _lambdaInfo.CustomPayload);
+            return instrumentedBody;
+        }
+
+
+        /// <summary>
+        ///   Wraps <paramref name="stateMachine"/>, the state machine of the instrumented body, in a <see cref="TracingStateMachine"/>.
+        /// </summary>
+        /// <param name="resultSlot">Where the body leaves its return value, or null if it has none.</param>
+        /// <param name="isAsyncEnumerable">Whether the state machine is driven by <c>AsyncHelpers.DriveAsyncEnumerable</c>.</param>
+        internal Expression TraceStateMachine(Expression stateMachine, ParameterExpression? resultSlot, bool isAsyncEnumerable) {
+            if (_functionInfo is null) {
+                return stateMachine;
+            }
+
+            return Expression.Block(
+                typeof(TracingStateMachine),
+                [_tracer],
+                Expression.Assign(
+                    _tracer,
+                    Expression.New(
+                        s_tracingStateMachineCtor,
+                        stateMachine,
+                        Expression.Constant(_debugContext),
+                        Expression.Constant(_functionInfo),
+                        (Expression?)resultSlot ?? Expression.Constant(null, typeof(StrongBox<object?>)),
+                        _exceptionSlot,
+                        Expression.Constant(isAsyncEnumerable))));
+        }
+
+
+        /// <summary>
+        ///   Rewrites the debug info and the exception handlers of an async body for <see cref="TracingStateMachine"/>.
+        /// </summary>
+        /// <remarks>
+        ///   Each <see cref="DebugInfoExpression"/> becomes a yield of a <see cref="DebugPoint"/> on the body's own yield label, recording
+        ///   the source span per cookie, and each catch block first reports the exception to the tracer. This mirrors what
+        ///   <c>DebugInfoRewriter</c> does to a synchronous body.
+        /// </remarks>
+        private sealed class DebugPointInjector : ExpressionVisitor {
+            private readonly DebugContext _debugContext;
+            private readonly LabelTarget _yieldLabel;
+            private readonly ParameterExpression _exceptionSlot;
+            private readonly ParameterExpression _tracer;
+            private readonly List<DebugSourceSpan> _spans = new();
+
+            internal DebugPointInjector(DebugContext debugContext, LabelTarget yieldLabel, ParameterExpression exceptionSlot, ParameterExpression tracer) {
+                _debugContext = debugContext;
+                _yieldLabel = yieldLabel;
+                _exceptionSlot = exceptionSlot;
+                _tracer = tracer;
+            }
+
+            internal DebugSourceSpan[] SequencePoints => _spans.ToArray();
+
+            internal IList<VariableInfo>[] VariableScopeMap {
+                get {
+                    var map = new IList<VariableInfo>[_spans.Count];
+                    for (int i = 0; i < map.Length; i++) {
+                        map[i] = Array.Empty<VariableInfo>();
+                    }
+                    return map;
+                }
+            }
+
+            protected override Expression VisitDebugInfo(DebugInfoExpression node) {
+                if (node.IsClear) {
+                    return Expression.Empty();
+                }
+
+                if (node.Document is null) {
+                    throw new InvalidOperationException(
+                        string.Format(CultureInfo.CurrentCulture, "DebugInfoExpression at sequence point {0} has no SymbolDocumentInfo", _spans.Count));
+                }
+
+                DebugSourceFile sourceFile = _debugContext.GetDebugSourceFile(
+                    string.IsNullOrEmpty(node.Document.FileName) ? "<compile>" : node.Document.FileName);
+
+                int cookie = _spans.Count;
+                _spans.Add(new DebugSourceSpan(sourceFile, node.StartLine, node.StartColumn, node.EndLine, node.EndColumn));
+
+                return Expression.Block(
+                    typeof(void),
+                    Utils.YieldReturn(_yieldLabel, Expression.Constant(new DebugPoint(cookie))),
+                    AsyncRewriter.RethrowIfPending(_exceptionSlot));
+            }
+
+            protected override Expression VisitTry(TryExpression node) {
+                var visited = (TryExpression)base.VisitTry(node);
+                if (visited.Handlers.Count == 0) {
+                    return visited;
+                }
+
+                var handlers = new List<CatchBlock>(visited.Handlers.Count);
+                foreach (CatchBlock handler in visited.Handlers) {
+                    ParameterExpression variable = handler.Variable ?? Expression.Variable(handler.Test, null);
+                    Expression exception = variable.Type == typeof(Exception) ? variable : Expression.Convert(variable, typeof(Exception));
+                    handlers.Add(Expression.MakeCatchBlock(
+                        handler.Test,
+                        variable,
+                        Expression.Block(Expression.Call(_tracer, s_onExceptionCaughtMethod, exception), handler.Body),
+                        handler.Filter));
+                }
+                return Expression.MakeTry(visited.Type, visited.Body, visited.Finally, visited.Fault, handlers);
+            }
+
+            protected override Expression VisitLambda<T>(Expression<T> node) {
+                // Don't descend into nested lambdas — they have their own debug context.
+                return node;
+            }
+
+            protected override Expression VisitExtension(Expression node) {
+                // A nested async node is sliced into its own state machine against its own label; leave it opaque.
+                if (node is AsyncExpression or GeneratorExpression
+#if NET
+                    or AsyncEnumerableExpression
+#endif
+                    ) {
+                    return node;
+                }
+                if (node is YieldExpression { Value: not null } yield && yield.Target == _yieldLabel) {
+                    return Expression.Block(
+                        base.VisitExtension(node),
+                        AsyncRewriter.RethrowIfPending(_exceptionSlot));
+                }
+                return base.VisitExtension(node);
+            }
+        }
+    }
+}

src/core/Microsoft.Dynamic/Ast/AsyncEnumerableExpression.cs

@@ -13,6 +13,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 
+using Microsoft.Scripting.Debugging.CompilerServices;
 using Microsoft.Scripting.Utils;
 
 namespace Microsoft.Scripting.Ast {
@@ -35,12 +36,33 @@ public sealed class AsyncEnumerableExpression : Expression {
 
         internal AsyncEnumerableExpression(string? name, Expression body, LabelTarget yieldLabel,
                                            Expression? cancellationToken = null,
-                                           Expression? cancellationException = null) {
+                                           Expression? cancellationException = null,
+                                           DebugContext? debugContext = null,
+                                           DebugLambdaInfo? debugLambdaInfo = null) {
             Name = name;
             Body = body;
             YieldLabel = yieldLabel;
             CancellationToken = cancellationToken ?? Utils.DefaultCancellationToken;
             CancellationException = cancellationException ?? Utils.DefaultCancellationException;
+            DebugContext = debugContext;
+            DebugLambdaInfo = debugLambdaInfo;
+        }
+
+        /// <summary>
+        ///   Optional tracing hook; see <see cref="AsyncExpression.DebugContext"/>. When both are non-null,
+        ///   <see cref="Reduce"/> instruments the body with debug points so it becomes traceable.
+        /// </summary>
+        internal DebugContext? DebugContext { get; }
+
+        internal DebugLambdaInfo? DebugLambdaInfo { get; }
+
+        /// <summary>
+        ///   Returns a copy of this node that instruments its body for tracing under the given debug context.
+        /// </summary>
+        public AsyncEnumerableExpression WithDebugInfo(DebugContext debugContext, DebugLambdaInfo debugLambdaInfo) {
+            ContractUtils.RequiresNotNull(debugContext, nameof(debugContext));
+            ContractUtils.RequiresNotNull(debugLambdaInfo, nameof(debugLambdaInfo));
+            return new AsyncEnumerableExpression(Name, Body, YieldLabel, CancellationToken, CancellationException, debugContext, debugLambdaInfo);
         }
 
         /// <summary>Optional diagnostic name (forwarded to the inner generator).</summary>
@@ -76,7 +98,7 @@ protected override Expression VisitChildren(ExpressionVisitor visitor) {
             Expression ct = visitor.Visit(CancellationToken);
             Expression ce = visitor.Visit(CancellationException);
             if (b == Body && ct == CancellationToken && ce == CancellationException) return this;
-            return new AsyncEnumerableExpression(Name, b, YieldLabel, ct, ce);
+            return new AsyncEnumerableExpression(Name, b, YieldLabel, ct, ce, DebugContext, DebugLambdaInfo);
         }
     }
 

src/core/Microsoft.Dynamic/Ast/AsyncEnumerableRewriter.cs

@@ -11,7 +11,6 @@
 using System.Linq.Expressions;
 using System.Reflection;
 using System.Runtime.CompilerServices;
-using System.Runtime.ExceptionServices;
 using System.Threading.Tasks;
 
 namespace Microsoft.Scripting.Ast {
@@ -28,8 +27,6 @@ private static readonly ConstructorInfo s_awaitPointCtor
             = typeof(Runtime.AwaitPoint).GetConstructor([typeof(Task)])!;
         private static readonly FieldInfo s_valueSlotField
             = typeof(StrongBox<object?>).GetField(nameof(StrongBox<object?>.Value))!;
-        private static readonly FieldInfo s_exceptionSlotField
-            = typeof(StrongBox<Exception?>).GetField(nameof(StrongBox<Exception?>.Value))!;
         private static readonly ConstructorInfo s_valueSlotCtor
             = typeof(StrongBox<object?>).GetConstructor(Type.EmptyTypes)!;
         private static readonly ConstructorInfo s_exceptionSlotCtor
@@ -48,8 +45,17 @@ public Expression Reduce() {
             ParameterExpression valueSlot = Expression.Variable(typeof(StrongBox<object?>), "$asyncValue");
             ParameterExpression exceptionSlot = Expression.Variable(typeof(StrongBox<Exception?>), "$awaitException");
 
+            // Optional tracing, as in AsyncRewriter: the debug points are injected on the same label the body's yields and awaits target,
+            // so they all land in one state machine, which is then wrapped in a TracingStateMachine.
+            Expression body = _node.Body;
+            AsyncDebugSupport? debugSupport = null;
+            if (_node.DebugContext is not null && _node.DebugLambdaInfo is not null) {
+                debugSupport = new AsyncDebugSupport(_node.DebugContext, _node.DebugLambdaInfo, _node.Name, exceptionSlot);
+                body = debugSupport.InjectDebugPoints(body, _node.YieldLabel);
+            }
+
             var rewriter = new AwaitToAwaitPointRewriter(_node.YieldLabel, valueSlot, exceptionSlot);
-            Expression rewrittenBody = rewriter.Visit(_node.Body);
+            Expression rewrittenBody = rewriter.Visit(body);
 
             // Coerce to void for Utils.Generator (the generator body's value is discarded).
             Expression generatorBody = rewrittenBody.Type == typeof(void)
@@ -62,12 +68,13 @@ public Expression Reduce() {
                 generatorBody,
                 typeof(IEnumerator<object>),
                 rewriteAssignments: false);
+            Expression states = debugSupport?.TraceStateMachine(generator, resultSlot: null, isAsyncEnumerable: true) ?? generator;
 
             // Argument order matches DriveAsyncEnumerable: ..., cancellationToken, cancellationException
             // (same as DriveAsync — cancellationToken is the last required parameter).
             Expression drive = Expression.Call(
                 s_driveMethod,
-                generator,
+                states,
                 valueSlot,
                 exceptionSlot,
                 _node.CancellationToken,
@@ -88,11 +95,6 @@ public Expression Reduce() {
         ///   value yielded by a language-level <c>yield</c>.
         /// </summary>
         private sealed class AwaitToAwaitPointRewriter : ExpressionVisitor {
-            private static readonly MethodInfo s_captureMethod
-                = typeof(ExceptionDispatchInfo).GetMethod(nameof(ExceptionDispatchInfo.Capture))!;
-            private static readonly MethodInfo s_throwMethod
-                = typeof(ExceptionDispatchInfo).GetMethod(nameof(ExceptionDispatchInfo.Throw), Type.EmptyTypes)!;
-
             private readonly LabelTarget _yieldLabel;
             private readonly ParameterExpression _valueSlot;
             private readonly ParameterExpression _exceptionSlot;
@@ -109,20 +111,12 @@ protected override Expression VisitExtension(Expression node) {
                     // Wrap the awaited Task in an AwaitPoint marker, then box to object for the yield.
                     Expression awaitPoint = Expression.New(s_awaitPointCtor, Expression.Convert(operand, typeof(Task)));
                     Expression yielded = Expression.Convert(awaitPoint, typeof(object));
-
-                    Expression readException = Expression.Field(_exceptionSlot, s_exceptionSlotField);
                     Expression readSlot = Expression.Field(_valueSlot, s_valueSlotField);
 
-                    Expression rethrow = Expression.IfThen(
-                        Expression.ReferenceNotEqual(readException, Expression.Constant(null, typeof(Exception))),
-                        Expression.Call(
-                            Expression.Call(s_captureMethod, readException),
-                            s_throwMethod));
-
                     return Expression.Block(
                         typeof(object),
                         Utils.YieldReturn(_yieldLabel, yielded),
-                        rethrow,
+                        AsyncRewriter.RethrowIfPending(_exceptionSlot),
                         readSlot);
                 }
                 return base.VisitExtension(node);