Allow overriding the AsyncMethodBuilder on methods

src/Compilers/CSharp/Portable/CSharpResources.resx

@@ -6585,6 +6585,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>

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

@@ -225,6 +225,7 @@ internal enum MessageID
         IDS_FeatureLambdaAttributes = MessageBase + 12800,
         IDS_FeatureWithOnAnonymousTypes = MessageBase + 12801,
         IDS_FeatureExtendedPropertyPatterns = MessageBase + 12802,
+        IDS_AsyncMethodBuilderOverride = MessageBase + 12803,
     }
 
     // Message IDs may refer to strings that need to be localized.
@@ -342,6 +343,7 @@ internal static LanguageVersion RequiredVersion(this MessageID feature)
                 case MessageID.IDS_FeatureInferredDelegateType: // semantic check
                 case MessageID.IDS_FeatureLambdaAttributes: // semantic check
                 case MessageID.IDS_FeatureExtendedPropertyPatterns:
+                case MessageID.IDS_AsyncMethodBuilderOverride: // semantic check
                     return LanguageVersion.Preview;
 
                 // C# 9.0 features.

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(
@@ -189,11 +189,33 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                 bool customBuilder = returnType.IsCustomTaskType(out builderArgument);
                 if (customBuilder)
                 {
-                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: false);
-                    if ((object)builderType != null)
+                    if (method.HasMethodLevelBuilder(out var methodLevelBuilder))
+                    {
+                        var initialBuilderType = ValidateBuilderType(F, methodLevelBuilder, returnType.DeclaredAccessibility, isGeneric: false, forOverride: true);
+                        if ((object)initialBuilderType != null)
+                        {
+                            // We only get the Create method from the initial builder type, then we'll use its return type as builder type
+                            createBuilderMethod = GetCustomCreateMethod(F, initialBuilderType, forOverride: true);
+                            builderType = createBuilderMethod?.ReturnType as NamedTypeSymbol;
+                            if ((object)builderType != null)
+                            {
+                                taskProperty = GetCustomTaskProperty(F, builderType, returnType);
+                                validateIgnoredBuilderType(F, returnType, builderArgument, isGeneric: false);
+                            }
+                        }
+                        else
+                        {
+                            builderType = null;
+                        }
+                    }
+                    else
                     {
-                        taskProperty = GetCustomTaskProperty(F, builderType, returnType);
-                        createBuilderMethod = GetCustomCreateMethod(F, builderType);
+                        builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: false);
+                        if ((object)builderType != null)
+                        {
+                            taskProperty = GetCustomTaskProperty(F, builderType, returnType);
+                            createBuilderMethod = GetCustomCreateMethod(F, builderType);
+                        }
                     }
                 }
                 else
@@ -217,7 +239,7 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     (object)createBuilderMethod == null ||
                     (object)taskProperty == null)
                 {
-                    collection = default(AsyncMethodBuilderMemberCollection);
+                    collection = default;
                     return false;
                 }
                 return TryCreate(
@@ -257,12 +279,35 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                 bool customBuilder = returnType.IsCustomTaskType(out builderArgument);
                 if (customBuilder)
                 {
-                    builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: true);
-                    if ((object)builderType != null)
+                    if (method.HasMethodLevelBuilder(out var methodLevelBuilder))
                     {
-                        builderType = builderType.ConstructedFrom.Construct(resultType);
-                        taskProperty = GetCustomTaskProperty(F, builderType, returnType);
-                        createBuilderMethod = GetCustomCreateMethod(F, builderType);
+                        var initialBuilderType = ValidateBuilderType(F, methodLevelBuilder, returnType.DeclaredAccessibility, isGeneric: true, forOverride: true);
+                        if ((object)initialBuilderType != null)
+                        {
+                            // We only get the Create method from the initial builder type, then we'll use its return type as builder type
+                            initialBuilderType = initialBuilderType.ConstructedFrom.Construct(resultType);
+                            createBuilderMethod = GetCustomCreateMethod(F, initialBuilderType, forOverride: true);
+                            builderType = createBuilderMethod?.ReturnType as NamedTypeSymbol;
+                            if ((object)builderType != null)
+                            {
+                                taskProperty = GetCustomTaskProperty(F, builderType, returnType);
+                                validateIgnoredBuilderType(F, returnType, builderArgument, isGeneric: true);
+                            }
+                        }
+                        else
+                        {
+                            builderType = null;
+                        }
+                    }
+                    else
+                    {
+                        builderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric: true);
+                        if ((object)builderType != null)
+                        {
+                            builderType = builderType.ConstructedFrom.Construct(resultType);
+                            taskProperty = GetCustomTaskProperty(F, builderType, returnType);
+                            createBuilderMethod = GetCustomCreateMethod(F, builderType);
+                        }
                     }
                 }
                 else
@@ -287,7 +332,7 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
                     (object)taskProperty == null ||
                     (object)createBuilderMethod == null)
                 {
-                    collection = default(AsyncMethodBuilderMemberCollection);
+                    collection = default;
                     return false;
                 }
                 return TryCreate(
@@ -307,16 +352,52 @@ internal static bool TryCreate(SyntheticBoundNodeFactory F, MethodSymbol method,
             }
 
             throw ExceptionUtilities.UnexpectedValue(method);
+
+            void validateIgnoredBuilderType(SyntheticBoundNodeFactory F, NamedTypeSymbol returnType, object builderArgument, bool isGeneric)
+            {
+                var ignoredBuilderType = ValidateBuilderType(F, builderArgument, returnType.DeclaredAccessibility, isGeneric);
+
+                if ((object)ignoredBuilderType != null)
+                {
+                    if (isGeneric)
+                    {
+                        var resultType = returnType.TypeArgumentsWithAnnotationsNoUseSiteDiagnostics.Single().Type;
+                        ignoredBuilderType = ignoredBuilderType.ConstructedFrom.Construct(resultType);
+                    }
+
+                    GetCustomCreateMethod(F, ignoredBuilderType, forOverride: false);
+                    GetCustomTaskProperty(F, ignoredBuilderType, returnType);
+
+                    _ = TryGetBuilderMember(F,
+                            isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__SetException : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__SetException,
+                            ignoredBuilderType, customBuilder: true, out MethodSymbol _) &&
+                        TryGetBuilderMember(F,
+                            isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__SetResult : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__SetResult,
+                            ignoredBuilderType, customBuilder: true, out MethodSymbol _) &&
+                        TryGetBuilderMember(F,
+                            isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__AwaitOnCompleted : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__AwaitOnCompleted,
+                            ignoredBuilderType, customBuilder: true, out MethodSymbol _) &&
+                        TryGetBuilderMember(F,
+                           isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__AwaitUnsafeOnCompleted : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__AwaitUnsafeOnCompleted,
+                           ignoredBuilderType, customBuilder: true, out MethodSymbol _) &&
+                        TryGetBuilderMember(F,
+                           isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__Start_T : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__Start_T,
+                           ignoredBuilderType, customBuilder: true, out MethodSymbol _) &&
+                        TryGetBuilderMember(F,
+                            isGeneric ? WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder_T__SetStateMachine : WellKnownMember.System_Runtime_CompilerServices_AsyncTaskMethodBuilder__SetStateMachine,
+                            ignoredBuilderType, customBuilder: true, out MethodSymbol _);
+                }
+            }
         }
 
-        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 forOverride = false)
         {
             var builderType = builderAttributeArgument as NamedTypeSymbol;
 
             if ((object)builderType != null &&
                  !builderType.IsErrorType() &&
                  !builderType.IsVoidType() &&
-                 builderType.DeclaredAccessibility == desiredAccessibility)
+                 (forOverride || builderType.DeclaredAccessibility == desiredAccessibility))
             {
                 bool isArityOk = isGeneric
                                  ? builderType.IsUnboundGenericType && builderType.ContainingType?.IsGenericType != true && builderType.Arity == 1
@@ -376,7 +457,7 @@ private static bool TryCreate(
                 return true;
             }
 
-            collection = default(AsyncMethodBuilderMemberCollection);
+            collection = default;
             return false;
         }
 
@@ -429,9 +510,12 @@ private static bool TryGetBuilderMember<TSymbol>(
             return true;
         }
 
+        // For method-level builders, we allow the `Create` method to return a different type.
+        // We'll just use that type as the final builder type.
         private static MethodSymbol GetCustomCreateMethod(
             SyntheticBoundNodeFactory F,
-            NamedTypeSymbol builderType)
+            NamedTypeSymbol builderType,
+            bool forOverride = false)
         {
             // The Create method's return type is expected to be builderType.
             // The WellKnownMembers routines aren't able to enforce that, which is why this method exists.
@@ -448,7 +532,8 @@ private static MethodSymbol GetCustomCreateMethod(
                     method.IsStatic &&
                     method.ParameterCount == 0 &&
                     !method.IsGenericMethod &&
-                    method.ReturnType.Equals(builderType, TypeCompareKind.AllIgnoreOptions))
+                    method.RefKind == RefKind.None &&
+                    (forOverride || method.ReturnType.Equals(builderType, TypeCompareKind.AllIgnoreOptions)))
                 {
                     return method;
                 }

src/Compilers/CSharp/Portable/Symbols/MethodSymbolExtensions.cs

@@ -12,6 +12,7 @@
 using System.Diagnostics;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using System.Linq;
+using System.Diagnostics.CodeAnalysis;
 
 namespace Microsoft.CodeAnalysis.CSharp.Symbols
 {
@@ -190,6 +191,32 @@ public static bool IsAsyncReturningIAsyncEnumerator(this MethodSymbol method, CS
                 && method.ReturnType.IsIAsyncEnumeratorType(compilation);
         }
 
+#nullable enable
+        /// <summary>
+        /// Returns true if the method has a [AsyncMethodBuilder(typeof(B))] attribute. If so it returns the "B".
+        /// Validation of builder type B is left for elsewhere. This method returns B without validation of any kind.
+        /// </summary>
+        internal static bool HasMethodLevelBuilder(this MethodSymbol method, [NotNullWhen(true)] out object? builderArgument)
+        {
+            Debug.Assert(method is not null);
+
+            // Find the AsyncMethodBuilder attribute.
+            foreach (var attr in method.GetAttributes())
+            {
+                if (attr.IsTargetAttribute(method, AttributeDescription.AsyncMethodBuilderAttribute)
+                    && attr.CommonConstructorArguments.Length == 1
+                    && attr.CommonConstructorArguments[0].Kind == TypedConstantKind.Type)
+                {
+                    builderArgument = attr.CommonConstructorArguments[0].ValueInternal!;
+                    return true;
+                }
+            }
+
+            builderArgument = null;
+            return false;
+        }
+#nullable disable
+
         internal static CSharpSyntaxNode ExtractReturnTypeSyntax(this MethodSymbol method)
         {
             if (method is SynthesizedSimpleProgramEntryPointSymbol synthesized)

src/Compilers/CSharp/Portable/Symbols/Source/SourceMethodSymbolWithAttributes.cs

@@ -1056,6 +1056,11 @@ protected void AsyncMethodChecks(BindingDiagnosticBag diagnostics)
                     hasErrors = true;
                 }
 
+                if (this.HasMethodLevelBuilder(out _))
+                {
+                    hasErrors = MessageID.IDS_AsyncMethodBuilderOverride.CheckFeatureAvailability(diagnostics, this.DeclaringCompilation, errorLocation);
+                }
+
                 for (NamedTypeSymbol curr = this.ContainingType; (object)curr != null; curr = curr.ContainingType)
                 {
                     var sourceNamedTypeSymbol = curr as SourceNamedTypeSymbol;

src/Compilers/CSharp/Portable/Symbols/TypeSymbolExtensions.cs

@@ -1638,7 +1638,7 @@ internal static bool IsCustomTaskType(this NamedTypeSymbol type, [NotNullWhen(tr
             var arity = type.Arity;
             if (arity < 2)
             {
-                // Find the AsyncBuilder attribute.
+                // Find the AsyncMethodBuilder attribute.
                 foreach (var attr in type.GetAttributes())
                 {
                     if (attr.IsTargetAttribute(type, AttributeDescription.AsyncMethodBuilderAttribute)