Repository: swiftlang/swift-syntax (swift)

Issue:
Unexpected trivia when removing attributes in `MacroSystem`
### Description

On `509.0.0` and `main` ([`8966ec7`](https://github.com/apple/swift-syntax/commit/8966ec7157f9774dd3b03b53520e75bc570a284a)), I'm running into the following unexpected trivia handling in a peer macro:

```swift
fileprivate struct NoOpPeerMacro: PeerMacro {
  static func expansion(
    of node: AttributeSyntax,
    providingPeersOf declaration: some DeclSyntaxProtocol,
    in context: some MacroExpansionContext
  ) throws -> [DeclSyntax] {
    return []
  }
}

final class PeerMacroTests: XCTestCase {
  // …
  func testEmptyBeforeAttributeOnSameLineAsMemberVariable() {
    assertMacroExpansion(
      """
      struct Foo {
        @Test @State var x: Int
      }
      """,
      expandedSource: """
        struct Foo {
          @State var x: Int
        }
        """,
      macros: ["Test": NoOpPeerMacro.self]
    )
  }
  // …
}
```

Output:
```diff
–struct Foo {
–  @State var x: Int
+struct Foo {@State 
+  var x: Int
 }
```

_Quick reference_: [`AttributeRemover`]( https://github.com/apple/swift-syntax/blob/8966ec7157f9774dd3b03b53520e75bc570a284a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift#L387C15-L387C31).

To work around, I can tailor the whitespace in my test in a specific manner as follows:

```swift
struct Foo {
  @Test
  @State 
  var x: Int
}
```

What is the expected trivia treatment in this case? Should trivia in test cases be deliberately crafted, or should the `MacroSystem` strive to normalize all syntactically valid trivia in some predictable way?

_Aside: Here are a couple other tests and outputs that I found surprising:_

<details>
  <summary>Click to expand!</summary>

```swift
  func testAttributeOnOwnLineThenEmptyBeforeVariable() {
    assertMacroExpansion(
      """
      @State
      @Test var x: Int
      """,
      expandedSource: """
        @State
        var x: Int
        """,
      macros: [
        "Test": NoOpPeerMacro.self
      ]
    )
  }
```

Output:
```diff
–@State
–var x: Int
+@Statevar x: Int
```
---

```swift
  func testEmptyAfterAttributeOnSameLineAsMemberVariable() {
    assertMacroExpansion(
      """
      struct Foo {
        @State @Test var x: Int
      }
      """,
      expandedSource: """
        struct Foo {
          @State var x: Int
        }
        """,
      macros: ["Test": NoOpPeerMacro.self]
    )
  }
```

Output:
```diff
 struct Foo {
–  @State var x: Int
+  @State 
+  var x: Int
 }
```

---

```swift
  func testEmptyOnOwnLineThenAttributedMemberVariable() {
    assertMacroExpansion(
      """
      struct Foo {
        @Test
        @State var x: Int
      }
      """,
      expandedSource: """
        struct Foo {
          @State var x: Int
        }
        """,
      macros: ["Test": NoOpPeerMacro.self]
    )
  }
```

Output:
```diff
 struct Foo {
–  @State var x: Int
+  @State 
+  var x: Int
 }
```
</details>


### Steps to Reproduce

Add the tests included in the description to "Tests/SwiftMacroExpansionTest/PeerMacroTests.swift" and run them.

Pull request:
Initially opened for discussion of #2203, and now a fix for it.
    
Also, adds `AttributeRemoverTests` covering a variety of attribute removal cases involving surrounding trivia.

Fixes #2203
rdar://115585919

Patch:
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
index f97f02b53a1..f1699830691 100644
--- a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
+++ b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
@@ -397,36 +397,97 @@ private class AttributeRemover: SyntaxRewriter {
     var filteredAttributes: [AttributeListSyntax.Element] = []
     for case .attribute(let attribute) in node {
       if attributesToRemove.contains(attribute) {
-        var leadingTrivia = node.leadingTrivia
-        if let lastNewline = leadingTrivia.pieces.lastIndex(where: { $0.isNewline }),
+        var leadingTrivia = attribute.leadingTrivia
+
+        // Don't leave behind an empty line when the attribute being removed is on its own line,
+        // based on the following conditions:
+        //  - Leading trivia ends with a newline followed by arbitrary number of spaces or tabs
+        //  - All leading trivia pieces after the last newline are just whitespace, ensuring
+        //    there are no comments or other non-whitespace characters on the same line
+        //    preceding the attribute.
+        //  - There is no trailing trivia and the next token has leading trivia.
+        if let lastNewline = leadingTrivia.pieces.lastIndex(where: \.isNewline),
           leadingTrivia.pieces[lastNewline...].allSatisfy(\.isWhitespace),
-          node.trailingTrivia.isEmpty,
-          node.nextToken(viewMode: .sourceAccurate)?.leadingTrivia.first?.isNewline ?? false
+          attribute.trailingTrivia.isEmpty,
+          let nextToken = attribute.nextToken(viewMode: .sourceAccurate),
+          !nextToken.leadingTrivia.isEmpty
         {
-          // If the attribute is on its own line based on the following conditions,
-          // remove the newline from it so we don’t end up with an empty line
-          //  - Trailing trivia ends with a newline followed by arbitrary number of spaces or tabs
-          //  - There is no trailing trivia and the next token starts on a new line
           leadingTrivia = Trivia(pieces: leadingTrivia.pieces[..<lastNewline])
         }
+
         // Drop any spaces or tabs from the trailing trivia because there’s no
         // more attribute they need to separate.
-        let trailingTrivia = Trivia(pieces: attribute.trailingTrivia.drop(while: { $0.isSpaceOrTab }))
+        let trailingTrivia = attribute.trailingTrivia.trimmingPrefix(while: \.isSpaceOrTab)
         triviaToAttachToNextToken += leadingTrivia + trailingTrivia
+
+        // If the attribute is not separated from the previous attribute by trivia, as in
+        // `@First@Second var x: Int` (yes, that's valid Swift), removing the `@Second`
+        // attribute and dropping all its trivia would cause `@First` and `var` to join
+        // without any trivia in between, which is invalid. In such cases, the trailing trivia
+        // of the attribute is significant and must be retained.
+        if triviaToAttachToNextToken.isEmpty,
+          let previousToken = attribute.previousToken(viewMode: .sourceAccurate),
+          previousToken.trailingTrivia.isEmpty
+        {
+          triviaToAttachToNextToken = attribute.trailingTrivia
+        }
       } else {
-        filteredAttributes.append(.attribute(attribute))
+        filteredAttributes.append(.attribute(prependAndClearAccumulatedTrivia(to: attribute)))
+      }
+    }
+
+    // Ensure that any horizontal whitespace trailing the attributes list is trimmed if the next
+    // token starts a new line.
+    if let nextToken = node.nextToken(viewMode: .sourceAccurate),
+      nextToken.leadingTrivia.startsWithNewline
+    {
+      if !triviaToAttachToNextToken.isEmpty {
+        triviaToAttachToNextToken = triviaToAttachToNextToken.trimmingSuffix(while: \.isSpaceOrTab)
+      } else if let lastAttribute = filteredAttributes.last {
+        filteredAttributes[filteredAttributes.count - 1].trailingTrivia = lastAttribute.trailingTrivia.trimmingSuffix(while: \.isSpaceOrTab)
       }
     }
     return AttributeListSyntax(filteredAttributes)
   }
 
   override func visit(_ token: TokenSyntax) -> TokenSyntax {
-    if !triviaToAttachToNextToken.isEmpty {
-      defer { triviaToAttachToNextToken = Trivia() }
-      return token.with(\.leadingTrivia, triviaToAttachToNextToken + token.leadingTrivia)
-    } else {
-      return token
-    }
+    return prependAndClearAccumulatedTrivia(to: token)
+  }
+
+  /// Prepends the accumulated trivia to the given node's leading trivia.
+  ///
+  /// To preserve correct formatting after attribute removal, this function reassigns
+  /// significant trivia accumulated from removed attributes to the provided subsequent node.
+  /// Once attached, the accumulated trivia is cleared.
+  ///
+  /// - Parameter node: The syntax node receiving the accumulated trivia.
+  /// - Returns: The modified syntax node with the prepended trivia.
+  private func prependAndClearAccumulatedTrivia<T: SyntaxProtocol>(to syntaxNode: T) -> T {
+    defer { triviaToAttachToNextToken = Trivia() }
+    return syntaxNode.with(\.leadingTrivia, triviaToAttachToNextToken + syntaxNode.leadingTrivia)
+  }
+}
+
+private extension Trivia {
+  func trimmingPrefix(
+    while predicate: (TriviaPiece) -> Bool
+  ) -> Trivia {
+    Trivia(pieces: self.drop(while: predicate))
+  }
+
+  func trimmingSuffix(
+    while predicate: (TriviaPiece) -> Bool
+  ) -> Trivia {
+    Trivia(
+      pieces: self[...]
+        .reversed()
+        .drop(while: predicate)
+        .reversed()
+    )
+  }
+
+  var startsWithNewline: Bool {
+    self.first?.isNewline ?? false
   }
 }
 


Test patch:
diff --git a/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift
new file mode 100644
index 00000000000..cd16b029bda
--- /dev/null
+++ b/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift
@@ -0,0 +1,514 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+//==========================================================================//
+// IMPORTANT: The macros defined in this file are intended to test the      //
+// behavior of MacroSystem. Many of them do not serve as good examples of   //
+// how macros should be written. In particular, they often lack error       //
+// handling because it is not needed in the few test cases in which these   //
+// macros are invoked.                                                      //
+//==========================================================================//
+
+import SwiftDiagnostics
+import SwiftSyntax
+import SwiftSyntaxMacroExpansion
+import SwiftSyntaxMacros
+import SwiftSyntaxMacrosTestSupport
+import XCTest
+
+fileprivate struct NoOpPeerMacro: PeerMacro {
+  static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return []
+  }
+}
+
+final class AttributeRemoverTests: XCTestCase {
+  func testEmptyOnSameLineAsVariable() {
+    assertMacroExpansion(
+      "@Test var x: Int",
+      expandedSource: "var x: Int",
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyTwiceOnSameLineAsVariable() {
+    assertMacroExpansion(
+      "@Test @Test var x: Int",
+      expandedSource: "var x: Int",
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyOnOwnLineBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @Test
+      var x: Int
+      """,
+      expandedSource: "var x: Int",
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyTwiceOnOwnLineBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @Test @Test
+      var x: Int
+      """,
+      expandedSource: "var x: Int",
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyAndAttributeOnOwnLineBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @Test @State
+      var x: Int
+      """,
+      expandedSource: """
+        @State
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testAttributeAndEmptyOnOwnLineBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @State @Test
+      var x: Int
+      """,
+      expandedSource: """
+        @State
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testAttributeAndEmptyAndCommentOnOwnLineBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @State @Test// comment
+      var x: Int
+      """,
+      expandedSource: """
+        @State // comment
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testAttributeAndEmptyAndCommentOnOwnLineBeforeVariable2() {
+    assertMacroExpansion(
+      """
+      @State @Test // comment
+      var x: Int
+      """,
+      expandedSource: """
+        @State // comment
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testCommentsAroundEmpty() {
+    assertMacroExpansion(
+      """
+      struct S {
+        /// Some doc comment
+        @Test /* trailing */
+        var value: Int
+      }
+      """,
+      expandedSource: """
+        struct S {
+          /// Some doc comment
+          /* trailing */
+          var value: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testBlockCommentNewlineEmpty() {
+    assertMacroExpansion(
+      """
+      /* comment */
+      @Test
+      var value: Int
+      """,
+      expandedSource: """
+        /* comment */
+        var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmpyNewlineBlockComment() {
+    assertMacroExpansion(
+      """
+      @Test
+      /* comment */
+      var value: Int
+      """,
+      expandedSource: """
+        /* comment */
+        var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testAttributeNewlineBlockCommentEmpty() {
+    assertMacroExpansion(
+      """
+      @State
+      /*doc comment*/@Test
+      var x: Int
+      """,
+      expandedSource: """
+        @State
+        /*doc comment*/
+        var x: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyBlockCommentEmpty() {
+    assertMacroExpansion(
+      """
+      @Test /* comment */ @Test var value: Int
+      """,
+      expandedSource: """
+        /* comment */ var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyBlockCommentEmptyNewline() {
+    assertMacroExpansion(
+      """
+      @Test /* comment */ @Test
+      var value: Int
+      """,
+      expandedSource: """
+        /* comment */
+        var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyBlockCommentEmptyBlockComment() {
+    assertMacroExpansion(
+      """
+      @Test /* comment1 */ @Test /* comment2 */ var value: Int
+      """,
+      expandedSource: """
+        /* comment1 */ /* comment2 */ var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyWithLeadingSpace_SpacePreserved() {
+    assertMacroExpansion(
+      """
+      \u{0020}@Test var value: Int
+      """,
+      expandedSource: """
+        \u{0020}var value: Int
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyWithLeadingSpaceOnMember_SpacePreserved() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        \u{0020}@Test var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          \u{0020}var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyAndAttributeMashedTogether() {
+    // NB: In Swift, attributes can validly cozy up without whitespace.
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @Test@State var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testAttributeAndEmptyMashedTogether() {
+    // NB: In Swift, attributes can validly cozy up without whitespace.
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @State@Test var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyAndAttributeMashedTogether_VariableNextLine() {
+    // NB: In Swift, attributes can validly cozy up without whitespace.
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @Test@State
+        var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State
+          var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testAttributeAndEmptyMashedTogether_VariableNextLine() {
+    // NB: In Swift, attributes can validly cozy up without whitespace.
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @State@Test
+        var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State
+          var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testAttributeAndEmptyOnOwnLinesBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @State
+      @Test
+      var x: Int
+      """,
+      expandedSource: """
+        @State
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyAndAttributeOnOwnLinesBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @Test
+      @State
+      var x: Int
+      """,
+      expandedSource: """
+        @State
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testAttributeOnOwnLineThenEmptyBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @State
+      @Test var x: Int
+      """,
+      expandedSource: """
+        @State
+        var x: Int
+        """,
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyOnOwnLineThenEmptyBeforeVariable() {
+    assertMacroExpansion(
+      """
+      @Test
+      @Test var x: Int
+      """,
+      expandedSource: "var x: Int",
+      macros: [
+        "Test": NoOpPeerMacro.self
+      ]
+    )
+  }
+
+  func testEmptyOnMemberVariable() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @Test var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyBeforeAttributeOnSameLineAsMemberVariable() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @Test @State var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyAfterAttributeOnSameLineAsMemberVariable() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @State @Test var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyAfterAttributeOnSameLineAsMemberVariable_AwkwardWhitespace() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @State \t  @Test \t  var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State \t  var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testEmptyOnOwnLineThenAttributedMemberVariable() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @Test
+        @State var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+
+  func testAttributeOnOwnLineThenEmptyOnMemberVariable() {
+    assertMacroExpansion(
+      """
+      struct Foo {
+        @State
+        @Test var x: Int
+      }
+      """,
+      expandedSource: """
+        struct Foo {
+          @State
+          var x: Int
+        }
+        """,
+      macros: ["Test": NoOpPeerMacro.self]
+    )
+  }
+}