Adjust locations for edits based on detached nodes when computing fixedSource

The Fix-Its returned by the macro are defined in terms of the nodes that it got passed and thus their positions are relative to the start of the detached nodes that get passed to the macro. We need to translate these offsets back to the the offsets in the original source file before computing edits and apply the offset-based edits.

Fixes #2332
rdar://118012820

Author: ahoppen

Package.swift

@@ -93,7 +93,7 @@ let package = Package(
 
     .target(
       name: "_SwiftSyntaxTestSupport",
-      dependencies: ["SwiftBasicFormat", "SwiftSyntax", "SwiftSyntaxBuilder"]
+      dependencies: ["SwiftBasicFormat", "SwiftSyntax", "SwiftSyntaxBuilder", "SwiftSyntaxMacroExpansion"]
     ),
 
     .testTarget(

Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift

@@ -325,7 +325,16 @@ public func assertMacroExpansion(
 
   // Applying Fix-Its
   if let expectedFixedSource = expectedFixedSource {
-    let fixedTree = FixItApplier.applyFixes(from: context.diagnostics, filterByMessages: applyFixIts, to: origSourceFile)
+    let messages = applyFixIts ?? context.diagnostics.compactMap { $0.fixIts.first?.message.message }
+
+    let edits =
+      context.diagnostics
+      .flatMap(\.fixIts)
+      .filter { messages.contains($0.message.message) }
+      .flatMap { $0.changes }
+      .map { $0.edit(in: context) }
+
+    let fixedTree = FixItApplier.apply(edits: edits, to: origSourceFile)
     let fixedTreeDescription = fixedTree.description
     assertStringsEqualWithDiff(
       fixedTreeDescription.trimmingTrailingWhitespace(),
@@ -335,3 +344,50 @@ public func assertMacroExpansion(
     )
   }
 }
+
+fileprivate extension FixIt.Change {
+  /// Returns the edit for this change, translating positions from detached nodes
+  /// to the corresponding locations in the original source file based on
+  /// `expansionContext`.
+  ///
+  /// - SeeAlso: `FixIt.Change.edit`
+  func edit(in expansionContext: BasicMacroExpansionContext) -> SourceEdit {
+    switch self {
+    case .replace(let oldNode, let newNode):
+      let start = expansionContext.position(of: oldNode.position, anchoredAt: oldNode)
+      let end = expansionContext.position(of: oldNode.endPosition, anchoredAt: oldNode)
+      return SourceEdit(
+        range: start..<end,
+        replacement: newNode.description
+      )
+
+    case .replaceLeadingTrivia(let token, let newTrivia):
+      let start = expansionContext.position(of: token.position, anchoredAt: token)
+      let end = expansionContext.position(of: token.positionAfterSkippingLeadingTrivia, anchoredAt: token)
+      return SourceEdit(
+        range: start..<end,
+        replacement: newTrivia.description
+      )
+
+    case .replaceTrailingTrivia(let token, let newTrivia):
+      let start = expansionContext.position(of: token.endPositionBeforeTrailingTrivia, anchoredAt: token)
+      let end = expansionContext.position(of: token.endPosition, anchoredAt: token)
+      return SourceEdit(
+        range: start..<end,
+        replacement: newTrivia.description
+      )
+    }
+  }
+}
+
+fileprivate extension BasicMacroExpansionContext {
+  /// Translates a position from a detached node to the corresponding position
+  /// in the original source file.
+  func position(
+    of position: AbsolutePosition,
+    anchoredAt node: some SyntaxProtocol
+  ) -> AbsolutePosition {
+    let location = self.location(for: position, anchoredAt: Syntax(node), fileName: "")
+    return AbsolutePosition(utf8Offset: location.offset)
+  }
+}

Sources/_SwiftSyntaxTestSupport/FixItApplier.swift

@@ -12,6 +12,7 @@
 
 import SwiftDiagnostics
 import SwiftSyntax
+import SwiftSyntaxMacroExpansion
 
 public enum FixItApplier {
   /// Applies selected or all Fix-Its from the provided diagnostics to a given syntax tree.
@@ -22,20 +23,34 @@ public enum FixItApplier {
   ///     If `nil`, the first Fix-It from each diagnostic is applied.
   ///   - tree: The syntax tree to which the Fix-Its will be applied.
   ///
-  /// - Returns: A ``String`` representation of the modified syntax tree after applying the Fix-Its.
+  /// - Returns: A `String` representation of the modified syntax tree after applying the Fix-Its.
   public static func applyFixes(
     from diagnostics: [Diagnostic],
     filterByMessages messages: [String]?,
     to tree: any SyntaxProtocol
   ) -> String {
     let messages = messages ?? diagnostics.compactMap { $0.fixIts.first?.message.message }
 
-    var edits =
+    let edits =
       diagnostics
       .flatMap(\.fixIts)
       .filter { messages.contains($0.message.message) }
       .flatMap(\.edits)
 
+    return self.apply(edits: edits, to: tree)
+  }
+
+  /// Apply the given edits to the syntax tree.
+  ///
+  /// - Parameters:
+  ///   - edits: The edits to apply to the syntax tree
+  ///   - tree: he syntax tree to which the edits should be applied.
+  /// - Returns: A `String` representation of the modified syntax tree after applying the edits.
+  public static func apply(
+    edits: [SourceEdit],
+    to tree: any SyntaxProtocol
+  ) -> String {
+    var edits = edits
     var source = tree.description
 
     while let edit = edits.first {

Tests/SwiftSyntaxMacroExpansionTest/PeerMacroTests.swift

@@ -247,4 +247,38 @@ final class PeerMacroTests: XCTestCase {
       indentationWidth: indentationWidth
     )
   }
+
+  func testAdjustFixItLocationsWhenComputingFixedSource() {
+    // Test that we adjust the locations of the Fix-Its to the original source
+    // before computing the `fixedSource` if the macro doesn't start at the
+    // start of the file.
+    assertMacroExpansion(
+      """
+      func other() {}
+
+      @addCompletionHandler
+      func f(a: Int, for b: String, _ value: Double) -> String { }
+      """,
+      expandedSource: """
+        func other() {}
+        func f(a: Int, for b: String, _ value: Double) -> String { }
+        """,
+      diagnostics: [
+        DiagnosticSpec(
+          message: "can only add a completion-handler variant to an 'async' function",
+          line: 4,
+          column: 1,
+          fixIts: [FixItSpec(message: "add 'async'")]
+        )
+      ],
+      macros: ["addCompletionHandler": AddCompletionHandler.self],
+      fixedSource: """
+        func other() {}
+
+        @addCompletionHandler
+        func f(a: Int, for b: String, _ value: Double) async-> String { }
+        """,
+      indentationWidth: indentationWidth
+    )
+  }
 }