From 8dee94c4b628bef28412d425aab48a3fb6056439 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Galen=20O=E2=80=99Hanlon?= <gohanlon@icloud.com>
Date: Thu, 11 Jul 2024 22:54:25 -0700
Subject: [PATCH] Add AttributeRemover600 and conditional usage

- AttributeRemover600 is vendored from SwiftSyntax
- Conditionally use AttributeRemover600 or AttributeRemover509
---
 Sources/MacroTesting/AssertMacro.swift        |  23 ++--
 .../Internal/AttributeRemover600.swift        | 114 ++++++++++++++++++
 2 files changed, 128 insertions(+), 9 deletions(-)
 create mode 100644 Sources/MacroTesting/Internal/AttributeRemover600.swift

diff --git a/Sources/MacroTesting/AssertMacro.swift b/Sources/MacroTesting/AssertMacro.swift
index e26aba1..377fdeb 100644
--- a/Sources/MacroTesting/AssertMacro.swift
+++ b/Sources/MacroTesting/AssertMacro.swift
@@ -278,17 +278,22 @@ public func assertMacro(
       // TODO: write a test where didExpand returns false
       // For now, covered in MemberwiseInitTests.testAppliedToEnum_FailsWithDiagnostic
       var didExpand: Bool {
-        let origSourceWithMacroAttributesRemoved = MacroTesting.AttributeRemover509(
-          removingWhere: {
-            guard let name = $0.attributeName.as(IdentifierTypeSyntax.self)?.name.text
-            else { return false }
-            return macros.keys.contains(name)
-          }
-        ).rewrite(origSourceFile)
+        let removingWhere: (AttributeSyntax) -> Bool = {
+          guard let name = $0.attributeName.as(IdentifierTypeSyntax.self)?.name.text
+          else { return false }
+          return macros.keys.contains(name)
+        }
+
+        #if canImport(SwiftSyntax600)
+          let remover = MacroTesting.AttributeRemover600(removingWhere: removingWhere)
+        #else
+          let remover = MacroTesting.AttributeRemover509(removingWhere: removingWhere)
+        #endif
+
+        let removedSource = remover.rewrite(origSourceFile)
 
         return expandedSourceFile.description.trimmingCharacters(in: .newlines)
-          != origSourceWithMacroAttributesRemoved.description.trimmingCharacters(
-            in: .newlines)
+          != removedSource.description.trimmingCharacters(in: .newlines)
       }
 
       if didExpand {
diff --git a/Sources/MacroTesting/Internal/AttributeRemover600.swift b/Sources/MacroTesting/Internal/AttributeRemover600.swift
new file mode 100644
index 0000000..e347761
--- /dev/null
+++ b/Sources/MacroTesting/Internal/AttributeRemover600.swift
@@ -0,0 +1,114 @@
+import SwiftSyntax
+
+public class AttributeRemover600: SyntaxRewriter {
+  let predicate: (AttributeSyntax) -> Bool
+
+  var triviaToAttachToNextToken: Trivia = Trivia()
+
+  /// Initializes an attribute remover with a given predicate to determine which attributes to remove.
+  ///
+  /// - Parameter predicate: A closure that determines whether a given `AttributeSyntax` should be removed.
+  ///   If this closure returns `true` for an attribute, that attribute will be removed.
+  public init(removingWhere predicate: @escaping (AttributeSyntax) -> Bool) {
+    self.predicate = predicate
+  }
+
+  public override func visit(_ node: AttributeListSyntax) -> AttributeListSyntax {
+    var filteredAttributes: [AttributeListSyntax.Element] = []
+    for case .attribute(let attribute) in node {
+      if self.predicate(attribute) {
+        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),
+          attribute.trailingTrivia.isEmpty,
+          let nextToken = attribute.nextToken(viewMode: .sourceAccurate),
+          !nextToken.leadingTrivia.isEmpty
+        {
+          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 = 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(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)
+  }
+
+  public override func visit(_ token: TokenSyntax) -> TokenSyntax {
+    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)
+  }
+}
+
+extension Trivia {
+  fileprivate func trimmingPrefix(
+    while predicate: (TriviaPiece) -> Bool
+  ) -> Trivia {
+    Trivia(pieces: self.drop(while: predicate))
+  }
+
+  fileprivate func trimmingSuffix(
+    while predicate: (TriviaPiece) -> Bool
+  ) -> Trivia {
+    Trivia(
+      pieces: self[...]
+        .reversed()
+        .drop(while: predicate)
+        .reversed()
+    )
+  }
+
+  fileprivate var startsWithNewline: Bool {
+    self.first?.isNewline ?? false
+  }
+}