Correctly change return type of partial definition part

Author: DoctorKrolic

src/Analyzers/CSharp/CodeFixes/MakeMethodAsynchronous/CSharpMakeMethodAsynchronousCodeFixProvider.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Composition;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
@@ -67,31 +68,44 @@ protected override bool IsAsyncReturnType(ITypeSymbol type, KnownTaskTypes known
         => IsIAsyncEnumerableOrEnumerator(type, knownTypes) ||
            knownTypes.IsTaskLike(type);
 
-    protected override SyntaxNode AddAsyncTokenAndFixReturnType(
+    protected override SyntaxNode FixMethodSignature(
+        bool addAsyncModifier,
         bool keepVoid,
         IMethodSymbol methodSymbol,
         SyntaxNode node,
         KnownTaskTypes knownTypes)
     {
+        // We currently fix signature without adding 'async' modifier
+        // only for a partial definitions part of partial methods
+        Debug.Assert(addAsyncModifier || node is MethodDeclarationSyntax);
+
         return node switch
         {
-            MethodDeclarationSyntax method => FixMethod(keepVoid, methodSymbol, method, knownTypes),
+            MethodDeclarationSyntax method => FixMethod(addAsyncModifier, keepVoid, methodSymbol, method, knownTypes),
             LocalFunctionStatementSyntax localFunction => FixLocalFunction(keepVoid, methodSymbol, localFunction, knownTypes),
             AnonymousFunctionExpressionSyntax anonymous => FixAnonymousFunction(anonymous),
             _ => node,
         };
     }
 
     private static MethodDeclarationSyntax FixMethod(
+        bool addAsyncModifier,
         bool keepVoid,
         IMethodSymbol methodSymbol,
         MethodDeclarationSyntax method,
         KnownTaskTypes knownTypes)
     {
-        var (newModifiers, newReturnType) = AddAsyncModifierWithCorrectedTrivia(
-            method.Modifiers,
-            FixMethodReturnType(keepVoid, methodSymbol, method.ReturnType, knownTypes));
-        return method.WithReturnType(newReturnType).WithModifiers(newModifiers);
+        var fixedReturnType = FixMethodReturnType(keepVoid, methodSymbol, method.ReturnType, knownTypes);
+
+        if (addAsyncModifier)
+        {
+            var (newModifiers, newReturnType) = AddAsyncModifierWithCorrectedTrivia(method.Modifiers, fixedReturnType);
+            return method.WithReturnType(newReturnType).WithModifiers(newModifiers);
+        }
+        else
+        {
+            return method.WithReturnType(fixedReturnType);
+        }
     }
 
     private static LocalFunctionStatementSyntax FixLocalFunction(

src/Analyzers/CSharp/Tests/MakeMethodAsynchronous/MakeMethodAsynchronousTests.cs

@@ -1406,7 +1406,7 @@ partial void M()
             
             public partial class C
             {
-                partial void MAsync();
+                partial Task MAsync();
             }
             
             public partial class C
@@ -1440,7 +1440,7 @@ public partial void M()
             
             public partial class C
             {
-                public partial void MAsync();
+                public partial Task MAsync();
             }
             
             public partial class C

src/Analyzers/Core/CodeFixes/MakeMethodAsynchronous/AbstractMakeMethodAsynchronousCodeFixProvider.cs

@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.CodeFixes;
+using Microsoft.CodeAnalysis.Editing;
 using Microsoft.CodeAnalysis.LanguageService;
 using Microsoft.CodeAnalysis.Rename;
 using Microsoft.CodeAnalysis.Shared.Extensions;
@@ -24,8 +25,12 @@ internal abstract partial class AbstractMakeMethodAsynchronousCodeFixProvider :
 
     protected abstract bool IsAsyncReturnType(ITypeSymbol type, KnownTaskTypes knownTypes);
 
-    protected abstract SyntaxNode AddAsyncTokenAndFixReturnType(
-        bool keepVoid, IMethodSymbol methodSymbol, SyntaxNode node, KnownTaskTypes knownTypes);
+    protected abstract SyntaxNode FixMethodSignature(
+        bool addAsyncModifier,
+        bool keepVoid,
+        IMethodSymbol methodSymbol,
+        SyntaxNode node,
+        KnownTaskTypes knownTypes);
 
     public override FixAllProvider GetFixAllProvider() => WellKnownFixAllProviders.BatchFixer;
 
@@ -119,7 +124,7 @@ private async Task<Solution> FixNodeAsync(
 
         return NeedsRename()
             ? await RenameThenAddAsyncTokenAsync(keepVoid, document, node, methodSymbol, knownTypes, cancellationToken).ConfigureAwait(false)
-            : await AddAsyncTokenAsync(keepVoid, document, methodSymbol, knownTypes, node, cancellationToken).ConfigureAwait(false);
+            : await FixRelatedSignaturesAsync(keepVoid, document, methodSymbol, knownTypes, node, cancellationToken).ConfigureAwait(false);
 
         bool NeedsRename()
         {
@@ -174,26 +179,39 @@ private async Task<Solution> RenameThenAddAsyncTokenAsync(
         {
             var semanticModel = await newDocument.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
             var newMethod = (IMethodSymbol)semanticModel.GetRequiredDeclaredSymbol(newNode, cancellationToken);
-            return await AddAsyncTokenAsync(keepVoid, newDocument, newMethod, knownTypes, newNode, cancellationToken).ConfigureAwait(false);
+            return await FixRelatedSignaturesAsync(keepVoid, newDocument, newMethod, knownTypes, newNode, cancellationToken).ConfigureAwait(false);
         }
 
         return newSolution;
     }
 
-    private async Task<Solution> AddAsyncTokenAsync(
+    private async Task<Solution> FixRelatedSignaturesAsync(
         bool keepVoid,
         Document document,
         IMethodSymbol methodSymbol,
         KnownTaskTypes knownTypes,
         SyntaxNode node,
         CancellationToken cancellationToken)
     {
-        var newNode = AddAsyncTokenAndFixReturnType(keepVoid, methodSymbol, node, knownTypes);
+        var newNode = FixMethodSignature(addAsyncModifier: true, keepVoid, methodSymbol, node, knownTypes);
+
+        var solution = document.Project.Solution;
+        var solutionEditor = new SolutionEditor(solution);
+        var mainDocumentEditor = await solutionEditor.GetDocumentEditorAsync(document.Id, cancellationToken).ConfigureAwait(false);
 
-        var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
-        var newRoot = root.ReplaceNode(node, newNode);
+        mainDocumentEditor.ReplaceNode(node, newNode);
+
+        if (!keepVoid && methodSymbol.PartialDefinitionPart is { } partialDefinition)
+        {
+            var partialDefinitionNode = partialDefinition.Locations[0].FindNode(cancellationToken);
+            var fixedPartialDefinitionNode = FixMethodSignature(addAsyncModifier: false, keepVoid, methodSymbol, partialDefinitionNode, knownTypes);
+
+            var partialDefinitionDocument = solution.GetDocument(partialDefinitionNode.SyntaxTree);
+            Contract.ThrowIfNull(partialDefinitionDocument);
+            var partialDefinitionDocumentEditor = await solutionEditor.GetDocumentEditorAsync(partialDefinitionDocument.Id, cancellationToken).ConfigureAwait(false);
+            partialDefinitionDocumentEditor.ReplaceNode(partialDefinitionNode, fixedPartialDefinitionNode);
+        }
 
-        var newDocument = document.WithSyntaxRoot(newRoot);
-        return newDocument.Project.Solution;
+        return solutionEditor.GetChangedSolution();
     }
 }

src/Analyzers/VisualBasic/CodeFixes/MakeMethodAsynchronous/VisualBasicMakeMethodAsynchronousCodeFixProvider.vb

@@ -57,12 +57,17 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.MakeMethodAsynchronous
             Return knownTypes.IsTaskLike(type)
         End Function
 
-        Protected Overrides Function AddAsyncTokenAndFixReturnType(
+        Protected Overrides Function FixMethodSignature(
+                addAsyncModifier As Boolean,
                 keepVoid As Boolean,
                 methodSymbolOpt As IMethodSymbol,
                 node As SyntaxNode,
                 knownTypes As KnownTaskTypes) As SyntaxNode
 
+            ' This flag can only be false when updating partial definition method signature.
+            ' Since partial methods cannot be async in VB, it cannot be false here
+            Debug.Assert(addAsyncModifier)
+
             If node.IsKind(SyntaxKind.SingleLineSubLambdaExpression) OrElse
                node.IsKind(SyntaxKind.SingleLineFunctionLambdaExpression) Then