Fix indentation issue if trailing trivia contains a newline at the end of an indentation scope

We were always adding indentation to the trailing trivia of a token. This could cause issues if the token at the end of an indentation scope (in most cases a `}` token) contained a trailing newline.

In that case we would add indentation after that final newline based on the indentation scope of `}`, which effectively indented the next token to the indentation scope inside the `}` and thus indenting it one level to far.

To fix this, don’t add indentation for trailing newlines of tokens. Instead, add the indentation as the leading trivia of the next token, which will pick up the correct indentation scope.

Author: ahoppen

Sources/SwiftBasicFormat/BasicFormat.swift

@@ -594,8 +594,19 @@ open class BasicFormat: SyntaxRewriter {
       trailingTriviaIndentation = anchorPointIndentation
     }
 
-    leadingTrivia = leadingTrivia.indented(indentation: leadingTriviaIndentation, isOnNewline: previousTokenIsStringLiteralEndingInNewline)
-    trailingTrivia = trailingTrivia.indented(indentation: trailingTriviaIndentation, isOnNewline: false)
+    leadingTrivia = leadingTrivia.indented(
+      indentation: leadingTriviaIndentation,
+      isOnNewline: previousTokenIsStringLiteralEndingInNewline || previousTokenWillEndWithNewline
+    )
+    trailingTrivia = trailingTrivia.indented(
+      indentation: trailingTriviaIndentation,
+      isOnNewline: false,
+      // Don't add indentation to the last newline.
+      // Its indentation will be added to the next token's leading trivia, which
+      // might have a different indentation scope than this token (in particular
+      // if this token is a closing brace).
+      addIndentationAfterLastNewline: false
+    )
 
     leadingTrivia = leadingTrivia.trimmingTrailingWhitespaceBeforeNewline(isBeforeNewline: leadingTriviaIsFollowedByNewline)
     trailingTrivia = trailingTrivia.trimmingTrailingWhitespaceBeforeNewline(isBeforeNewline: nextTokenWillStartWithNewline)

Sources/SwiftBasicFormat/Trivia+FormatExtensions.swift

@@ -62,7 +62,20 @@ extension Trivia {
   }
 
   /// Adds `indentation` after every newline in this trivia.
-  func indented(indentation: Trivia, isOnNewline: Bool) -> Trivia {
+  ///
+  /// - Parameters:
+  ///   - indentation: The amount of indentation to add
+  ///   - isOnNewline: Whether this token starts on a new line.
+  ///     This causes the indentation to get added at the start of the trivia.
+  ///   - addIndentationAfterLastNewline: Whether to add indentation after newline
+  ///     if the newline is the last piece of trivia. Not doing this makes sense
+  ///     if the indentation will get added to the next token's leading trivia
+  ///     via `isOnNewline`.
+  func indented(
+    indentation: Trivia,
+    isOnNewline: Bool,
+    addIndentationAfterLastNewline: Bool = true
+  ) -> Trivia {
     guard !isEmpty else {
       if isOnNewline {
         return indentation
@@ -72,13 +85,13 @@ extension Trivia {
 
     var indentedPieces: [TriviaPiece] = []
     if isOnNewline {
-      indentedPieces.append(contentsOf: indentation)
+      indentedPieces += indentation
     }
 
-    for piece in pieces {
+    for (index, piece) in pieces.enumerated() {
       indentedPieces.append(piece)
-      if piece.isNewline {
-        indentedPieces.append(contentsOf: indentation)
+      if piece.isNewline && !(index == pieces.count - 1 && !addIndentationAfterLastNewline) {
+        indentedPieces += indentation
       }
     }
 

Tests/SwiftBasicFormatTest/BasicFormatTests.swift

@@ -566,4 +566,48 @@ final class BasicFormatTest: XCTestCase {
       """
     assertFormatted(source: source, expected: source)
   }
+
+  func testNewlineInTrailingTriviaAtEndOfIndentationScope() throws {
+    assertFormatted(
+      tree: try FunctionDeclSyntax("func test()") {
+        CodeBlockItemSyntax("Task {\n}\n")
+      },
+      expected: """
+        func test() {
+            Task {
+            }
+        }
+        """
+    )
+
+    assertFormatted(
+      tree: try FunctionDeclSyntax("func test()") {
+        CodeBlockItemSyntax("Task {\n}\n\n")
+      },
+      expected: """
+        func test() {
+            Task {
+            }
+
+        }
+        """
+    )
+
+    assertFormatted(
+      tree: try FunctionDeclSyntax("func bar()") {
+        try FunctionDeclSyntax("func test()") {
+          CodeBlockItemSyntax("Task {\n}\n")
+        }
+      },
+      expected: """
+        func bar() {
+            func test() {
+                Task {
+                }
+            }
+        }
+        """
+    )
+  }
+
 }