Allow overriding the AsyncMethodBuilder on methods

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

@@ -8640,15 +8640,15 @@ private BoundConditionalAccess BindConditionalAccessExpression(ConditionalAccess
                             break;
 
                         case SyntaxKind.SimpleLambdaExpression:
-                            resultIsUsed = (((SimpleLambdaExpressionSyntax)parent).Body != node) || ContainingMethodOrLambdaRequiresValue();
+                            resultIsUsed = (((SimpleLambdaExpressionSyntax)parent).Body != node) || MethodOrLambdaRequiresValue(ContainingMemberOrLambda, Compilation);
                             break;
 
                         case SyntaxKind.ParenthesizedLambdaExpression:
-                            resultIsUsed = (((ParenthesizedLambdaExpressionSyntax)parent).Body != node) || ContainingMethodOrLambdaRequiresValue();
+                            resultIsUsed = (((ParenthesizedLambdaExpressionSyntax)parent).Body != node) || MethodOrLambdaRequiresValue(ContainingMemberOrLambda, Compilation);
                             break;
 
                         case SyntaxKind.ArrowExpressionClause:
-                            resultIsUsed = (((ArrowExpressionClauseSyntax)parent).Expression != node) || ContainingMethodOrLambdaRequiresValue();
+                            resultIsUsed = (((ArrowExpressionClauseSyntax)parent).Expression != node) || MethodOrLambdaRequiresValue(ContainingMemberOrLambda, Compilation);
                             break;
 
                         case SyntaxKind.ForStatement:
@@ -8678,13 +8678,11 @@ private BoundConditionalAccess BindConditionalAccessExpression(ConditionalAccess
             return new BoundConditionalAccess(node, receiver, access, accessType);
         }
 
-        private bool ContainingMethodOrLambdaRequiresValue()
+        internal static bool MethodOrLambdaRequiresValue(Symbol symbol, CSharpCompilation compilation)
         {
-            var containingMethod = ContainingMemberOrLambda as MethodSymbol;
-            return
-                (object)containingMethod == null ||
-                    !containingMethod.ReturnsVoid &&
-                    !containingMethod.IsAsyncReturningTask(this.Compilation);
+            return symbol is MethodSymbol method &&
+                !method.ReturnsVoid &&
+                !method.IsAsyncEffectivelyReturningTask(compilation);
         }
 
         private BoundConditionalAccess GenerateBadConditionalAccessNodeError(ConditionalAccessExpressionSyntax node, BoundExpression receiver, BoundExpression access, BindingDiagnosticBag diagnostics)

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

@@ -623,7 +623,7 @@ BoundBlock runAnalysis(BoundBlock block, BindingDiagnosticBag blockDiagnostics)
 
         private bool ImplicitReturnIsOkay(MethodSymbol method)
         {
-            return method.ReturnsVoid || method.IsIterator || method.IsAsyncReturningTask(this.Compilation);
+            return method.ReturnsVoid || method.IsIterator || method.IsAsyncEffectivelyReturningTask(this.Compilation);
         }
 
         public BoundStatement BindExpressionStatement(ExpressionStatementSyntax node, BindingDiagnosticBag diagnostics)
@@ -2656,16 +2656,16 @@ protected bool IsInAsyncMethod()
             return IsInAsyncMethod(this.ContainingMemberOrLambda as MethodSymbol);
         }
 
-        protected bool IsTaskReturningAsyncMethod()
+        protected bool IsEffectivelyTaskReturningAsyncMethod()
         {
             var symbol = this.ContainingMemberOrLambda;
-            return symbol?.Kind == SymbolKind.Method && ((MethodSymbol)symbol).IsAsyncReturningTask(this.Compilation);
+            return symbol?.Kind == SymbolKind.Method && ((MethodSymbol)symbol).IsAsyncEffectivelyReturningTask(this.Compilation);
         }
 
-        protected bool IsGenericTaskReturningAsyncMethod()
+        protected bool IsEffectivelyGenericTaskReturningAsyncMethod()
         {
             var symbol = this.ContainingMemberOrLambda;
-            return symbol?.Kind == SymbolKind.Method && ((MethodSymbol)symbol).IsAsyncReturningGenericTask(this.Compilation);
+            return symbol?.Kind == SymbolKind.Method && ((MethodSymbol)symbol).IsAsyncEffectivelyReturningGenericTask(this.Compilation);
         }
 
         protected bool IsIAsyncEnumerableOrIAsyncEnumeratorReturningAsyncMethod()
@@ -2768,7 +2768,7 @@ private BoundStatement BindReturn(ReturnStatementSyntax syntax, BindingDiagnosti
             // on a lambda expression of unknown return type.
             if ((object)retType != null)
             {
-                if (retType.IsVoidType() || IsTaskReturningAsyncMethod())
+                if (retType.IsVoidType() || IsEffectivelyTaskReturningAsyncMethod())
                 {
                     if (arg != null)
                     {
@@ -2808,7 +2808,7 @@ private BoundStatement BindReturn(ReturnStatementSyntax syntax, BindingDiagnosti
                     if (arg == null)
                     {
                         // Error case: non-void-returning or Task<T>-returning method or lambda but just have "return;"
-                        var requiredType = IsGenericTaskReturningAsyncMethod()
+                        var requiredType = IsEffectivelyGenericTaskReturningAsyncMethod()
                             ? retType.GetMemberTypeArgumentsNoUseSiteDiagnostics().Single()
                             : retType;
 
@@ -2851,7 +2851,7 @@ internal BoundExpression CreateReturnConversion(
             {
                 Debug.Assert(returnRefKind == RefKind.None);
 
-                if (!IsGenericTaskReturningAsyncMethod())
+                if (!IsEffectivelyGenericTaskReturningAsyncMethod())
                 {
                     conversion = Conversion.NoConversion;
                     badAsyncReturnAlreadyReported = true;
@@ -2888,7 +2888,8 @@ internal BoundExpression CreateReturnConversion(
                     if (!badAsyncReturnAlreadyReported)
                     {
                         RefKind unusedRefKind;
-                        if (IsGenericTaskReturningAsyncMethod() && TypeSymbol.Equals(argument.Type, this.GetCurrentReturnType(out unusedRefKind), TypeCompareKind.ConsiderEverything2))
+                        if (IsEffectivelyGenericTaskReturningAsyncMethod()
+                            && TypeSymbol.Equals(argument.Type, this.GetCurrentReturnType(out unusedRefKind), TypeCompareKind.ConsiderEverything2))
                         {
                             // Since this is an async method, the return expression must be of type '{0}' rather than 'Task<{0}>'
                             Error(diagnostics, ErrorCode.ERR_BadAsyncReturnExpression, argument.Syntax, returnType);
@@ -3177,7 +3178,7 @@ internal BoundBlock CreateBlockFromExpression(CSharpSyntaxNode node, ImmutableAr
                     expression = BindToTypeForErrorRecovery(expression);
                     statement = new BoundReturnStatement(syntax, RefKind.None, expression) { WasCompilerGenerated = true };
                 }
-                else if (returnType.IsVoidType() || IsTaskReturningAsyncMethod())
+                else if (returnType.IsVoidType() || IsEffectivelyTaskReturningAsyncMethod())
                 {
                     // If the return type is void then the expression is required to be a legal
                     // statement expression.

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

@@ -570,21 +570,6 @@ private static TypeWithAnnotations DelegateReturnTypeWithAnnotations(MethodSymbo
             return invokeMethod.ReturnTypeWithAnnotations;
         }
 
-        private bool DelegateNeedsReturn(MethodSymbol? invokeMethod)
-        {
-            if (invokeMethod is null || invokeMethod.ReturnsVoid)
-            {
-                return false;
-            }
-
-            if (IsAsync && invokeMethod.ReturnType.IsNonGenericTaskType(this.Binder.Compilation))
-            {
-                return false;
-            }
-
-            return true;
-        }
-
         internal NamedTypeSymbol? InferDelegateType(ref CompoundUseSiteInfo<AssemblySymbol> useSiteInfo)
         {
             Debug.Assert(Binder.ContainingMemberOrLambda is { });
@@ -716,7 +701,7 @@ private BoundLambda ReallyBind(NamedTypeSymbol delegateType)
             bool reachableEndpoint = ControlFlowPass.Analyze(compilation, lambdaSymbol, block, diagnostics.DiagnosticBag);
             if (reachableEndpoint)
             {
-                if (DelegateNeedsReturn(invokeMethod))
+                if (Binder.MethodOrLambdaRequiresValue(lambdaSymbol, this.Binder.Compilation))
                 {
                     // Not all code paths return a value in {0} of type '{1}'
                     diagnostics.Add(ErrorCode.ERR_AnonymousReturnExpected, lambdaSymbol.DiagnosticLocation, this.MessageID.Localize(), delegateType);
@@ -731,8 +716,8 @@ private BoundLambda ReallyBind(NamedTypeSymbol delegateType)
             {
                 if (returnType.HasType && // Can be null if "delegateType" is not actually a delegate type.
                     !returnType.IsVoidType() &&
-                    !returnType.Type.IsNonGenericTaskType(compilation) &&
-                    !returnType.Type.IsGenericTaskType(compilation))
+                    !lambdaSymbol.IsAsyncEffectivelyReturningTask(compilation) &&
+                    !lambdaSymbol.IsAsyncEffectivelyReturningGenericTask(compilation))
                 {
                     // Cannot convert async {0} to delegate type '{1}'. An async {0} may return void, Task or Task&lt;T&gt;, none of which are convertible to '{1}'.
                     diagnostics.Add(ErrorCode.ERR_CantConvAsyncAnonFuncReturns, lambdaSymbol.DiagnosticLocation, lambdaSymbol.MessageID.Localize(), delegateType);

src/Compilers/CSharp/Portable/CSharpResources.resx

@@ -6588,6 +6588,9 @@ To remove the warning, you can use /reference instead (set the Embed Interop Typ
   <data name="IDS_FeatureWithOnAnonymousTypes" xml:space="preserve">
     <value>with on anonymous types</value>
   </data>
+  <data name="IDS_AsyncMethodBuilderOverride" xml:space="preserve">
+    <value>async method builder override</value>
+  </data>
   <data name="IDS_FeaturePositionalFieldsInRecords" xml:space="preserve">
     <value>positional fields in records</value>
   </data>
@@ -6713,4 +6716,7 @@ To remove the warning, you can use /reference instead (set the Embed Interop Typ
   <data name="HDN_DuplicateWithGlobalUsing_Title" xml:space="preserve">
     <value>The using directive appeared previously as global using</value>
   </data>
+  <data name="ERR_BuilderAttributeDisallowed" xml:space="preserve">
+    <value>The AsyncMethodBuilder attribute is disallowed on anonymous methods without an explicit return type.</value>
+  </data>
 </root>

src/Compilers/CSharp/Portable/Errors/ErrorCode.cs

@@ -1957,10 +1957,10 @@ internal enum ErrorCode
         ERR_InterfaceImplementedByUnmanagedCallersOnlyMethod = 8932,
         HDN_DuplicateWithGlobalUsing = 8933,
         ERR_CantConvAnonMethReturnType = 8934,
+        ERR_BuilderAttributeDisallowed = 8935,
 
         #endregion
 
-
         // Note: you will need to re-generate compiler code after adding warnings (eng\generate-compiler-code.cmd)
     }
 }

src/Compilers/CSharp/Portable/Errors/MessageID.cs

@@ -227,6 +227,7 @@ internal enum MessageID
         IDS_FeatureExtendedPropertyPatterns = MessageBase + 12802,
         IDS_FeatureStaticAbstractMembersInInterfaces = MessageBase + 12803,
         IDS_FeatureLambdaReturnType = MessageBase + 12804,
+        IDS_AsyncMethodBuilderOverride = MessageBase + 12805,
     }
 
     // Message IDs may refer to strings that need to be localized.
@@ -346,6 +347,7 @@ internal static LanguageVersion RequiredVersion(this MessageID feature)
                 case MessageID.IDS_FeatureExtendedPropertyPatterns:
                 case MessageID.IDS_FeatureStaticAbstractMembersInInterfaces: // semantic check
                 case MessageID.IDS_FeatureLambdaReturnType: // semantic check
+                case MessageID.IDS_AsyncMethodBuilderOverride: // semantic check
                     return LanguageVersion.Preview;
 
                 // C# 9.0 features.

src/Compilers/CSharp/Portable/FlowAnalysis/FlowAnalysisPass.cs

@@ -40,7 +40,7 @@ public static BoundBlock Rewrite(
 #endif
             var compilation = method.DeclaringCompilation;
 
-            if (method.ReturnsVoid || method.IsIterator || method.IsAsyncReturningTask(compilation))
+            if (method.ReturnsVoid || method.IsIterator || method.IsAsyncEffectivelyReturningTask(compilation))
             {
                 // we don't analyze synthesized void methods.
                 if ((method.IsImplicitlyDeclared && !method.IsScriptInitializer) ||

src/Compilers/CSharp/Portable/FlowAnalysis/NullableWalker.cs

@@ -2558,7 +2558,7 @@ private bool TryGetReturnType(out TypeWithAnnotations type, out FlowAnalysisAnno
                 return true;
             }
 
-            if (returnType.Type.IsGenericTaskType(compilation))
+            if (method.IsAsyncEffectivelyReturningGenericTask(compilation))
             {
                 type = ((NamedTypeSymbol)returnType.Type).TypeArgumentsWithAnnotationsNoUseSiteDiagnostics.Single();
                 annotations = FlowAnalysisAnnotations.None;

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

@@ -159,7 +159,7 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     out createBuilderMethod);
                 if ((object)createBuilderMethod == null)
                 {
-                    collection = default(AsyncMethodBuilderMemberCollection);
+                    collection = default;
                     return false;
                 }
                 return TryCreate(
@@ -178,18 +178,30 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     collection: out collection);
             }
 
-            if (method.IsAsyncReturningTask(F.Compilation))
+            object methodLevelBuilder = null;
+            if (method.IsAsyncEffectivelyReturningTask(F.Compilation))
             {
                 var returnType = (NamedTypeSymbol)method.ReturnType;
                 NamedTypeSymbol builderType;
                 MethodSymbol createBuilderMethod = null;
                 PropertySymbol taskProperty = null;
-
+                bool useMethodLevelBuilder = method.HasAsyncMethodBuilderAttribute(out methodLevelBuilder);
+                bool customBuilder;
                 object builderArgument;
-                bool customBuilder = returnType.IsCustomTaskType(out builderArgument);
+
+                if (useMethodLevelBuilder)
+                {
+                    customBuilder = true;
+                    builderArgument = methodLevelBuilder;
+                }
+                else
+                {
+                    customBuilder = returnType.IsCustomTaskType(out builderArgument);
+                }
+
                 if (customBuilder)
                 {
-                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: false);
+                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: false, useMethodLevelBuilder);
                     if ((object)builderType != null)
                     {
                         taskProperty = GetCustomTaskProperty(F, builderType, returnType);
@@ -213,13 +225,15 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                         customBuilder,
                         out taskProperty);
                 }
+
                 if ((object)builderType == null ||
                     (object)createBuilderMethod == null ||
                     (object)taskProperty == null)
                 {
-                    collection = default(AsyncMethodBuilderMemberCollection);
+                    collection = default;
                     return false;
                 }
+
                 return TryCreate(
                     F,
                     customBuilder: customBuilder,
@@ -236,7 +250,7 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     collection: out collection);
             }
 
-            if (method.IsAsyncReturningGenericTask(F.Compilation))
+            if (method.IsAsyncEffectivelyReturningGenericTask(F.Compilation))
             {
                 var returnType = (NamedTypeSymbol)method.ReturnType;
                 var resultType = returnType.TypeArgumentsWithAnnotationsNoUseSiteDiagnostics.Single().Type;
@@ -252,12 +266,23 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                 NamedTypeSymbol builderType;
                 MethodSymbol createBuilderMethod = null;
                 PropertySymbol taskProperty = null;
-
+                bool useMethodLevelBuilder = method.HasAsyncMethodBuilderAttribute(out methodLevelBuilder);
+                bool customBuilder;
                 object builderArgument;
-                bool customBuilder = returnType.IsCustomTaskType(out builderArgument);
+
+                if (useMethodLevelBuilder)
+                {
+                    customBuilder = true;
+                    builderArgument = methodLevelBuilder;
+                }
+                else
+                {
+                    customBuilder = returnType.IsCustomTaskType(out builderArgument);
+                }
+
                 if (customBuilder)
                 {
-                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: true);
+                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: true, useMethodLevelBuilder);
                     if ((object)builderType != null)
                     {
                         builderType = builderType.ConstructedFrom.Construct(resultType);
@@ -283,13 +308,15 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                         customBuilder,
                         out taskProperty);
                 }
+
                 if ((object)builderType == null ||
                     (object)taskProperty == null ||
                     (object)createBuilderMethod == null)
                 {
-                    collection = default(AsyncMethodBuilderMemberCollection);
+                    collection = default;
                     return false;
                 }
+
                 return TryCreate(
                     F,
                     customBuilder: customBuilder,
@@ -309,14 +336,14 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
             throw ExceptionUtilities.UnexpectedValue(method);
         }
 
-        private static NamedTypeSymbol ValidateBuilderType(SyntheticBoundNodeFactory F, object builderAttributeArgument, Accessibility desiredAccessibility, bool isGeneric)
+        private static NamedTypeSymbol ValidateBuilderType(SyntheticBoundNodeFactory F, object builderAttributeArgument, Accessibility desiredAccessibility, bool isGeneric, bool forMethodLevelBuilder = false)
         {
             var builderType = builderAttributeArgument as NamedTypeSymbol;
 
             if ((object)builderType != null &&
                  !builderType.IsErrorType() &&
                  !builderType.IsVoidType() &&
-                 builderType.DeclaredAccessibility == desiredAccessibility)
+                 (forMethodLevelBuilder || builderType.DeclaredAccessibility == desiredAccessibility))
             {
                 bool isArityOk = isGeneric
                                  ? builderType.IsUnboundGenericType && builderType.ContainingType?.IsGenericType != true && builderType.Arity == 1
@@ -376,7 +403,7 @@ private static bool TryCreate(
                 return true;
             }
 
-            collection = default(AsyncMethodBuilderMemberCollection);
+            collection = default;
             return false;
         }
 
@@ -448,6 +475,7 @@ private static MethodSymbol GetCustomCreateMethod(
                     method.IsStatic &&
                     method.ParameterCount == 0 &&
                     !method.IsGenericMethod &&
+                    method.RefKind == RefKind.None &&
                     method.ReturnType.Equals(builderType, TypeCompareKind.AllIgnoreOptions))
                 {
                     return method;