Async-streams: Call AsyncIteratorMethodBuilder.Complete (#39436)

docs/features/async-streams.md

@@ -176,6 +176,7 @@ IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken token)
     if (initialThreadId == /*managedThreadId*/ && state == StateMachineStates.FinishedStateMachine)
     {
         state = InitialState; // -3
+        builder = AsyncIteratorMethodBuilder.Create();
         disposeMode = false;
         result = this;
     }

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

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
-using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Diagnostics;
 using Microsoft.CodeAnalysis.CodeGen;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
@@ -71,6 +71,7 @@ protected override BoundStatement GenerateSetResultCall()
 
             // if (this.combinedTokens != null) { this.combinedTokens.Dispose(); this.combinedTokens = null; } // for enumerables only
             // this.promiseOfValueOrEnd.SetResult(false);
+            // this.builder.Complete();
             // return;
             // _exprReturnLabelTrue:
             // this.promiseOfValueOrEnd.SetResult(true);
@@ -86,6 +87,7 @@ protected override BoundStatement GenerateSetResultCall()
             builder.AddRange(
                 // this.promiseOfValueOrEnd.SetResult(false);
                 generateSetResultOnPromise(false),
+                GenerateCompleteOnBuilder(),
                 F.Return(),
                 F.Label(_exprReturnLabelTrue),
                 // this.promiseOfValueOrEnd.SetResult(true);
@@ -102,6 +104,17 @@ BoundExpressionStatement generateSetResultOnPromise(bool result)
             }
         }
 
+        private BoundExpressionStatement GenerateCompleteOnBuilder()
+        {
+            // Produce:
+            // this.builder.Complete();
+            return F.ExpressionStatement(
+                F.Call(
+                    F.Field(F.This(), _asyncMethodBuilderField),
+                        _asyncMethodBuilderMemberCollection.SetResult, // AsyncIteratorMethodBuilder.Complete is the corresponding method to AsyncTaskMethodBuilder.SetResult
+                        ImmutableArray<BoundExpression>.Empty));
+        }
+
         private void AddDisposeCombinedTokensIfNeeded(ArrayBuilder<BoundStatement> builder)
         {
             // if (this.combinedTokens != null) { this.combinedTokens.Dispose(); this.combinedTokens = null; } // for enumerables only
@@ -131,6 +144,9 @@ protected override BoundStatement GenerateSetExceptionCall(LocalSymbol exception
                 _asyncIteratorInfo.SetExceptionMethod,
                 F.Local(exceptionLocal))));
 
+            // this.builder.Complete();
+            builder.Add(GenerateCompleteOnBuilder());
+
             return F.Block(builder.ToImmutableAndFree());
         }
 

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

@@ -133,7 +133,7 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     createBuilderMethod: createBuilderMethod,
                     taskProperty: null,
                     setException: null, // unused
-                    setResult: WellKnownMember.System_Runtime_CompilerServices_AsyncIteratorMethodBuilder__Complete,
+                    setResult: WellKnownMember.System_Runtime_CompilerServices_AsyncIteratorMethodBuilder__Complete, // AsyncIteratorMethodBuilder.Complete is the corresponding method to AsyncTaskMethodBuilder.SetResult
                     awaitOnCompleted: WellKnownMember.System_Runtime_CompilerServices_AsyncIteratorMethodBuilder__AwaitOnCompleted,
                     awaitUnsafeOnCompleted: WellKnownMember.System_Runtime_CompilerServices_AsyncIteratorMethodBuilder__AwaitUnsafeOnCompleted,
                     start: WellKnownMember.System_Runtime_CompilerServices_AsyncIteratorMethodBuilder__MoveNext_T,

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

@@ -28,12 +28,12 @@ internal class AsyncMethodToStateMachineRewriter : MethodToStateMachineRewriter
         /// <see cref="AsyncVoidMethodBuilder"/>, <see cref="AsyncTaskMethodBuilder"/>, or <see cref="AsyncTaskMethodBuilder{TResult}"/> depending on the
         /// return type of the async method.
         /// </summary>
-        private readonly FieldSymbol _asyncMethodBuilderField;
+        protected readonly FieldSymbol _asyncMethodBuilderField;
 
         /// <summary>
         /// A collection of well-known members for the current async method builder.
         /// </summary>
-        private readonly AsyncMethodBuilderMemberCollection _asyncMethodBuilderMemberCollection;
+        protected readonly AsyncMethodBuilderMemberCollection _asyncMethodBuilderMemberCollection;
 
         /// <summary>
         /// The exprReturnLabel is used to label the return handling code at the end of the async state-machine

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

@@ -148,7 +148,7 @@ protected override void GenerateConstructor()
                 // {
                 //     this.state = state;
                 //     this.initialThreadId = {managedThreadId};
-                //     this.builder = System.Runtime.CompilerServices.AsyncVoidMethodBuilder.Create();
+                //     this.builder = System.Runtime.CompilerServices.AsyncIteratorMethodBuilder.Create();
                 // }
                 Debug.Assert(stateMachineType.Constructor is IteratorConstructor);
 
@@ -164,23 +164,25 @@ protected override void GenerateConstructor()
                     bodyBuilder.Add(F.Assignment(F.InstanceField(initialThreadIdField), managedThreadId));
                 }
 
-                // this.builder = System.Runtime.CompilerServices.AsyncVoidMethodBuilder.Create();
-                AsyncMethodBuilderMemberCollection methodScopeAsyncMethodBuilderMemberCollection;
-                bool found = AsyncMethodBuilderMemberCollection.TryCreate(F, method, typeMap: null, out methodScopeAsyncMethodBuilderMemberCollection);
-                Debug.Assert(found);
-
-                bodyBuilder.Add(
-                    F.Assignment(
-                        F.InstanceField(_builderField),
-                        F.StaticCall(
-                            null,
-                            methodScopeAsyncMethodBuilderMemberCollection.CreateBuilder)));
+                // this.builder = System.Runtime.CompilerServices.AsyncIteratorMethodBuilder.Create();
+                bodyBuilder.Add(GenerateCreateAndAssignBuilder());
 
                 bodyBuilder.Add(F.Return());
                 F.CloseMethod(F.Block(bodyBuilder.ToImmutableAndFree()));
                 bodyBuilder = null;
             }
 
+            private BoundExpressionStatement GenerateCreateAndAssignBuilder()
+            {
+                // Produce:
+                // this.builder = System.Runtime.CompilerServices.AsyncIteratorMethodBuilder.Create();
+                return F.Assignment(
+                    F.InstanceField(_builderField),
+                    F.StaticCall(
+                        null,
+                        _asyncMethodBuilderMemberCollection.CreateBuilder));
+            }
+
             protected override void InitializeStateMachine(ArrayBuilder<BoundStatement> bodyBuilder, NamedTypeSymbol frameType, LocalSymbol stateMachineLocal)
             {
                 // var stateMachineLocal = new {StateMachineType}({initialState})
@@ -632,10 +634,23 @@ private void GenerateIAsyncEnumerableImplementation_GetAsyncEnumerator()
                 GenerateIteratorGetEnumerator(IAsyncEnumerableOfElementType_GetEnumerator, ref managedThreadId, initialState: StateMachineStates.InitialAsyncIteratorStateMachine);
             }
 
-            protected override BoundStatement GetExtraResetForIteratorGetEnumerator()
+            protected override void GenerateResetInstance(ArrayBuilder<BoundStatement> builder, int initialState)
             {
-                // disposeMode = false;
-                return F.Assignment(F.InstanceField(_disposeModeField), F.Literal(false));
+                // this.state = {initialState};
+                // this.builder = System.Runtime.CompilerServices.AsyncIteratorMethodBuilder.Create();
+                // this.disposeMode = false;
+
+                builder.Add(
+                    // this.state = {initialState};
+                    F.Assignment(F.Field(F.This(), stateField), F.Literal(initialState)));
+
+                builder.Add(
+                    // this.builder = System.Runtime.CompilerServices.AsyncIteratorMethodBuilder.Create();
+                    GenerateCreateAndAssignBuilder());
+
+                builder.Add(
+                    // disposeMode = false;
+                    F.Assignment(F.InstanceField(_disposeModeField), F.Literal(false)));
             }
 
             protected override void GenerateMoveNext(SynthesizedImplementationMethod moveNextMethod)

src/Compilers/CSharp/Portable/Lowering/StateMachineRewriter/StateMachineRewriter.cs

@@ -402,20 +402,12 @@ protected SynthesizedImplementationMethod GenerateIteratorGetEnumerator(MethodSy
                 managedThreadId = MakeCurrentThreadId();
 
                 var thenBuilder = ArrayBuilder<BoundStatement>.GetInstance(4);
-                thenBuilder.Add(
-                    // this.state = {initialState};
-                    F.Assignment(F.Field(F.This(), stateField), F.Literal(initialState)));
+                GenerateResetInstance(thenBuilder, initialState);
 
                 thenBuilder.Add(
                     // result = this;
                     F.Assignment(F.Local(resultVariable), F.This()));
 
-                var extraReset = GetExtraResetForIteratorGetEnumerator();
-                if (extraReset != null)
-                {
-                    thenBuilder.Add(extraReset);
-                }
-
                 if (method.IsStatic || method.ThisParameter.Type.IsReferenceType)
                 {
                     // if this is a reference type, no need to copy it since it is not assignable
@@ -473,6 +465,16 @@ protected SynthesizedImplementationMethod GenerateIteratorGetEnumerator(MethodSy
             return getEnumerator;
         }
 
+        /// <summary>
+        /// Generate logic to reset the current instance (rather than creating a new instance)
+        /// </summary>
+        protected virtual void GenerateResetInstance(ArrayBuilder<BoundStatement> builder, int initialState)
+        {
+            builder.Add(
+                // this.state = {initialState};
+                F.Assignment(F.Field(F.This(), stateField), F.Literal(initialState)));
+        }
+
         protected virtual BoundStatement InitializeParameterField(MethodSymbol getEnumeratorMethod, ParameterSymbol parameter, BoundExpression resultParameter, BoundExpression parameterProxy)
         {
             Debug.Assert(!method.IsIterator || !method.IsAsync); // an override handles async-iterators
@@ -481,12 +483,6 @@ protected virtual BoundStatement InitializeParameterField(MethodSymbol getEnumer
             return F.Assignment(resultParameter, parameterProxy);
         }
 
-        /// <summary>
-        /// Async-iterator methods use a GetAsyncEnumerator method just like the GetEnumerator of iterator methods.
-        /// But they need to do a bit more work (to reset the dispose mode).
-        /// </summary>
-        protected virtual BoundStatement GetExtraResetForIteratorGetEnumerator() => null;
-
         /// <summary>
         /// Returns true if either Thread.ManagedThreadId or Environment.CurrentManagedThreadId are available
         /// </summary>