Repository: apple/swift-syntax (swift)

Issue:
MacroSystem should add newline at the end of added members for an empty type
### Description

Given a macro `AddMembers` that does as its name suggests, MacroSystem will currently expand the following:
```
@AddMembers struct Foo {}
```

To:
```
struct Foo {
  func addedMember() {
    ...
  }}
```

Note the lack of a newline after the added member. This usually makes sense as there would already be one, but when there isn't, we should add it.

rdar://116773975

Pull request:
Fixes #2307

Patch:
diff --git a/Sources/SwiftParser/ParseSourceFile.swift b/Sources/SwiftParser/ParseSourceFile.swift
index 99015ac69a4..98b227966b1 100644
--- a/Sources/SwiftParser/ParseSourceFile.swift
+++ b/Sources/SwiftParser/ParseSourceFile.swift
@@ -152,8 +152,9 @@ extension Parser {
     }
 
     let remainingTokens = self.consumeRemainingTokens()
+
     if remainingTokens.isEmpty {
-      return into
+      return R.init(transferTrailingTrivaFromEndOfFileIfPresent(raw: into.raw))!
     }
 
     let existingUnexpected: [RawSyntax]
@@ -166,6 +167,33 @@ extension Parser {
     let unexpected = RawUnexpectedNodesSyntax(elements: existingUnexpected + remainingTokens, arena: self.arena)
 
     let withUnexpected = layout.replacingChild(at: layout.children.count - 1, with: unexpected.raw, arena: self.arena)
-    return R.init(withUnexpected)!
+
+    return R.init(transferTrailingTrivaFromEndOfFileIfPresent(raw: withUnexpected))!
+  }
+
+  /// Parses the end-of-file token and appends its leading trivia to the provided `RawSyntax`.
+  /// - Parameter raw: The raw syntax node to which the leading trivia of the end-of-file token will be appended.
+  /// - Returns: A new `RawSyntax` instance with trailing trivia transferred from the end-of-file token if present, otherwise it will return the raw parameter..
+  private mutating func transferTrailingTrivaFromEndOfFileIfPresent(raw: RawSyntax) -> RawSyntax {
+    guard let endOfFileToken = self.consume(if: .endOfFile),
+      !endOfFileToken.leadingTriviaPieces.isEmpty,
+      let raw = raw.withTrailingTrivia(
+        Trivia(
+          rawPieces: (raw.trailingTriviaPieces ?? []) + endOfFileToken.leadingTriviaPieces
+        ),
+        arena: self.arena
+      )
+    else {
+      return raw
+    }
+
+    return raw
+  }
+}
+
+private extension Trivia {
+  init(rawPieces: [RawTriviaPiece]) {
+    let pieces = rawPieces.map(TriviaPiece.init(raw:))
+    self.init(pieces: pieces)
   }
 }
diff --git a/Sources/SwiftSyntax/Raw/RawSyntax.swift b/Sources/SwiftSyntax/Raw/RawSyntax.swift
index 583de2c3298..d083bf476d5 100644
--- a/Sources/SwiftSyntax/Raw/RawSyntax.swift
+++ b/Sources/SwiftSyntax/Raw/RawSyntax.swift
@@ -302,7 +302,8 @@ extension RawSyntax {
   /// - Parameters:
   ///   - leadingTrivia: The trivia to attach.
   ///   - arena: SyntaxArena to the result node data resides.
-  func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
+  @_spi(RawSyntax)
+  public func withLeadingTrivia(_ leadingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
     switch view {
     case .token(let tokenView):
       return .makeMaterializedToken(
@@ -328,7 +329,8 @@ extension RawSyntax {
   /// - Parameters:
   ///   - trailingTrivia: The trivia to attach.
   ///   - arena: SyntaxArena to the result node data resides.
-  func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
+  @_spi(RawSyntax)
+  public func withTrailingTrivia(_ trailingTrivia: Trivia, arena: SyntaxArena) -> RawSyntax? {
     switch view {
     case .token(let tokenView):
       return .makeMaterializedToken(
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
index 7d5777e0f5e..0e3265fd14d 100644
--- a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
+++ b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
@@ -762,10 +762,8 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
     return CodeBlockItemListSyntax(newItems)
   }
 
-  override func visit(_ node: MemberBlockItemListSyntax) -> MemberBlockItemListSyntax {
+  override func visit(_ node: MemberBlockSyntax) -> MemberBlockSyntax {
     let parentDeclGroup = node
-      .parent?
-      .as(MemberBlockSyntax.self)?
       .parent?
       .as(DeclSyntax.self)
     var newItems: [MemberBlockItemSyntax] = []
@@ -792,7 +790,7 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
       extensions += expandExtensions(of: node.decl)
     }
 
-    for var item in node {
+    for var item in node.members {
       // Expand member attribute members attached to the declaration context.
       // Note that MemberAttribute macros are _not_ applied to generated members
       if let parentDeclGroup, let decl = item.decl.asProtocol(WithAttributesSyntax.self) {
@@ -825,7 +823,29 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
       }
     }
 
-    return .init(newItems)
+    /// Returns an leading trivia for the member blocks closing brace.
+    /// It will add a leading newline, if there is none.
+    var leadingTriviaForClosingBrace: Trivia {
+      if newItems.isEmpty {
+        return node.rightBrace.leadingTrivia
+      }
+
+      if node.rightBrace.leadingTrivia.contains(where: { $0.isNewline }) {
+        return node.rightBrace.leadingTrivia
+      }
+
+      if newItems.last?.trailingTrivia.pieces.last?.isNewline ?? false {
+        return node.rightBrace.leadingTrivia
+      } else {
+        return .newline + node.rightBrace.leadingTrivia
+      }
+    }
+
+    return MemberBlockSyntax(
+      leftBrace: node.leftBrace,
+      members: MemberBlockItemListSyntax(newItems),
+      rightBrace: node.rightBrace.with(\.leadingTrivia, leadingTriviaForClosingBrace)
+    )
   }
 
   override func visit(_ node: VariableDeclSyntax) -> DeclSyntax {


Test patch:
diff --git a/Examples/Tests/MacroExamples/Implementation/Peer/AddAsyncMacroTests.swift b/Examples/Tests/MacroExamples/Implementation/Peer/AddAsyncMacroTests.swift
index 89227d3bda8..79178633770 100644
--- a/Examples/Tests/MacroExamples/Implementation/Peer/AddAsyncMacroTests.swift
+++ b/Examples/Tests/MacroExamples/Implementation/Peer/AddAsyncMacroTests.swift
@@ -43,6 +43,7 @@ final class AddAsyncMacroTests: XCTestCase {
               }
             }
           }
+
         }
         """#,
       macros: macros,
@@ -70,6 +71,7 @@ final class AddAsyncMacroTests: XCTestCase {
               continuation.resume(returning: returnValue)
             }
           }
+
         }
         """,
       macros: macros,
diff --git a/Examples/Tests/MacroExamples/Implementation/Peer/AddCompletionHandlerMacroTests.swift b/Examples/Tests/MacroExamples/Implementation/Peer/AddCompletionHandlerMacroTests.swift
index a6025fc1179..eab2c5521d6 100644
--- a/Examples/Tests/MacroExamples/Implementation/Peer/AddCompletionHandlerMacroTests.swift
+++ b/Examples/Tests/MacroExamples/Implementation/Peer/AddCompletionHandlerMacroTests.swift
@@ -35,6 +35,7 @@ final class AddCompletionHandlerMacroTests: XCTestCase {
           Task {
             completionHandler(await f(a: a, for: b, value))
           }
+
         }
         """,
       macros: macros,
diff --git a/Tests/SwiftBasicFormatTest/BasicFormatTests.swift b/Tests/SwiftBasicFormatTest/BasicFormatTests.swift
index 97ab4bb2ca5..15e5fa1d582 100644
--- a/Tests/SwiftBasicFormatTest/BasicFormatTests.swift
+++ b/Tests/SwiftBasicFormatTest/BasicFormatTests.swift
@@ -588,6 +588,7 @@ final class BasicFormatTest: XCTestCase {
         func test() {
             Task {
             }
+
         }
         """
     )
diff --git a/Tests/SwiftParserTest/DeclarationTests.swift b/Tests/SwiftParserTest/DeclarationTests.swift
index f06ad334a26..d7a6363c242 100644
--- a/Tests/SwiftParserTest/DeclarationTests.swift
+++ b/Tests/SwiftParserTest/DeclarationTests.swift
@@ -3141,4 +3141,17 @@ final class DeclarationTests: ParserTestCase {
       experimentalFeatures: .nonescapableTypes
     )
   }
+
+  func testDeclarationEndingWithNewline() {
+    let inputs: [UInt: String] = [
+      #line: "var x = 0\n",
+      #line: "var x = 0 garbage\n",
+      #line: "var x = 0 \n",
+    ]
+
+    for (line, input) in inputs {
+      let decl = DeclSyntax(stringLiteral: input)
+      XCTAssertEqual(decl.description, input, line: line)
+    }
+  }
 }
diff --git a/Tests/SwiftSyntaxMacroExpansionTest/MemberMacroTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/MemberMacroTests.swift
index a281cc69744..95e783fec41 100644
--- a/Tests/SwiftSyntaxMacroExpansionTest/MemberMacroTests.swift
+++ b/Tests/SwiftSyntaxMacroExpansionTest/MemberMacroTests.swift
@@ -299,4 +299,129 @@ final class MemberMacroTests: XCTestCase {
       ]
     )
   }
+
+  func testAddMemberToEmptyDeclaration() {
+    struct TestMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        return [DeclSyntax("var x = 0")]
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @Test
+      struct Foo {}
+      """,
+      expandedSource: """
+        struct Foo {
+
+          var x = 0
+        }
+        """,
+      macros: [
+        "Test": TestMacro.self
+      ],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testAddTwoMembersToEmptyDeclaration() {
+    struct TestMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        return [DeclSyntax("var x = 0"), DeclSyntax("var x = 0")]
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @Test
+      struct Foo {}
+      """,
+      expandedSource: """
+        struct Foo {
+
+          var x = 0
+
+          var x = 0
+        }
+        """,
+      macros: [
+        "Test": TestMacro.self
+      ],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testAddMemberToEmptyDeclarationWithEndingNewline() {
+    struct TestMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        return [DeclSyntax("var x = 0\n")]
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @Test
+      struct Foo {}
+      """,
+      expandedSource: """
+        struct Foo {
+
+          var x = 0
+        }
+        """,
+      macros: [
+        "Test": TestMacro.self
+      ],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testAddMemberToDeclarationWithASingleVariable() {
+    struct TestMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        conformingTo protocols: [TypeSyntax],
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        return [DeclSyntax("var x = 0\n")]
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      @Test
+      struct Foo {
+        var y = 0
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          var y = 0
+
+          var x = 0
+        }
+        """,
+      macros: [
+        "Test": TestMacro.self
+      ],
+      indentationWidth: indentationWidth
+    )
+  }
 }