Can't figure out why ValueTask won't work here

src/EditorFeatures/CSharpTest/CodeActions/IntroduceVariable/IntroduceVariableTests.cs

@@ -7544,6 +7544,124 @@ void M()
             string {|Rename:v|} = (x.ToString());
         };
     }
+}");
+        }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestMissingReturnStatementInAsyncTaskMethod()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, Task> f = async x => await [|M2()|];
+    }
+
+    async Task M2()
+    {
+    }
+}",
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, Task> f = async x =>
+        {
+            Task {|Rename:task|} = M2();
+            await task;
+        };
+    }
+
+    async Task M2()
+    {
+    }
+}");
+        }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestMissingReturnStatementInAsyncValueTaskMethod()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask> f = async x => await [|M2()|];
+    }
+
+    async ValueTask M2()
+    {
+    }
+}",
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask> f = async x =>
+        {
+            ValueTask {|Rename:task|} = M2();
+            await task;
+        };
+    }
+
+    async ValueTask M2()
+    {
+    }
+}");
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestReturnStatementInAsyncTaskTypeMethod()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, Task<int>> f = async x => await [|M2()|];
+    }
+
+    async Task<int> M2()
+    {
+        return 0;
+    }
+}",
+@"using System;
+using System.Threading.Tasks;
+
+class Program
+{
+    void M()
+    {
+        Func<int, Task<int>> f = async x =>
+        {
+            Task<int> {|Rename:task|} = M2();
+            return await task;
+        };
+    }
+
+    async Task<int> M2()
+    {
+        return 0;
+    }
 }");
         }
     }

src/Features/CSharp/Portable/IntroduceVariable/CSharpIntroduceVariableService_IntroduceLocal.cs

@@ -96,76 +96,97 @@ internal partial class CSharpIntroduceVariableService
             var rewrittenBody = Rewrite(
                 document, expression, newLocalName, document, oldBody, allOccurrences, cancellationToken);
 
-            var doesDelegateMethodReturnVoid =
-                document.SemanticModel.GetTypeInfo(oldLambda, cancellationToken).ConvertedType is INamedTypeSymbol delegateType &&
-                delegateType.DelegateInvokeMethod != null &&
-                delegateType.DelegateInvokeMethod.ReturnsVoid;
-
+            var shouldIncludeReturnStatement = ShouldIncludeReturnStatement(document, oldLambda, cancellationToken);
             var newBody = GetNewBlockBodyForLambda(
-                declarationStatement, isEntireLambdaBodySelected, rewrittenBody, doesDelegateMethodReturnVoid);
+                declarationStatement, isEntireLambdaBodySelected, rewrittenBody, shouldIncludeReturnStatement);
+
+            // Add an elastic newline so that the formatter will place this new lambda body across multiple lines.
+            newBody = newBody.WithOpenBraceToken(newBody.OpenBraceToken.WithAppendedTrailingTrivia(SyntaxFactory.ElasticCarriageReturnLineFeed))
+                             .WithAdditionalAnnotations(Formatter.Annotation);
+
             var newLambda = oldLambda.WithBody(newBody);
 
             var newRoot = document.Root.ReplaceNode(oldLambda, newLambda);
             return document.Document.WithSyntaxRoot(newRoot);
         }
 
+        private static bool ShouldIncludeReturnStatement(
+            SemanticDocument document,
+            LambdaExpressionSyntax oldLambda,
+            CancellationToken cancellationToken)
+        {
+            if (document.SemanticModel.GetTypeInfo(oldLambda, cancellationToken).ConvertedType is INamedTypeSymbol delegateType &&
+                delegateType.DelegateInvokeMethod != null)
+            {
+                if (delegateType.DelegateInvokeMethod.ReturnsVoid)
+                {
+                    return false;
+                }
+
+                if (oldLambda.AsyncKeyword != default)
+                {
+                    if (document.SemanticModel.Compilation.TaskType().Equals(delegateType.DelegateInvokeMethod.ReturnType) ||
+                        document.SemanticModel.Compilation.ValueTaskOfTType().Equals(delegateType.DelegateInvokeMethod.ReturnType))
+                    {
+                        return false;
+                    }
+                }
+            }
+
+            return true;
+        }
+
         private static BlockSyntax GetNewBlockBodyForLambda(
             LocalDeclarationStatementSyntax declarationStatement,
             bool isEntireLambdaBodySelected,
             ExpressionSyntax rewrittenBody,
-            bool doesDelegateMethodReturnVoid)
+            bool includeReturnStatement)
         {
-            // For lambdas with void return types, we don't need to include the rewritten body if the entire lambda body
-            // was originally selected for refactoring, as the rewritten body should already be encompassed within the
-            // declaration statement.
-
-            // Case 1a: The lambda has a void return type, and the user selects the entire lambda body.
-            // e.g.:
-            //     Action<int> goo = x => [|x.ToString()|];
-            //
-            // After refactoring:
-            //     Action<int> goo = x =>
-            //     {
-            //         string v = x.ToString();
-            //     };
-            var newBody = SyntaxFactory.Block(declarationStatement);
-
-            if (doesDelegateMethodReturnVoid && !isEntireLambdaBodySelected)
+            if (includeReturnStatement)
             {
-                // Case 1b: The lambda has a void return type, and the user didn't select the entire lambda body.
+                // Case 1: The lambda has a non-void return type.
                 // e.g.:
-                //     Task.Run(() => File.Copy("src", [|Path.Combine("dir", "file")|]));
+                //     Func<int, int> f = x => [|x + 1|];
                 //
                 // After refactoring:
-                //     Task.Run(() =>
+                //     Func<int, int> f = x =>
                 //     {
-                //         string destFileName = Path.Combine("dir", "file");
-                //         File.Copy("src", destFileName);
-                //     });
-                newBody = newBody.AddStatements(
-                    SyntaxFactory.ExpressionStatement(rewrittenBody, SyntaxFactory.Token(SyntaxKind.SemicolonToken)));
+                //         var v = x + 1;
+                //         return v;
+                //     };
+                return SyntaxFactory.Block(declarationStatement, SyntaxFactory.ReturnStatement(rewrittenBody));
             }
-            else if (!doesDelegateMethodReturnVoid)
+
+            // For lambdas with void return types, we don't need to include the rewritten body if the entire lambda body
+            // was originally selected for refactoring, as the rewritten body should already be encompassed within the
+            // declaration statement.
+            if (isEntireLambdaBodySelected)
             {
-                // Case 2: The lambda has a non-void return type.
+                // Case 2a: The lambda has a void return type, and the user selects the entire lambda body.
                 // e.g.:
-                //     Func<int, int> f = x => [|x + 1|];
+                //     Action<int> goo = x => [|x.ToString()|];
                 //
                 // After refactoring:
-                //     Func<int, int> f = x =>
+                //     Action<int> goo = x =>
                 //     {
-                //         var v = x + 1;
-                //         return v;
+                //         string v = x.ToString();
                 //     };
-                newBody = newBody.AddStatements(
-                    SyntaxFactory.ReturnStatement(rewrittenBody));
+                return SyntaxFactory.Block(declarationStatement);
             }
 
-            // Add an elastic newline so that the formatter will place this new lambda body across multiple lines.
-            newBody = newBody.WithOpenBraceToken(newBody.OpenBraceToken.WithAppendedTrailingTrivia(SyntaxFactory.ElasticCarriageReturnLineFeed))
-                             .WithAdditionalAnnotations(Formatter.Annotation);
-
-            return newBody;
+            // Case 2b: The lambda has a void return type, and the user didn't select the entire lambda body.
+            // e.g.:
+            //     Task.Run(() => File.Copy("src", [|Path.Combine("dir", "file")|]));
+            //
+            // After refactoring:
+            //     Task.Run(() =>
+            //     {
+            //         string destFileName = Path.Combine("dir", "file");
+            //         File.Copy("src", destFileName);
+            //     });
+            return SyntaxFactory.Block(
+                declarationStatement,
+                SyntaxFactory.ExpressionStatement(rewrittenBody, SyntaxFactory.Token(SyntaxKind.SemicolonToken)));
         }
 
         private TypeSyntax GetTypeSyntax(SemanticDocument document, ExpressionSyntax expression, CancellationToken cancellationToken)