Add nullability information to extract method

src/EditorFeatures/CSharpTest/CodeActions/ExtractMethod/ExtractMethodTests.cs

@@ -2661,5 +2661,53 @@ class C
         return x;
     }
 }");
+
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractMethod)]
+        public Task TestFlowNullable_LambdaWithReturn()
+            => TestInRegularAndScriptAsync(
+@"#nullable enable
+
+using System;
+
+class C
+{
+    public string? M()
+    {
+        [|string? x = null;
+        Func<string?> returnNull = () =>
+        {
+            return null;
+        };
+
+        x = returnNull() ?? string.Empty;|]
+
+        return x;
+    }
+}",
+@"#nullable enable
+
+using System;
+
+class C
+{
+    public string? M()
+    {
+        string? x = {|Rename:NewMethod|}();
+
+        return x;
+    }
+
+    private static string NewMethod()
+    {
+        string? x = null;
+        Func<string?> returnNull = () =>
+        {
+            return null;
+        };
+
+        x = returnNull() ?? string.Empty;
+        return x;
+    }
+}");
     }
 }

src/Features/CSharp/Portable/ExtractMethod/CSharpMethodExtractor.Analyzer.cs

@@ -147,37 +147,42 @@ protected override ITypeSymbol GetSymbolType(SemanticModel semanticModel, ISymbo
                         var references = selectionOperation.DescendantsAndSelf()
                             .Where(IsSymbolReferencedByOperation);
 
-                        var flowState = GetFlowStateFromLocations(references);
-                        return base.GetSymbolType(semanticModel, symbol).WithNullability(flowState);
+                        if (AreAllReferencesNotNull(references))
+                        {
+                            return base.GetSymbolType(semanticModel, symbol).WithNullability(NullableAnnotation.NotAnnotated);
+                        }
+
+                        return base.GetSymbolType(semanticModel, symbol);
 
                     default:
                         return base.GetSymbolType(semanticModel, symbol);
                 }
 
-                NullableFlowState GetFlowStateFromLocations(IEnumerable<IOperation> references)
+                bool AreAllReferencesNotNull(IEnumerable<IOperation> references)
                 {
                     foreach (var reference in references)
                     {
                         var typeInfo = semanticModel.GetTypeInfo(reference.Syntax);
 
                         switch (typeInfo.Nullability.FlowState)
                         {
-                            case NullableFlowState.MaybeNull: return NullableFlowState.MaybeNull;
-                            case NullableFlowState.None: return NullableFlowState.None;
+                            case NullableFlowState.MaybeNull:
+                            case NullableFlowState.None:
+                                return false;
                         }
                     }
 
-                    return NullableFlowState.NotNull;
+                    return true;
                 }
 
                 bool IsSymbolReferencedByOperation(IOperation operation)
-                => operation switch
-                {
-                    ILocalReferenceOperation localReference => localReference.Local == symbol,
-                    IParameterReferenceOperation parameterReference => parameterReference.Parameter == symbol,
-                    IAssignmentOperation assignment => IsSymbolReferencedByOperation(assignment.Target),
-                    _ => false
-                };
+                    => operation switch
+                    {
+                        ILocalReferenceOperation localReference => localReference.Local == symbol,
+                        IParameterReferenceOperation parameterReference => parameterReference.Parameter == symbol,
+                        IAssignmentOperation assignment => IsSymbolReferencedByOperation(assignment.Target),
+                        _ => false
+                    };
             }
         }
     }

src/Features/CSharp/Portable/ExtractMethod/CSharpMethodExtractor.CSharpCodeGenerator.cs

@@ -677,6 +677,11 @@ protected override async Task<SemanticDocument> UpdateMethodAfterGenerationAsync
 
                 foreach (var returnOperation in returnOperations)
                 {
+                    if (!ReturnOperationBelongsToMethod(returnOperation.Syntax, methodOperation.Syntax))
+                    {
+                        continue;
+                    }
+
                     var syntax = returnOperation.ReturnedValue?.Syntax ?? returnOperation.Syntax;
                     var returnTypeInfo = semanticModel.GetTypeInfo(syntax, cancellationToken);
                     if (returnTypeInfo.Nullability.FlowState == NullableFlowState.MaybeNull)
@@ -694,6 +699,19 @@ protected override async Task<SemanticDocument> UpdateMethodAfterGenerationAsync
 
                 var newDocument = originalDocument.Document.WithSyntaxRoot(newRoot);
                 return await SemanticDocument.CreateAsync(newDocument, cancellationToken).ConfigureAwait(false);
+
+                static bool ReturnOperationBelongsToMethod(SyntaxNode returnOperationSyntax, SyntaxNode methodSyntax)
+                {
+                    var enclosingMethod = returnOperationSyntax.FirstAncestorOrSelf<SyntaxNode>(n => n switch
+                    {
+                        BaseMethodDeclarationSyntax _ => true,
+                        AnonymousFunctionExpressionSyntax _ => true,
+                        LocalFunctionStatementSyntax _ => true,
+                        _ => false
+                    });
+
+                    return enclosingMethod == methodSyntax;
+                }
             }
         }
     }

src/Features/Core/Portable/ExtractMethod/MethodExtractor.Analyzer.cs

@@ -220,7 +220,7 @@ private void WrapReturnTypeInTask(SemanticModel model, ref ITypeSymbol returnTyp
                 {
                     // check whether current selection contains return statement
                     var parameters = GetMethodParameters(variableInfoMap.Values);
-                    var returnType = GetReturnTypeFromStatement(model);
+                    var returnType = SelectionResult.GetContainingScopeType() ?? compilation.GetSpecialType(SpecialType.System_Object);
 
                     var unsafeAddressTakenUsed = ContainsVariableUnsafeAddressTaken(dataFlowAnalysisData, variableInfoMap.Keys);
                     return (parameters, returnType, default(VariableInfo), unsafeAddressTakenUsed);
@@ -231,23 +231,14 @@ private void WrapReturnTypeInTask(SemanticModel model, ref ITypeSymbol returnTyp
                     var parameters = MarkVariableInfoToUseAsReturnValueIfPossible(GetMethodParameters(variableInfoMap.Values));
                     var variableToUseAsReturnValue = parameters.FirstOrDefault(v => v.UseAsReturnValue);
                     var returnType = variableToUseAsReturnValue != null
-                        ? GetReturnTypeFromReturnVariable(_semanticDocument, variableToUseAsReturnValue)
+                        ? variableToUseAsReturnValue.GetVariableType(_semanticDocument)
                         : compilation.GetSpecialType(SpecialType.System_Void);
 
                     var unsafeAddressTakenUsed = ContainsVariableUnsafeAddressTaken(dataFlowAnalysisData, variableInfoMap.Keys);
                     return (parameters, returnType, variableToUseAsReturnValue, unsafeAddressTakenUsed);
                 }
             }
 
-            protected ITypeSymbol GetReturnTypeFromStatement(SemanticModel semanticModel)
-            {
-                var compilation = semanticModel.Compilation;
-                return SelectionResult.GetContainingScopeType() ?? compilation.GetSpecialType(SpecialType.System_Object);
-            }
-
-            protected ITypeSymbol GetReturnTypeFromReturnVariable(SemanticDocument semanticDocument, VariableInfo variableToUseAsReturnValue)
-                => variableToUseAsReturnValue.GetVariableType(semanticDocument);
-
             private bool IsInExpressionOrHasReturnStatement(SemanticModel model)
             {
                 var isInExpressionOrHasReturnStatement = SelectionResult.SelectionInExpression;