Merge pull request #41377 from allisonchou/ExtractLocalFunctionBug4055

Fix extract local function crash

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

@@ -4282,6 +4282,158 @@ static void NewMethod()
             await TestInRegularAndScriptAsync(input, expected, CodeActionIndex);
         }
 
+        [WorkItem(40555, "https://github.com/dotnet/roslyn/issues/40555")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractLocalFunction)]
+        public async Task TestOnLocalFunctionHeader_Parameter()
+        {
+            var input = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M(() =>
+        {
+            void F(int [|x|])
+            {
+            }
+        });
+    }
+}";
+            var expected = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M({|Rename:NewMethod|}());
+
+        static Action NewMethod()
+        {
+            return () =>
+            {
+                void F(int x)
+                {
+                }
+            };
+        }
+    }
+}";
+            await TestInRegularAndScriptAsync(input, expected, CodeActionIndex);
+        }
+
+        [WorkItem(40555, "https://github.com/dotnet/roslyn/issues/40555")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractLocalFunction)]
+        public async Task TestOnLocalFunctionHeader_Parameter_ExpressionBody()
+        {
+            var input = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M(() =>
+        {
+            int F(int [|x|]) => 1;
+        });
+    }
+}";
+            var expected = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M({|Rename:NewMethod|}());
+
+        static Action NewMethod()
+        {
+            return () =>
+            {
+                int F(int x) => 1;
+            };
+        }
+    }
+}";
+            await TestInRegularAndScriptAsync(input, expected, CodeActionIndex);
+        }
+
+        [WorkItem(40555, "https://github.com/dotnet/roslyn/issues/40555")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractLocalFunction)]
+        public async Task TestOnLocalFunctionHeader_Identifier()
+        {
+            var input = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M(() =>
+        {
+            void [|F|](int x)
+            {
+            }
+        });
+    }
+}";
+            var expected = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M({|Rename:NewMethod|}());
+
+        static Action NewMethod()
+        {
+            return () =>
+            {
+                void F(int x)
+                {
+                }
+            };
+        }
+    }
+}";
+            await TestInRegularAndScriptAsync(input, expected, CodeActionIndex);
+        }
+
+        [WorkItem(40555, "https://github.com/dotnet/roslyn/issues/40555")]
+        [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractLocalFunction)]
+        public async Task TestOnLocalFunctionHeader_Identifier_ExpressionBody()
+        {
+            var input = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M(() =>
+        {
+            int [|F|](int x) => 1;
+        });
+    }
+}";
+            var expected = @"
+using System;
+class C
+{
+    void M(Action a)
+    {
+        M({|Rename:NewMethod|}());
+
+        static Action NewMethod()
+        {
+            return () =>
+            {
+                int F(int x) => 1;
+            };
+        }
+    }
+}";
+            await TestInRegularAndScriptAsync(input, expected, CodeActionIndex);
+        }
+
         [WorkItem(40654, "https://github.com/dotnet/roslyn/issues/40654")]
         [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsExtractLocalFunction)]
         public async Task TestMissingOnUsingStatement()

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

@@ -28,21 +28,23 @@ public CSharpMethodExtractor(CSharpSelectionResult result, bool localFunction)
         protected override Task<AnalyzerResult> AnalyzeAsync(SelectionResult selectionResult, bool localFunction, CancellationToken cancellationToken)
             => CSharpAnalyzer.AnalyzeAsync(selectionResult, localFunction, cancellationToken);
 
-        protected override async Task<InsertionPoint> GetInsertionPointAsync(SemanticDocument document, int position, CancellationToken cancellationToken)
+        protected override async Task<InsertionPoint> GetInsertionPointAsync(SemanticDocument document, CancellationToken cancellationToken)
         {
-            Contract.ThrowIfFalse(position >= 0);
+            var originalSpanStart = OriginalSelectionResult.OriginalSpan.Start;
+            Contract.ThrowIfFalse(originalSpanStart >= 0);
 
             var root = await document.Document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
-            var basePosition = root.FindToken(position);
+            var basePosition = root.FindToken(originalSpanStart);
 
             if (LocalFunction)
             {
                 // If we are extracting a local function and are within a local function, then we want the new function to be created within the
                 // existing local function instead of the overarching method.
-                var localMethodNode = basePosition.GetAncestor<LocalFunctionStatementSyntax>(node => node.SpanStart != basePosition.SpanStart);
-                if (localMethodNode is object)
+                var localFunctionNode = basePosition.GetAncestor<LocalFunctionStatementSyntax>(node => (node.Body != null && node.Body.Span.Contains(OriginalSelectionResult.OriginalSpan)) ||
+                                                                                                       (node.ExpressionBody != null && node.ExpressionBody.Span.Contains(OriginalSelectionResult.OriginalSpan)));
+                if (localFunctionNode is object)
                 {
-                    return await InsertionPoint.CreateAsync(document, localMethodNode, cancellationToken).ConfigureAwait(false);
+                    return await InsertionPoint.CreateAsync(document, localFunctionNode, cancellationToken).ConfigureAwait(false);
                 }
             }
 

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

@@ -27,7 +27,7 @@ public MethodExtractor(SelectionResult selectionResult, bool localFunction)
         }
 
         protected abstract Task<AnalyzerResult> AnalyzeAsync(SelectionResult selectionResult, bool localFunction, CancellationToken cancellationToken);
-        protected abstract Task<InsertionPoint> GetInsertionPointAsync(SemanticDocument document, int position, CancellationToken cancellationToken);
+        protected abstract Task<InsertionPoint> GetInsertionPointAsync(SemanticDocument document, CancellationToken cancellationToken);
         protected abstract Task<TriviaResult> PreserveTriviaAsync(SelectionResult selectionResult, CancellationToken cancellationToken);
         protected abstract Task<SemanticDocument> ExpandAsync(SelectionResult selection, CancellationToken cancellationToken);
 
@@ -53,7 +53,7 @@ public async Task<ExtractMethodResult> ExtractMethodAsync(CancellationToken canc
                 return new FailedExtractMethodResult(operationStatus);
             }
 
-            var insertionPoint = await GetInsertionPointAsync(analyzeResult.SemanticDocument, OriginalSelectionResult.OriginalSpan.Start, cancellationToken).ConfigureAwait(false);
+            var insertionPoint = await GetInsertionPointAsync(analyzeResult.SemanticDocument, cancellationToken).ConfigureAwait(false);
             cancellationToken.ThrowIfCancellationRequested();
 
             var triviaResult = await PreserveTriviaAsync(OriginalSelectionResult.With(insertionPoint.SemanticDocument), cancellationToken).ConfigureAwait(false);

src/Features/VisualBasic/Portable/ExtractMethod/VisualBasicMethodExtractor.vb

@@ -24,11 +24,12 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.ExtractMethod
             Return VisualBasicAnalyzer.AnalyzeResultAsync(selectionResult, cancellationToken)
         End Function
 
-        Protected Overrides Async Function GetInsertionPointAsync(document As SemanticDocument, position As Integer, cancellationToken As CancellationToken) As Task(Of InsertionPoint)
-            Contract.ThrowIfFalse(position >= 0)
+        Protected Overrides Async Function GetInsertionPointAsync(document As SemanticDocument, cancellationToken As CancellationToken) As Task(Of InsertionPoint)
+            Dim originalSpanStart = OriginalSelectionResult.OriginalSpan.Start
+            Contract.ThrowIfFalse(originalSpanStart >= 0)
 
             Dim root = Await document.Document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(False)
-            Dim basePosition = root.FindToken(position)
+            Dim basePosition = root.FindToken(originalSpanStart)
 
             Dim enclosingTopLevelNode As SyntaxNode = basePosition.GetAncestor(Of PropertyBlockSyntax)()
             If enclosingTopLevelNode Is Nothing Then