From 4d550d83487e738c97e8dcdac6d4a58e663a29e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Galen=20O=E2=80=99Hanlon?= <gohanlon@icloud.com>
Date: Fri, 29 Dec 2023 09:26:53 -0800
Subject: [PATCH] Don't replace parent syntax collection when targeting its
 first child

This pulls in the latest `FixItApplier` from swift-syntax `main`
(d647052), which is now String-based.

Fixes #15.
---
 Sources/MacroTesting/AssertMacro.swift        | 124 +++++++++---------
 .../MacroTesting/SwiftSyntax/SourceEdit.swift |  74 +++++++++++
 .../FixItApplier.swift                        | 109 +++++++++++++++
 Tests/MacroTestingTests/FixItTests.swift      |   6 +-
 4 files changed, 247 insertions(+), 66 deletions(-)
 create mode 100644 Sources/MacroTesting/SwiftSyntax/SourceEdit.swift
 create mode 100644 Sources/MacroTesting/_SwiftSyntaxTestSupport/FixItApplier.swift

diff --git a/Sources/MacroTesting/AssertMacro.swift b/Sources/MacroTesting/AssertMacro.swift
index a602827..755244c 100644
--- a/Sources/MacroTesting/AssertMacro.swift
+++ b/Sources/MacroTesting/AssertMacro.swift
@@ -250,10 +250,16 @@ public func assertMacro(
     if !allDiagnostics.isEmpty && allDiagnostics.allSatisfy({ !$0.fixIts.isEmpty }) {
       offset += 1
 
+      let edits =
+        context.diagnostics
+        .flatMap(\.fixIts)
+        .flatMap { $0.changes }
+        .map { $0.edit(in: context) }
+
       var fixedSourceFile = origSourceFile
       fixedSourceFile = Parser.parse(
-        source: FixItApplier.applyFixes(
-          context: context, in: allDiagnostics.map(anchor), to: origSourceFile
+        source: FixItApplier.apply(
+          edits: edits, to: origSourceFile
         )
         .description
       )
@@ -343,6 +349,57 @@ public func assertMacro(
   }
 }
 
+// From: https://github.com/apple/swift-syntax/blob/d647052/Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift
+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`
+  fileprivate 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
+      )
+    }
+  }
+}
+
+// From: https://github.com/apple/swift-syntax/blob/d647052/Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift
+extension BasicMacroExpansionContext {
+  /// Translates a position from a detached node to the corresponding position
+  /// in the original source file.
+  fileprivate 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)
+  }
+}
+
 /// Asserts that a given Swift source string matches an expected string with all macros expanded.
 ///
 /// See ``assertMacro(_:indentationWidth:record:of:diagnostics:fixes:expansion:file:function:line:column:)-pkfi``
@@ -619,69 +676,6 @@ extension Dictionary where Key == String, Value == Macro.Type {
   }
 }
 
-private class FixItApplier: SyntaxRewriter {
-  let context: BasicMacroExpansionContext
-  let diagnostics: [Diagnostic]
-
-  init(context: BasicMacroExpansionContext, diagnostics: [Diagnostic]) {
-    self.context = context
-    self.diagnostics = diagnostics
-    super.init(viewMode: .all)
-  }
-
-  public override func visitAny(_ node: Syntax) -> Syntax? {
-    for diagnostic in diagnostics {
-      for fixIts in diagnostic.fixIts {
-        for change in fixIts.changes {
-          switch change {
-          case .replace(let oldNode, let newNode):
-            let offset =
-              context
-              .location(for: oldNode.position, anchoredAt: oldNode, fileName: "")
-              .offset
-            if node.position.utf8Offset == offset {
-              return newNode
-            }
-          default:
-            break
-          }
-        }
-      }
-    }
-    return nil
-  }
-
-  override func visit(_ node: TokenSyntax) -> TokenSyntax {
-    var modifiedNode = node
-    for diagnostic in diagnostics {
-      for fixIts in diagnostic.fixIts {
-        for change in fixIts.changes {
-          switch change {
-          case .replaceLeadingTrivia(token: let changedNode, let newTrivia)
-          where changedNode.id == node.id:
-            modifiedNode = node.with(\.leadingTrivia, newTrivia)
-          case .replaceTrailingTrivia(token: let changedNode, let newTrivia)
-          where changedNode.id == node.id:
-            modifiedNode = node.with(\.trailingTrivia, newTrivia)
-          default:
-            break
-          }
-        }
-      }
-    }
-    return modifiedNode
-  }
-
-  public static func applyFixes(
-    context: BasicMacroExpansionContext,
-    in diagnostics: [Diagnostic],
-    to tree: some SyntaxProtocol
-  ) -> Syntax {
-    let applier = FixItApplier(context: context, diagnostics: diagnostics)
-    return applier.rewrite(tree)
-  }
-}
-
 private let oldPrefix = "\u{2212}"
 private let newPrefix = "+"
 private let prefix = "\u{2007}"
diff --git a/Sources/MacroTesting/SwiftSyntax/SourceEdit.swift b/Sources/MacroTesting/SwiftSyntax/SourceEdit.swift
new file mode 100644
index 0000000..50852b7
--- /dev/null
+++ b/Sources/MacroTesting/SwiftSyntax/SourceEdit.swift
@@ -0,0 +1,74 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+/// A textual edit to the original source represented by a range and a
+/// replacement.
+public struct SourceEdit: Equatable {
+  /// The half-open range that this edit applies to.
+  public let range: Range<AbsolutePosition>
+  /// The text to replace the original range with. Empty for a deletion.
+  public let replacement: String
+
+  /// Length of the original source range that this edit applies to. Zero if
+  /// this is an addition.
+  public var length: SourceLength {
+    return SourceLength(utf8Length: range.lowerBound.utf8Offset - range.upperBound.utf8Offset)
+  }
+
+  /// Create an edit to replace `range` in the original source with
+  /// `replacement`.
+  public init(range: Range<AbsolutePosition>, replacement: String) {
+    self.range = range
+    self.replacement = replacement
+  }
+
+  /// Convenience function to create a textual addition after the given node
+  /// and its trivia.
+  public static func insert(_ newText: String, after node: some SyntaxProtocol) -> SourceEdit {
+    return SourceEdit(range: node.endPosition..<node.endPosition, replacement: newText)
+  }
+
+  /// Convenience function to create a textual addition before the given node
+  /// and its trivia.
+  public static func insert(_ newText: String, before node: some SyntaxProtocol) -> SourceEdit {
+    return SourceEdit(range: node.position..<node.position, replacement: newText)
+  }
+
+  /// Convenience function to create a textual replacement of the given node,
+  /// including its trivia.
+  public static func replace(_ node: some SyntaxProtocol, with replacement: String) -> SourceEdit {
+    return SourceEdit(range: node.position..<node.endPosition, replacement: replacement)
+  }
+
+  /// Convenience function to create a textual deletion the given node and its
+  /// trivia.
+  public static func remove(_ node: some SyntaxProtocol) -> SourceEdit {
+    return SourceEdit(range: node.position..<node.endPosition, replacement: "")
+  }
+}
+
+extension SourceEdit: CustomDebugStringConvertible {
+  public var debugDescription: String {
+    let hasNewline = replacement.contains { $0.isNewline }
+    if hasNewline {
+      return #"""
+        \#(range.lowerBound.utf8Offset)-\#(range.upperBound.utf8Offset)
+        """
+        \#(replacement)
+        """
+        """#
+    }
+    return "\(range.lowerBound.utf8Offset)-\(range.upperBound.utf8Offset) \"\(replacement)\""
+  }
+}
diff --git a/Sources/MacroTesting/_SwiftSyntaxTestSupport/FixItApplier.swift b/Sources/MacroTesting/_SwiftSyntaxTestSupport/FixItApplier.swift
new file mode 100644
index 0000000..ec2da3b
--- /dev/null
+++ b/Sources/MacroTesting/_SwiftSyntaxTestSupport/FixItApplier.swift
@@ -0,0 +1,109 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftDiagnostics
+import SwiftSyntax
+import SwiftSyntaxMacroExpansion
+
+public enum FixItApplier {
+  /// Applies selected or all Fix-Its from the provided diagnostics to a given syntax tree.
+  ///
+  /// - Parameters:
+  ///   - diagnostics: An array of `Diagnostic` objects, each containing one or more Fix-Its.
+  ///   - filterByMessages: An optional array of message strings to filter which Fix-Its to apply.
+  ///     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.
+//  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 }
+//
+//    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 {
+      edits = Array(edits.dropFirst())
+
+      let startIndex = source.utf8.index(source.utf8.startIndex, offsetBy: edit.startUtf8Offset)
+      let endIndex = source.utf8.index(source.utf8.startIndex, offsetBy: edit.endUtf8Offset)
+
+      source.replaceSubrange(startIndex..<endIndex, with: edit.replacement)
+
+      edits = edits.compactMap { remainingEdit -> SourceEdit? in
+        if remainingEdit.replacementRange.overlaps(edit.replacementRange) {
+          // The edit overlaps with the previous edit. We can't apply both
+          // without conflicts. Apply the one that's listed first and drop the
+          // later edit.
+          return nil
+        }
+
+        // If the remaining edit starts after or at the end of the edit that we just applied,
+        // shift it by the current edit's difference in length.
+        if edit.endUtf8Offset <= remainingEdit.startUtf8Offset {
+          let startPosition = AbsolutePosition(
+            utf8Offset: remainingEdit.startUtf8Offset - edit.replacementRange.count
+              + edit.replacementLength)
+          let endPosition = AbsolutePosition(
+            utf8Offset: remainingEdit.endUtf8Offset - edit.replacementRange.count
+              + edit.replacementLength)
+          return SourceEdit(
+            range: startPosition..<endPosition, replacement: remainingEdit.replacement)
+        }
+
+        return remainingEdit
+      }
+    }
+
+    return source
+  }
+}
+
+extension SourceEdit {
+  fileprivate var startUtf8Offset: Int {
+    return range.lowerBound.utf8Offset
+  }
+
+  fileprivate var endUtf8Offset: Int {
+    return range.upperBound.utf8Offset
+  }
+
+  fileprivate var replacementLength: Int {
+    return replacement.utf8.count
+  }
+
+  fileprivate var replacementRange: Range<Int> {
+    return startUtf8Offset..<endUtf8Offset
+  }
+}
diff --git a/Tests/MacroTestingTests/FixItTests.swift b/Tests/MacroTestingTests/FixItTests.swift
index bc8ba95..232aa4a 100644
--- a/Tests/MacroTestingTests/FixItTests.swift
+++ b/Tests/MacroTestingTests/FixItTests.swift
@@ -52,7 +52,7 @@ final class FixItTests: BaseTestCase {
     }
   }
 
-  func testReplaceFirstMember_IncorrectlyReplacesParent() {
+  func testReplaceFirstMember() {
     assertMacro {
       """
       @ReplaceFirstMember
@@ -79,12 +79,16 @@ final class FixItTests: BaseTestCase {
       @ReplaceFirstMember
       struct FooBar {
         let oye: Oye
+        let bar: Bar
+        let baz: Baz
       }
       """
     } expansion: {
       """
       struct FooBar {
         let oye: Oye
+        let bar: Bar
+        let baz: Baz
       }
       """
     }