Merge pull request #2215 from gohanlon/improve-attribute-remover-trivia

Fix trivia handling in `AttributeRemover` of `MacroSystem`

Author: ahoppen

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
   }
 }