Merge pull request #2048 from ahoppen/ahoppen/incrorrect-trivia-wrapping

Fix trivia transfer bugs for freestanding expression macros

Author: ahoppen

Sources/SwiftSyntaxMacroExpansion/IndentationUtils.swift

@@ -85,6 +85,13 @@ extension Trivia {
     }
     return Trivia(pieces: resultPieces)
   }
+
+  /// Remove all indentation from the last line of this trivia
+  var removingIndentationOnLastLine: Trivia {
+    let lastNewlineIndex = pieces.lastIndex(where: \.isNewline) ?? pieces.startIndex
+
+    return Trivia(pieces: pieces[..<lastNewlineIndex]) + Trivia(pieces: pieces[lastNewlineIndex...]).removingIndentation
+  }
 }
 
 extension SyntaxProtocol {

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -66,7 +66,7 @@ private func expandFreestandingMemberDeclList(
   let indentedSource =
     expanded
     .indented(by: node.indentationOfFirstLine)
-    .wrappingInNonCommentTrivia(from: node)
+    .wrappingInTrivia(from: node)
   return "\(raw: indentedSource)"
 }
 
@@ -94,7 +94,7 @@ private func expandFreestandingCodeItemList(
   let indentedSource =
     expanded
     .indented(by: node.indentationOfFirstLine)
-    .wrappingInNonCommentTrivia(from: node)
+    .wrappingInTrivia(from: node)
   return "\(raw: indentedSource)"
 }
 
@@ -119,7 +119,7 @@ private func expandFreestandingExpr(
   let indentedSource =
     expanded
     .indented(by: node.indentationOfFirstLine)
-    .wrappingInNonCommentTrivia(from: node)
+    .wrappingInTrivia(from: node)
   return "\(raw: indentedSource)"
 }
 
@@ -947,17 +947,18 @@ private extension AccessorBlockSyntax {
 }
 
 private extension String {
-  /// Add any non-comment trivia from `node` before/after this string.
+  /// Add any trivia from `node` before/after this string.
   ///
-  /// We need to do this because the macro is responsible for copying trivia
-  /// from the freestanding macro to the generated declaration.
-  ///
-  /// Essentially, what we want to keep any empty newlines in front of the
-  /// freestanding macro that separate it from the previous declarations.
-  func wrappingInNonCommentTrivia(from node: some SyntaxProtocol) -> String {
-    return node.leadingTrivia.removingComments.removingIndentation.description
+  /// We need to do this because we are replacing the entire macro expansion
+  /// declaration / expression with the expanded source but semantically, the
+  /// user should think about it as just replacing the `#...` expression without
+  /// any trivia.
+  func wrappingInTrivia(from node: some SyntaxProtocol) -> String {
+    // We need to remove the indentation from the last line because the macro
+    // expansion buffer already contains the indentation.
+    return node.leadingTrivia.removingIndentationOnLastLine.description
       + self
-      + node.leadingTrivia.removingComments.removingIndentation.description
+      + node.trailingTrivia.description
   }
 }
 
@@ -972,37 +973,3 @@ private extension SyntaxProtocol {
     return self.detached
   }
 }
-
-private extension Trivia {
-  /// Drop all comments from the trivia.
-  ///
-  /// If a comment is the only entry on a line, drop the entire line instead of
-  /// leaving an empty line.
-  var removingComments: Trivia {
-    var result: [TriviaPiece] = []
-
-    var lineContainedComment = false
-    var lineContainedNonWhitespaceNonComment = false
-    for piece in self.pieces {
-      switch piece {
-      case .spaces, .tabs:
-        result.append(piece)
-      case .backslashes, .formfeeds, .pounds, .unexpectedText, .verticalTabs:
-        lineContainedNonWhitespaceNonComment = true
-        result.append(piece)
-      case .blockComment, .docBlockComment, .docLineComment, .lineComment:
-        lineContainedComment = true
-      case .carriageReturns, .carriageReturnLineFeeds, .newlines:
-        if lineContainedComment && !lineContainedNonWhitespaceNonComment {
-          continue
-        } else {
-          result.append(piece)
-        }
-        lineContainedComment = false
-        lineContainedNonWhitespaceNonComment = false
-      }
-    }
-
-    return Trivia(pieces: result)
-  }
-}

Tests/SwiftSyntaxMacroExpansionTest/MacroSystemTests.swift

@@ -701,6 +701,25 @@ final class MacroSystemTests: XCTestCase {
     )
   }
 
+  func testTriviaTransferOnExpressionMacro() {
+    assertMacroExpansion(
+      """
+      // Ignore me
+      \t
+      // Capture me
+      #stringify(x)
+      """,
+      expandedSource: """
+        // Ignore me
+        \t
+        // Capture me
+        (x, "x")
+        """,
+      macros: testMacros,
+      indentationWidth: indentationWidth
+    )
+  }
+
   func testCommentsOnExpressionMacro() {
     assertMacroExpansion(
       """
@@ -709,7 +728,7 @@ final class MacroSystemTests: XCTestCase {
       """,
       expandedSource: """
         let b = 
-        (x + y, "x + y")
+        /*leading */ (x + y, "x + y") /*trailing*/
         """,
       macros: testMacros,
       indentationWidth: indentationWidth
@@ -1441,10 +1460,11 @@ final class MacroSystemTests: XCTestCase {
       ) /* trailing comment */
       """,
       expandedSource: """
+        // some comment
         func foo() {
         }
         func bar() {
-        }
+        } /* trailing comment */
         """,
       macros: ["decls": DeclsFromStringsMacro.self],
       indentationWidth: indentationWidth
@@ -1464,10 +1484,11 @@ final class MacroSystemTests: XCTestCase {
       """,
       expandedSource: """
         struct Foo {
+          // some comment
           func foo() {
           }
           func bar() {
-          }
+          } /* trailing comment */
         }
         """,
       macros: ["decls": DeclsFromStringsMacro.self],