From 80113cd9b4b8eae8bb0ba9626ea81d4f0c4cee43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Galen=20O=E2=80=99Hanlon?= <gohanlon@icloud.com>
Date: Thu, 28 Sep 2023 16:41:04 -0700
Subject: [PATCH] Isolate `AttributeRemoverTests` to `AttributeRemover`

* Add a private assert helper: `assertSyntaxRemovingTestAttributes`
* Make `AttributeRemover` SPI (`@_spi(Testing)`)
* Rather than depending on `SyntaxNode`'s id-only `Equatable`
  conformance, refactor `AttributeRemover` to remove attributes matching
  a predicate: `AttributeRemover.init(removingWhere:)` of type
  `(AttributeSyntax) -> Bool)`. This change allows `AttributeRemover` to
  be used without first plucking the `AttributeSyntax` values to be
  removed from a live tree.

Note: Unlike `assertMacroExpansion`, the new
`assertSyntaxRemovingTestAttributes` does not trim newlines before
comparing values. The trimming by `assertMacroExpansion` was masking a
(minor) bug in `AttributeRemover` where an extra leading newline remains
when the removed attribute has no proceeding tokens. For now, this
commit "fixes" the failing tests by including the unwanted leading
newlines in the expected test output, while also adding "FIXME" comments
to draw attention to the need for a future fix.
---
 .../MacroSystem.swift                         |  21 +-
 .../AttributeRemoverTests.swift               | 293 ++++++++----------
 2 files changed, 139 insertions(+), 175 deletions(-)

diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
index f1699830691..0e4a54a8d6c 100644
--- a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
+++ b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
@@ -384,19 +384,24 @@ struct MacroSystem {
 // MARK: - MacroApplication
 
 /// Removes attributes from a syntax tree while maintaining their surrounding trivia.
-private class AttributeRemover: SyntaxRewriter {
-  var attributesToRemove: [AttributeSyntax]
+@_spi(Testing)
+public class AttributeRemover: SyntaxRewriter {
+  let predicate: (AttributeSyntax) -> Bool
 
   var triviaToAttachToNextToken: Trivia = Trivia()
 
-  init(attributesToRemove: [AttributeSyntax]) {
-    self.attributesToRemove = attributesToRemove
+  /// 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
   }
 
-  override func visit(_ node: AttributeListSyntax) -> AttributeListSyntax {
+  public override func visit(_ node: AttributeListSyntax) -> AttributeListSyntax {
     var filteredAttributes: [AttributeListSyntax.Element] = []
     for case .attribute(let attribute) in node {
-      if attributesToRemove.contains(attribute) {
+      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,
@@ -450,7 +455,7 @@ private class AttributeRemover: SyntaxRewriter {
     return AttributeListSyntax(filteredAttributes)
   }
 
-  override func visit(_ token: TokenSyntax) -> TokenSyntax {
+  public override func visit(_ token: TokenSyntax) -> TokenSyntax {
     return prependAndClearAccumulatedTrivia(to: token)
   }
 
@@ -573,7 +578,7 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
 
       let attributesToRemove = self.macroAttributes(attachedTo: visitedNode).map(\.attributeNode)
 
-      return AttributeRemover(attributesToRemove: attributesToRemove).rewrite(visitedNode)
+      return AttributeRemover(removingWhere: { attributesToRemove.contains($0) }).rewrite(visitedNode)
     }
 
     return nil
diff --git a/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift
index cd16b029bda..1858bbd92bd 100644
--- a/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift
+++ b/Tests/SwiftSyntaxMacroExpansionTest/AttributeRemoverTests.swift
@@ -10,144 +10,131 @@
 //
 //===----------------------------------------------------------------------===//
 
-//==========================================================================//
-// 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 SwiftParser
 import SwiftSyntax
-import SwiftSyntaxMacroExpansion
-import SwiftSyntaxMacros
-import SwiftSyntaxMacrosTestSupport
+@_spi(Testing) import SwiftSyntaxMacroExpansion
 import XCTest
-
-fileprivate struct NoOpPeerMacro: PeerMacro {
-  static func expansion(
-    of node: AttributeSyntax,
-    providingPeersOf declaration: some DeclSyntaxProtocol,
-    in context: some MacroExpansionContext
-  ) throws -> [DeclSyntax] {
-    return []
-  }
+import _SwiftSyntaxTestSupport
+
+fileprivate func assertSyntaxRemovingTestAttributes(
+  _ originalSource: String,
+  reduction expectedReducedSource: String,
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  let attributeToRemove = AttributeSyntax(stringLiteral: "@Test")
+
+  let reducedSource = AttributeRemover(
+    removingWhere: { $0.trimmedDescription == attributeToRemove.trimmedDescription }
+  )
+  .rewrite(
+    Parser.parse(source: originalSource)
+  )
+
+  assertStringsEqualWithDiff(
+    reducedSource.description,
+    expectedReducedSource,
+    "Attribute removal did not produce the expected reduced source",
+    additionalInfo: """
+      Actual reduced source:
+      \(reducedSource)
+      """,
+    file: file,
+    line: line
+  )
 }
 
 final class AttributeRemoverTests: XCTestCase {
   func testEmptyOnSameLineAsVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       "@Test var x: Int",
-      expandedSource: "var x: Int",
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+      reduction: "var x: Int"
     )
   }
 
   func testEmptyTwiceOnSameLineAsVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       "@Test @Test var x: Int",
-      expandedSource: "var x: Int",
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+      reduction: "var x: Int"
     )
   }
 
+  // FIXME: `AttributeRemover` should not leave a leading newline.
   func testEmptyOnOwnLineBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test
       var x: Int
       """,
-      expandedSource: "var x: Int",
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+      reduction: "\nvar x: Int"
     )
   }
 
+  // FIXME: `AttributeRemover` should not leave a leading newline.
   func testEmptyTwiceOnOwnLineBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test @Test
       var x: Int
       """,
-      expandedSource: "var x: Int",
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+      reduction: "\nvar x: Int"
     )
   }
 
   func testEmptyAndAttributeOnOwnLineBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test @State
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
   func testAttributeAndEmptyOnOwnLineBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State @Test
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
   func testAttributeAndEmptyAndCommentOnOwnLineBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State @Test// comment
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State // comment
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
   func testAttributeAndEmptyAndCommentOnOwnLineBeforeVariable2() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State @Test // comment
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State // comment
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
   func testCommentsAroundEmpty() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct S {
         /// Some doc comment
@@ -155,360 +142,332 @@ final class AttributeRemoverTests: XCTestCase {
         var value: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct S {
           /// Some doc comment
           /* trailing */
           var value: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testBlockCommentNewlineEmpty() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       /* comment */
       @Test
       var value: Int
       """,
-      expandedSource: """
+      reduction: """
         /* comment */
         var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
-  func testEmpyNewlineBlockComment() {
-    assertMacroExpansion(
+  // FIXME: `AttributeRemover` should not leave a leading newline.
+  func testEmptyNewlineBlockComment() {
+    assertSyntaxRemovingTestAttributes(
       """
       @Test
       /* comment */
       var value: Int
       """,
-      expandedSource: """
-        /* comment */
+      reduction: """
+        \n/* comment */
         var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testAttributeNewlineBlockCommentEmpty() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State
       /*doc comment*/@Test
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State
         /*doc comment*/
         var x: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyBlockCommentEmpty() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test /* comment */ @Test var value: Int
       """,
-      expandedSource: """
+      reduction: """
         /* comment */ var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyBlockCommentEmptyNewline() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test /* comment */ @Test
       var value: Int
       """,
-      expandedSource: """
+      reduction: """
         /* comment */
         var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyBlockCommentEmptyBlockComment() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test /* comment1 */ @Test /* comment2 */ var value: Int
       """,
-      expandedSource: """
+      reduction: """
         /* comment1 */ /* comment2 */ var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyWithLeadingSpace_SpacePreserved() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       \u{0020}@Test var value: Int
       """,
-      expandedSource: """
+      reduction: """
         \u{0020}var value: Int
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyWithLeadingSpaceOnMember_SpacePreserved() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         \u{0020}@Test var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           \u{0020}var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyAndAttributeMashedTogether() {
     // NB: In Swift, attributes can validly cozy up without whitespace.
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @Test@State var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testAttributeAndEmptyMashedTogether() {
     // NB: In Swift, attributes can validly cozy up without whitespace.
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @State@Test var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyAndAttributeMashedTogether_VariableNextLine() {
     // NB: In Swift, attributes can validly cozy up without whitespace.
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @Test@State
         var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State
           var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testAttributeAndEmptyMashedTogether_VariableNextLine() {
     // NB: In Swift, attributes can validly cozy up without whitespace.
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @State@Test
         var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State
           var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testAttributeAndEmptyOnOwnLinesBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State
       @Test
       var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
+  // FIXME: `AttributeRemover` should not leave a leading newline.
   func testEmptyAndAttributeOnOwnLinesBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test
       @State
       var x: Int
       """,
-      expandedSource: """
-        @State
+      reduction: """
+        \n@State
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
   func testAttributeOnOwnLineThenEmptyBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @State
       @Test var x: Int
       """,
-      expandedSource: """
+      reduction: """
         @State
         var x: Int
-        """,
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+        """
     )
   }
 
+  // FIXME: `AttributeRemover` should not leave a leading newline.
   func testEmptyOnOwnLineThenEmptyBeforeVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       @Test
       @Test var x: Int
       """,
-      expandedSource: "var x: Int",
-      macros: [
-        "Test": NoOpPeerMacro.self
-      ]
+      reduction: "\nvar x: Int"
     )
   }
 
   func testEmptyOnMemberVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @Test var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyBeforeAttributeOnSameLineAsMemberVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @Test @State var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyAfterAttributeOnSameLineAsMemberVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @State @Test var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyAfterAttributeOnSameLineAsMemberVariable_AwkwardWhitespace() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @State \t  @Test \t  var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State \t  var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testEmptyOnOwnLineThenAttributedMemberVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @Test
         @State var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 
   func testAttributeOnOwnLineThenEmptyOnMemberVariable() {
-    assertMacroExpansion(
+    assertSyntaxRemovingTestAttributes(
       """
       struct Foo {
         @State
         @Test var x: Int
       }
       """,
-      expandedSource: """
+      reduction: """
         struct Foo {
           @State
           var x: Int
         }
-        """,
-      macros: ["Test": NoOpPeerMacro.self]
+        """
     )
   }
 }