Merge pull request #43672 from allisonchou/IntroduceLocalLambdaBug

Fix for introduce local variable in lambda deletes unrelated code

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

@@ -7459,5 +7459,275 @@ void C(int a)
 }",
                 index: 3);
         }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestKeepExistingNonTrivialCodeInLambda()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System.IO;
+using System.Threading.Tasks;
+
+class C
+{
+    void M()
+    {
+        Task.Run(() => File.Copy(""src"", [|Path.Combine(""dir"", ""file"")|]));
+    }
+}",
+@"using System.IO;
+using System.Threading.Tasks;
+
+class C
+{
+    void M()
+    {
+        Task.Run(() =>
+        {
+            string {|Rename:destFileName|} = Path.Combine(""dir"", ""file"");
+            File.Copy(""src"", destFileName);
+        });
+    }
+}");
+        }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestIntroVarInActionSelectingInsideParens()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+
+class Program
+{
+    void M()
+    {
+        Action<int> goo = x => ([|x.ToString()|]);
+    }
+}",
+@"using System;
+
+class Program
+{
+    void M()
+    {
+        Action<int> goo = x =>
+        {
+            string {|Rename:v|} = x.ToString();
+        };
+    }
+}");
+        }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestIntroVarInActionSelectingParens()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+
+class Program
+{
+    void M()
+    {
+        Action<int> goo = x => [|(x.ToString())|];
+    }
+}",
+@"using System;
+
+class Program
+{
+    void M()
+    {
+        Action<int> goo = x =>
+        {
+            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;
+
+namespace System.Threading.Tasks {
+    struct ValueTask
+    {
+    }
+}
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask> f = async x => await [|M2()|];
+    }
+
+    async ValueTask M2()
+    {
+    }
+}",
+@"using System;
+using System.Threading.Tasks;
+
+namespace System.Threading.Tasks {
+    struct ValueTask
+    {
+    }
+}
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask> f = async x =>
+        {
+            ValueTask {|Rename:valueTask|} = M2();
+            await valueTask;
+        };
+    }
+
+    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;
+    }
+}");
+        }
+
+        [WorkItem(40745, "https://github.com/dotnet/roslyn/issues/40745")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsIntroduceVariable)]
+        public async Task TestReturnStatementInAsyncValueTaskTypeMethod()
+        {
+            await TestInRegularAndScriptAsync(
+@"using System;
+using System.Threading.Tasks;
+
+namespace System.Threading.Tasks {
+    struct ValueTask<T>
+    {
+    }
+}
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask<int>> f = async x => await [|M2()|];
+    }
+
+    async ValueTask<int> M2()
+    {
+        return 0;
+    }
+}",
+@"using System;
+using System.Threading.Tasks;
+
+namespace System.Threading.Tasks {
+    struct ValueTask<T>
+    {
+    }
+}
+
+class Program
+{
+    void M()
+    {
+        Func<int, ValueTask<int>> f = async x =>
+        {
+            ValueTask<int> {|Rename:valueTask|} = M2();
+            return await valueTask;
+        };
+    }
+
+    async ValueTask<int> M2()
+    {
+        return 0;
+    }
+}");
+        }
     }
 }

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

@@ -91,17 +91,14 @@ internal partial class CSharpIntroduceVariableService
             CancellationToken cancellationToken)
         {
             var oldBody = (ExpressionSyntax)oldLambda.Body;
+            var isEntireLambdaBodySelected = oldBody.Equals(expression.WalkUpParentheses());
 
             var rewrittenBody = Rewrite(
                 document, expression, newLocalName, document, oldBody, allOccurrences, cancellationToken);
 
-
-            var newBody =
-                document.SemanticModel.GetTypeInfo(oldLambda, cancellationToken).ConvertedType is INamedTypeSymbol delegateType
-                && delegateType.DelegateInvokeMethod != null
-                && delegateType.DelegateInvokeMethod.ReturnsVoid
-                    ? SyntaxFactory.Block(declarationStatement)
-                    : SyntaxFactory.Block(declarationStatement, SyntaxFactory.ReturnStatement(rewrittenBody));
+            var shouldIncludeReturnStatement = ShouldIncludeReturnStatement(document, oldLambda, cancellationToken);
+            var newBody = GetNewBlockBodyForLambda(
+                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))
@@ -113,6 +110,97 @@ internal partial class CSharpIntroduceVariableService
             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;
+                }
+
+                // Async lambdas with a Task or ValueTask return type don't need a return statement.
+                // e.g.:
+                //     Func<int, Task> f = async x => await M2();
+                //
+                // After refactoring:
+                //     Func<int, Task> f = async x =>
+                //     {
+                //         Task task = M2();
+                //         await task;
+                //     };
+                var compilation = document.SemanticModel.Compilation;
+                var delegateReturnType = delegateType.DelegateInvokeMethod.ReturnType;
+                if (oldLambda.AsyncKeyword != default && delegateReturnType != null)
+                {
+                    if ((compilation.TaskType() != null && delegateReturnType.Equals(compilation.TaskType())) ||
+                        (compilation.ValueTaskType() != null && delegateReturnType.Equals(compilation.ValueTaskType())))
+                    {
+                        return false;
+                    }
+                }
+            }
+
+            return true;
+        }
+
+        private static BlockSyntax GetNewBlockBodyForLambda(
+            LocalDeclarationStatementSyntax declarationStatement,
+            bool isEntireLambdaBodySelected,
+            ExpressionSyntax rewrittenBody,
+            bool includeReturnStatement)
+        {
+            if (includeReturnStatement)
+            {
+                // Case 1: The lambda has a non-void return type.
+                // e.g.:
+                //     Func<int, int> f = x => [|x + 1|];
+                //
+                // After refactoring:
+                //     Func<int, int> f = x =>
+                //     {
+                //         var v = x + 1;
+                //         return v;
+                //     };
+                return SyntaxFactory.Block(declarationStatement, SyntaxFactory.ReturnStatement(rewrittenBody));
+            }
+
+            // 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 2a: 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();
+                //     };
+                return SyntaxFactory.Block(declarationStatement);
+            }
+
+            // 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)
         {
             var typeSymbol = GetTypeSymbol(document, expression, cancellationToken);

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Extensions/ICompilationExtensions.cs

@@ -112,8 +112,12 @@ public static ImmutableArray<IAssemblySymbol> GetReferencedAssemblySymbols(this
         public static INamedTypeSymbol? TaskOfTType(this Compilation compilation)
             => compilation.GetTypeByMetadataName(typeof(Task<>).FullName!);
 
+        public static INamedTypeSymbol? ValueTaskType(this Compilation compilation)
+            => compilation.GetTypeByMetadataName("System.Threading.Tasks.ValueTask");
+
         public static INamedTypeSymbol? ValueTaskOfTType(this Compilation compilation)
             => compilation.GetTypeByMetadataName("System.Threading.Tasks.ValueTask`1");
+
         public static INamedTypeSymbol? IEnumerableOfTType(this Compilation compilation)
             => compilation.GetTypeByMetadataName(typeof(IEnumerable<>).FullName!);