Fix directive placement when removing nodes using syntax API. (#76275)

src/Compilers/CSharp/Portable/Syntax/SyntaxNodeRemover.cs

@@ -45,7 +45,7 @@ internal static class SyntaxNodeRemover
             return (TRoot?)result;
         }
 
-        private class SyntaxRemover : CSharpSyntaxRewriter
+        private sealed class SyntaxRemover : CSharpSyntaxRewriter
         {
             private readonly HashSet<SyntaxNode> _nodesToRemove;
             private readonly SyntaxRemoveOptions _options;
@@ -432,8 +432,8 @@ private void AddDirectives(SyntaxNode node, TextSpan span)
                     }
 
                     var directivesInSpan = node.DescendantTrivia(span, n => n.ContainsDirectives, descendIntoTrivia: true)
-                                         .Where(tr => tr.IsDirective)
-                                         .Select(tr => (DirectiveTriviaSyntax)tr.GetStructure()!);
+                                               .Where(tr => tr.IsDirective)
+                                               .Select(tr => (DirectiveTriviaSyntax)tr.GetStructure()!);
 
                     foreach (var directive in directivesInSpan)
                     {
@@ -465,10 +465,32 @@ private void AddDirectives(SyntaxNode node, TextSpan span)
 
                         if (_directivesToKeep.Contains(directive))
                         {
-                            AddResidualTrivia(SyntaxFactory.TriviaList(directive.ParentTrivia), requiresNewLine: true);
+                            var parentTrivia = directive.ParentTrivia;
+                            var (triviaList, directiveTriviaListIndex) = getTriviaListAndIndex(parentTrivia);
+
+                            // If we're keeping a directive, and it's not at the start of the line, keep the indentation
+                            // whitespace that precedes it as well so that the directive stays in the right location.
+                            if (directiveTriviaListIndex >= 1 && triviaList[directiveTriviaListIndex - 1] is { RawKind: (int)SyntaxKind.WhitespaceTrivia } indentationTrivia)
+                            {
+                                AddResidualTrivia(SyntaxFactory.TriviaList(indentationTrivia, parentTrivia), requiresNewLine: true);
+                            }
+                            else
+                            {
+                                AddResidualTrivia(SyntaxFactory.TriviaList(parentTrivia), requiresNewLine: true);
+                            }
                         }
                     }
                 }
+
+                static (SyntaxTriviaList triviaList, int index) getTriviaListAndIndex(SyntaxTrivia trivia)
+                {
+                    var parentToken = trivia.Token;
+
+                    var index = parentToken.LeadingTrivia.IndexOf(trivia);
+                    return index >= 0
+                        ? (parentToken.LeadingTrivia, index)
+                        : (parentToken.TrailingTrivia, parentToken.TrailingTrivia.IndexOf(trivia));
+                }
             }
 
             private static bool HasRelatedDirectives(DirectiveTriviaSyntax directive)

src/Compilers/CSharp/Test/Syntax/Syntax/SyntaxNodeTests.cs

@@ -2802,6 +2802,44 @@ class C
             });
         }
 
+        [Fact, WorkItem("https://github.com/dotnet/roslyn/issues/19613")]
+        public void TestRemove_KeepUnbalancedDirectives_Indented()
+        {
+            var inputText = """
+                class C
+                {
+                    // before
+                    #region Fred
+                    // more before
+                    void M()
+                    {
+                    } // after
+                    #endregion
+                }
+                """;
+
+            var expectedText = """
+                class C
+                {
+
+                    #region Fred
+                    #endregion
+                }
+                """;
+
+            TestWithWindowsAndUnixEndOfLines(inputText, expectedText, (cu, expected) =>
+            {
+                var m = cu.DescendantNodes().OfType<MethodDeclarationSyntax>().FirstOrDefault();
+                Assert.NotNull(m);
+
+                var cu2 = cu.RemoveNode(m, SyntaxRemoveOptions.KeepUnbalancedDirectives);
+
+                var text = cu2.ToFullString();
+
+                Assert.Equal(expected, text);
+            });
+        }
+
         [Fact]
         public void TestRemove_KeepDirectives()
         {

src/Compilers/VisualBasic/Portable/Syntax/SyntaxNodeRemover.vb

@@ -27,7 +27,7 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.Syntax
             Return DirectCast(result, TRoot)
         End Function
 
-        Private Class SyntaxRemover
+        Private NotInheritable Class SyntaxRemover
             Inherits VisualBasicSyntaxRewriter
 
             Private ReadOnly _nodesToRemove As HashSet(Of SyntaxNode)
@@ -299,8 +299,8 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.Syntax
                     End If
 
                     Dim directivesInSpan = node.DescendantTrivia(span, Function(n) n.ContainsDirectives, descendIntoTrivia:=True) _
-                                                .Where(Function(tr) tr.IsDirective) _
-                                                .Select(Function(tr) DirectCast(tr.GetStructure(), DirectiveTriviaSyntax))
+                                               .Where(Function(tr) tr.IsDirective) _
+                                               .Select(Function(tr) DirectCast(tr.GetStructure(), DirectiveTriviaSyntax))
 
                     For Each directive In directivesInSpan
                         If (Me._options And SyntaxRemoveOptions.KeepDirectives) <> 0 Then
@@ -319,12 +319,33 @@ Namespace Microsoft.CodeAnalysis.VisualBasic.Syntax
                         End If
 
                         If Me._directivesToKeep.Contains(directive) Then
-                            AddResidualTrivia(SyntaxFactory.TriviaList(directive.ParentTrivia), requiresNewLine:=True)
+                            Dim parentTrivia = directive.ParentTrivia
+                            Dim triviaListAndIndex = GetTriviaListAndIndex(parentTrivia)
+                            Dim triviaList = triviaListAndIndex.triviaList
+                            Dim directiveTriviaListIndex = triviaListAndIndex.Index
+
+                            ' If we're keeping a directive, and it's not at the start of the line, keep the whitespace
+                            ' that precedes it as well.
+                            If directiveTriviaListIndex >= 1 AndAlso triviaList(directiveTriviaListIndex - 1).Kind() = SyntaxKind.WhitespaceTrivia Then
+                                AddResidualTrivia(SyntaxFactory.TriviaList(
+                                    triviaList(directiveTriviaListIndex - 1), parentTrivia), requiresNewLine:=True)
+                            Else
+                                AddResidualTrivia(SyntaxFactory.TriviaList(parentTrivia), requiresNewLine:=True)
+                            End If
                         End If
                     Next
                 End If
             End Sub
 
+            Private Shared Function GetTriviaListAndIndex(trivia As SyntaxTrivia) As (triviaList As SyntaxTriviaList, Index As Integer)
+                Dim parentToken = trivia.Token
+
+                Dim index = parentToken.LeadingTrivia.IndexOf(trivia)
+                Return If(index >= 0,
+                    (parentToken.LeadingTrivia, index),
+                    (parentToken.TrailingTrivia, parentToken.TrailingTrivia.IndexOf(trivia)))
+            End Function
+
             Private Shared Function HasRelatedDirectives(directive As DirectiveTriviaSyntax) As Boolean
                 Select Case directive.Kind
                     Case SyntaxKind.IfDirectiveTrivia,

src/Compilers/VisualBasic/Test/Syntax/TestSyntaxNodes.vb

@@ -2726,7 +2726,35 @@ End Class
             Dim result = cu2.ToFullString()
 
             Assert.Equal(expected, result)
+        End Sub
+
+        <Fact, WorkItem("https://github.com/dotnet/roslyn/issues/19613")>
+        Public Sub TestRemove_KeepUnbalancedDirectives_Indented()
+            Dim text = <![CDATA[
+Class C
+    #Region "A Region"
+    Sub Goo()
+    End Sub
+    #End Region
+End Class
+]]>.Value.Replace(vbLf, vbCrLf)
+
+            Dim expected = <![CDATA[
+Class C
 
+    #Region "A Region"
+    #End Region
+End Class
+]]>.Value.Replace(vbLf, vbCrLf)
+
+            Dim cu = SyntaxFactory.ParseCompilationUnit(text)
+            Dim n = cu.DescendantTokens().Where(Function(t) t.ToString() = "Goo").Select(Function(t) t.Parent.FirstAncestorOrSelf(Of MethodBlockSyntax)()).FirstOrDefault()
+
+            Dim cu2 = cu.RemoveNode(n, SyntaxRemoveOptions.KeepUnbalancedDirectives)
+
+            Dim result = cu2.ToFullString()
+
+            Assert.Equal(expected, result)
         End Sub
 
         <Fact, WorkItem(530316, "http://vstfdevdiv:8080/DevDiv2/DevDiv/_workitems/edit/530316")>

src/EditorFeatures/CSharpTest/CodeActions/MoveType/MoveTypeTests.MoveToNewFile.cs

@@ -1179,22 +1179,22 @@ public class [||]Inner
             #endif
             """;
         var codeAfterMove =
-        """
-        using System;
+            """
+            using System;
 
-        namespace N
-        {
-            class Program
+            namespace N
             {
-                static void Main()
+                class Program
                 {
+                    static void Main()
+                    {
+                    }
                 }
             }
-        }
 
-        #if true
-        #endif
-        """;
+            #if true
+            #endif
+            """;
 
         var expectedDocumentName = "Inner.cs";
         var destinationDocumentText =
@@ -1237,22 +1237,22 @@ public class [||]Inner
             }
             """;
         var codeAfterMove =
-        """
-        using System;
+            """
+            using System;
 
-        namespace N
-        {
-            partial class Program
+            namespace N
             {
-                static void Main()
+                partial class Program
                 {
-                }
+                    static void Main()
+                    {
+                    }
 
-        #if true
-        #endif
+            #if true
+            #endif
+                }
             }
-        }
-        """;
+            """;
 
         var expectedDocumentName = "Inner.cs";
         var destinationDocumentText =
@@ -1275,6 +1275,99 @@ await TestMoveTypeToNewFileAsync(
             code, codeAfterMove, expectedDocumentName, destinationDocumentText);
     }
 
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/19613")]
+    public async Task MoveTypeWithDirectives3()
+    {
+        var code =
+            """
+            public class Goo
+            {
+                #region Region
+                public class [||]Bar
+                {
+                }
+                #endregion
+            }
+            """;
+        var codeAfterMove =
+            """
+            public partial class Goo
+            {
+
+                #region Region
+                #endregion
+            }
+            """;
+
+        var expectedDocumentName = "Bar.cs";
+        var destinationDocumentText =
+            """
+            public partial class Goo
+            {
+                #region Region
+                public class Bar
+                {
+                }
+                #endregion
+            }
+            """;
+
+        await TestMoveTypeToNewFileAsync(
+            code, codeAfterMove, expectedDocumentName, destinationDocumentText);
+    }
+
+    [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/19613")]
+    public async Task MoveTypeWithDirectives4()
+    {
+        var code =
+            """
+            public class Goo
+            {
+                #region Region1
+                public class NotBar
+                {
+                }
+                #endregion
+
+                #region Region2
+                public class [||]Bar
+                {
+                }
+                #endregion
+            }
+            """;
+        var codeAfterMove =
+            """
+            public partial class Goo
+            {
+                #region Region1
+                public class NotBar
+                {
+                }
+
+                #endregion
+                #region Region2
+                #endregion
+            }
+            """;
+
+        var expectedDocumentName = "Bar.cs";
+        var destinationDocumentText =
+            """
+            public partial class Goo
+            {
+                #region Region2
+                public class Bar
+                {
+                }
+                #endregion
+            }
+            """;
+
+        await TestMoveTypeToNewFileAsync(
+            code, codeAfterMove, expectedDocumentName, destinationDocumentText);
+    }
+
     [WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/21456")]
     public async Task TestLeadingBlankLines1()
     {

src/EditorFeatures/DiagnosticsTestUtilities/MoveType/AbstractMoveTypeTest.cs

@@ -172,7 +172,7 @@ private protected async Task TestMoveTypeToNewFileAsync(
 
                 var modifiedSourceDocument = newSolution.GetDocument(sourceDocumentId);
                 var actualSourceTextAfterRefactoring = (await modifiedSourceDocument.GetTextAsync()).ToString();
-                Assert.Equal(expectedSourceTextAfterRefactoring, actualSourceTextAfterRefactoring);
+                AssertEx.Equal(expectedSourceTextAfterRefactoring, actualSourceTextAfterRefactoring);
             }
             else
             {