EnC: Generate correct state mapping for methods containing await fore…

…ach (#69806)

src/Compilers/CSharp/Portable/Binder/Binder_Await.cs

@@ -39,7 +39,7 @@ private BoundAwaitExpression BindAwait(BoundExpression expression, SyntaxNode no
             // non-void return type T, the await-expression is classified as a value of type T.
             TypeSymbol awaitExpressionType = info.GetResult?.ReturnType ?? (hasErrors ? CreateErrorType() : Compilation.DynamicType);
 
-            return new BoundAwaitExpression(node, expression, info, awaitExpressionType, hasErrors);
+            return new BoundAwaitExpression(node, expression, info, debugInfo: default, awaitExpressionType, hasErrors);
         }
 
         internal void ReportBadAwaitDiagnostics(SyntaxNodeOrToken nodeOrToken, BindingDiagnosticBag diagnostics, ref bool hasErrors)

src/Compilers/CSharp/Portable/BoundTree/BoundAwaitExpressionDebugInfo.cs

@@ -0,0 +1,36 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.CodeAnalysis.CodeGen;
+
+namespace Microsoft.CodeAnalysis.CSharp;
+
+/// <summary>
+/// Debug info associated with <see cref="BoundAwaitExpression"/> to support EnC.
+/// </summary>
+/// <param name="AwaitId">The id of the generated await.</param>
+/// <param name="ReservedStateMachineCount">
+/// The number of async state machine states to reserve.
+/// 
+/// Any time multiple <see cref="BoundAwaitExpression"/>s might be associated with the same syntax node
+/// we need to make sure that the same number of state machine states gets allocated for the node,
+/// regardless of the actual number of <see cref="BoundAwaitExpression"/>s that get emitted.
+/// 
+/// To do so one or more of the emitted <see cref="BoundAwaitExpression"/>s may
+/// reserve additional dummy state machine states so that the total number of states
+/// (one for each <see cref="BoundAwaitExpression"/> plus total reserved states) is constant
+/// regardless of semantics of the syntax node.
+/// 
+/// E.g. `await foreach` produces at least one and at most two <see cref="BoundAwaitExpression"/>s:
+/// one for MoveNextAsync and the other for DisposeAsync.
+/// 
+/// If the enumerator is async-disposable it produces two <see cref="BoundAwaitExpression"/>s with 
+/// <paramref name="ReservedStateMachineCount"/> set to 0.
+/// 
+/// If the enumerator is not async-disposable it produces a single <see cref="BoundAwaitExpression"/> with 
+/// <paramref name="ReservedStateMachineCount"/> set to 1.
+/// 
+/// The states are only reserved in DEBUG builds.
+/// </param>
+internal readonly record struct BoundAwaitExpressionDebugInfo(AwaitDebugId AwaitId, byte ReservedStateMachineCount);

src/Compilers/CSharp/Portable/BoundTree/BoundNodes.xml

@@ -690,6 +690,7 @@
 
     <Field Name="Expression" Type="BoundExpression"/>
     <Field Name="AwaitableInfo" Type="BoundAwaitableInfo" Null="disallow"/>
+    <Field Name="DebugInfo" Type="BoundAwaitExpressionDebugInfo" Null="NotApplicable"/>
   </Node>
 
   <AbstractNode Name="BoundTypeOf" Base="BoundExpression">

src/Compilers/CSharp/Portable/Lowering/AsyncRewriter/AsyncIteratorMethodToStateMachineRewriter.cs

@@ -262,7 +262,7 @@ public override BoundNode VisitYieldReturnStatement(BoundYieldReturnStatement no
             //  _promiseOfValueOrEnd.SetResult(true);
             //  return;
 
-            AddResumableState(_iteratorStateAllocator, node.Syntax, out var stateNumber, out GeneratedLabelSymbol resumeLabel);
+            AddResumableState(_iteratorStateAllocator, node.Syntax, awaitId: default, out var stateNumber, out GeneratedLabelSymbol resumeLabel);
 
             var rewrittenExpression = (BoundExpression)Visit(node.Expression);
             var blockBuilder = ArrayBuilder<BoundStatement>.GetInstance();

src/Compilers/CSharp/Portable/Lowering/AsyncRewriter/AsyncMethodToStateMachineRewriter.cs

@@ -380,7 +380,7 @@ private BoundBlock VisitAwaitExpression(BoundAwaitExpression node, BoundExpressi
                     // if(!($awaiterTemp.IsCompleted)) { ... }
                     F.If(
                         condition: F.Not(GenerateGetIsCompleted(awaiterTemp, isCompletedMethod)),
-                        thenClause: GenerateAwaitForIncompleteTask(awaiterTemp)));
+                        thenClause: GenerateAwaitForIncompleteTask(awaiterTemp, node.DebugInfo)));
             BoundExpression getResultCall = MakeCallMaybeDynamic(
                 F.Local(awaiterTemp),
                 getResult,
@@ -450,9 +450,10 @@ private BoundExpression GenerateGetIsCompleted(LocalSymbol awaiterTemp, MethodSy
             return F.Call(F.Local(awaiterTemp), getIsCompletedMethod);
         }
 
-        private BoundBlock GenerateAwaitForIncompleteTask(LocalSymbol awaiterTemp)
+        private BoundBlock GenerateAwaitForIncompleteTask(LocalSymbol awaiterTemp, BoundAwaitExpressionDebugInfo debugInfo)
         {
-            AddResumableState(awaiterTemp.GetDeclaratorSyntax(), out StateMachineState stateNumber, out GeneratedLabelSymbol resumeLabel);
+            var awaitSyntax = awaiterTemp.GetDeclaratorSyntax();
+            AddResumableState(awaitSyntax, debugInfo.AwaitId, out StateMachineState stateNumber, out GeneratedLabelSymbol resumeLabel);
 
             TypeSymbol awaiterFieldType = awaiterTemp.Type.IsVerifierReference()
                 ? F.SpecialType(SpecialType.System_Object)
@@ -485,6 +486,15 @@ private BoundBlock GenerateAwaitForIncompleteTask(LocalSymbol awaiterTemp)
             blockBuilder.Add(
                 GenerateReturn(false));
 
+            if (F.Compilation.Options.EnableEditAndContinue)
+            {
+                for (int i = 0; i < debugInfo.ReservedStateMachineCount; i++)
+                {
+                    AddResumableState(awaitSyntax, new AwaitDebugId((byte)(debugInfo.AwaitId.RelativeStateOrdinal + 1 + i)), out _, out var dummyResumeLabel);
+                    blockBuilder.Add(F.Label(dummyResumeLabel));
+                }
+            }
+
             blockBuilder.Add(
                 F.Label(resumeLabel));
 

src/Compilers/CSharp/Portable/Lowering/IteratorRewriter/IteratorMethodToStateMachineRewriter.cs

@@ -333,7 +333,7 @@ public override BoundNode VisitYieldReturnStatement(BoundYieldReturnStatement no
             //     <next_state_label>: ;
             //     <hidden sequence point>
             //     this.state = finalizeState;
-            AddResumableState(node.Syntax, out StateMachineState stateNumber, out GeneratedLabelSymbol resumeLabel);
+            AddResumableState(node.Syntax, awaitId: default, out StateMachineState stateNumber, out GeneratedLabelSymbol resumeLabel);
             _currentFinallyFrame.AddState(stateNumber);
 
             var rewrittenExpression = (BoundExpression)Visit(node.Expression);
@@ -464,12 +464,12 @@ private IteratorFinallyFrame PushFrame(BoundTryStatement statement)
         {
             var syntax = statement.Syntax;
 
-            if (slotAllocatorOpt?.TryGetPreviousStateMachineState(syntax, out var finalizeState) != true)
+            if (slotAllocatorOpt?.TryGetPreviousStateMachineState(syntax, awaitId: default, out var finalizeState) != true)
             {
                 finalizeState = _nextFinalizeState--;
             }
 
-            AddStateDebugInfo(syntax, finalizeState);
+            AddStateDebugInfo(syntax, awaitId: default, finalizeState);
 
             var finallyMethod = MakeSynthesizedFinally(finalizeState);
             var newFrame = new IteratorFinallyFrame(_currentFinallyFrame, finalizeState, finallyMethod, _yieldsInTryAnalysis.Labels(statement));

src/Compilers/CSharp/Portable/Lowering/LocalRewriter/LocalRewriter_Await.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Immutable;
+using Microsoft.CodeAnalysis.CodeGen;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
 
 namespace Microsoft.CodeAnalysis.CSharp
@@ -19,9 +20,9 @@ public BoundExpression VisitAwaitExpression(BoundAwaitExpression node, bool used
             return RewriteAwaitExpression((BoundExpression)base.VisitAwaitExpression(node)!, used);
         }
 
-        private BoundExpression RewriteAwaitExpression(SyntaxNode syntax, BoundExpression rewrittenExpression, BoundAwaitableInfo awaitableInfo, TypeSymbol type, bool used)
+        private BoundExpression RewriteAwaitExpression(SyntaxNode syntax, BoundExpression rewrittenExpression, BoundAwaitableInfo awaitableInfo, TypeSymbol type, BoundAwaitExpressionDebugInfo debugInfo, bool used)
         {
-            return RewriteAwaitExpression(new BoundAwaitExpression(syntax, rewrittenExpression, awaitableInfo, type) { WasCompilerGenerated = true }, used);
+            return RewriteAwaitExpression(new BoundAwaitExpression(syntax, rewrittenExpression, awaitableInfo, debugInfo, type) { WasCompilerGenerated = true }, used);
         }
 
         /// <summary>