From 1c4c1cd6cdf1fbeb5cc3329540085e55b6b740d8 Mon Sep 17 00:00:00 2001
From: Kai Lau <kai@appapp.works>
Date: Fri, 26 Jul 2024 20:20:59 -0700
Subject: [PATCH] introduced `.replaceChild` to `FixIt.Change` to allow
 replacing an optional child node with a proper node

- introduced `ReplacingChildData` as the type-erased payload for `.replaceChild`
- added an entry to 601.md describing this API-breaking change
---
 Release Notes/601.md                          |  6 +++
 .../Diagnostics.swift                         |  5 +-
 .../PluginMacroExpansionContext.swift         | 19 +++++---
 Sources/SwiftDiagnostics/Convenience.swift    | 16 +++++++
 Sources/SwiftDiagnostics/FixIt.swift          | 47 +++++++++++++++++++
 .../Assertions.swift                          |  9 ++++
 .../PeerMacroTests.swift                      | 17 ++-----
 7 files changed, 99 insertions(+), 20 deletions(-)

diff --git a/Release Notes/601.md b/Release Notes/601.md
index 7036e503b33..349874500d8 100644
--- a/Release Notes/601.md	
+++ b/Release Notes/601.md	
@@ -34,6 +34,12 @@
   - Description: Allows retrieving the radix value from the `literal.text`.
   - Issue: https://github.com/apple/swift-syntax/issues/405
   - Pull Request: https://github.com/apple/swift-syntax/pull/2605
+  
+- `FixIt.Change` gained a new case `replaceChild(data:)`.
+  - Description: The new case covers the replacement of a child node with another node.
+  - Issue: https://github.com/swiftlang/swift-syntax/issues/2205
+  - Pull Request: https://github.com/swiftlang/swift-syntax/pull/2758
+  - Migration steps: In exhaustive switches over `FixIt.Change`, cover the new case.
 
 ## Template
 
diff --git a/Sources/SwiftCompilerPluginMessageHandling/Diagnostics.swift b/Sources/SwiftCompilerPluginMessageHandling/Diagnostics.swift
index 027ec68838b..73a1d5dc034 100644
--- a/Sources/SwiftCompilerPluginMessageHandling/Diagnostics.swift
+++ b/Sources/SwiftCompilerPluginMessageHandling/Diagnostics.swift
@@ -130,12 +130,15 @@ extension PluginMessage.Diagnostic {
               to: .afterTrailingTrivia
             )
             text = newTrivia.description
+          case .replaceChild(let replaceChildData):
+            range = sourceManager.range(replaceChildData.replacementRange, in: replaceChildData.parent)
+            text = replaceChildData.newChild.description
           #if RESILIENT_LIBRARIES
           @unknown default:
             fatalError()
           #endif
           }
-          guard let range = range else {
+          guard let range else {
             return nil
           }
           return .init(
diff --git a/Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift b/Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift
index 8914d159823..4c4a5fc317b 100644
--- a/Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift
+++ b/Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift
@@ -169,20 +169,25 @@ class SourceManager {
     of node: Syntax,
     from startKind: PositionInSyntaxNode = .afterLeadingTrivia,
     to endKind: PositionInSyntaxNode = .beforeTrailingTrivia
+  ) -> SourceRange? {
+    range(node.position(at: startKind)..<node.position(at: endKind), in: node)
+  }
+
+  /// Get ``SourceRange`` (file name + UTF-8 offset range) of `localRange` in `node`'s root node, which must be one
+  /// of the returned values from `add(_:)`.
+  func range(
+    _ localRange: @autoclosure () -> Range<AbsolutePosition>,
+    in node: some SyntaxProtocol
   ) -> SourceRange? {
     guard let base = self.knownSourceSyntax[node.root.id] else {
       return nil
     }
-    let localStartPosition = node.position(at: startKind)
-    let localEndPosition = node.position(at: endKind)
-    precondition(localStartPosition <= localEndPosition)
-
     let positionOffset = base.location.offset
-
+    let localRange = localRange()
     return SourceRange(
       fileName: base.location.fileName,
-      startUTF8Offset: localStartPosition.advanced(by: positionOffset).utf8Offset,
-      endUTF8Offset: localEndPosition.advanced(by: positionOffset).utf8Offset
+      startUTF8Offset: localRange.lowerBound.advanced(by: positionOffset).utf8Offset,
+      endUTF8Offset: localRange.upperBound.advanced(by: positionOffset).utf8Offset
     )
   }
 
diff --git a/Sources/SwiftDiagnostics/Convenience.swift b/Sources/SwiftDiagnostics/Convenience.swift
index 4262272201e..74225b635bd 100644
--- a/Sources/SwiftDiagnostics/Convenience.swift
+++ b/Sources/SwiftDiagnostics/Convenience.swift
@@ -51,4 +51,20 @@ extension FixIt {
       ]
     )
   }
+
+  public static func replaceChild<Parent: SyntaxProtocol, Child: SyntaxProtocol>(
+    message: FixItMessage,
+    parent: Parent,
+    replacingChildAt keyPath: WritableKeyPath<Parent, Child?> & Sendable,
+    with newChild: Child
+  ) -> Self {
+    FixIt(
+      message: message,
+      changes: [
+        .replaceChild(
+          data: FixIt.Change.ReplacingOptionalChildData(parent: parent, newChild: newChild, keyPath: keyPath)
+        )
+      ]
+    )
+  }
 }
diff --git a/Sources/SwiftDiagnostics/FixIt.swift b/Sources/SwiftDiagnostics/FixIt.swift
index 3dbabed7dbf..05e43f05012 100644
--- a/Sources/SwiftDiagnostics/FixIt.swift
+++ b/Sources/SwiftDiagnostics/FixIt.swift
@@ -27,15 +27,56 @@ public protocol FixItMessage: Sendable {
   var fixItID: MessageID { get }
 }
 
+/// Types conforming to this protocol provide the data required for replacing a child node of a parent node.
+///
+/// Conforming types should ensure the child of ``parent`` to be replaced at ``replacementRange`` is type-compatible
+/// with ``newChild``. Conforming types are stored as type-erased existentials (i.e. `any ReplacingChildData`) in
+///  ``FixIt/Change/replaceChild(data:)`` to keep ``FixIt`` type-erased.
+public protocol ReplacingChildData: Sendable {
+  associatedtype Parent: SyntaxProtocol
+  associatedtype Child: SyntaxProtocol
+
+  /// The node whose child node at ``replacementRange`` to be replaced by ``newChild``.
+  var parent: Parent { get }
+
+  /// The node to replace the child node of ``parent`` at ``replacementRange``.
+  var newChild: Child { get }
+
+  /// The absolute position range of the child node to be replaced.
+  ///
+  /// If a nil child node is to be replaced, conforming types should provide a zero-length range with both bounds
+  /// denoting the start position of ``newChild`` in ``parent`` after replacement.
+  var replacementRange: Range<AbsolutePosition> { get }
+}
+
 /// A Fix-It that can be applied to resolve a diagnostic.
 public struct FixIt: Sendable {
   public enum Change: Sendable {
+    struct ReplacingOptionalChildData<Parent: SyntaxProtocol, Child: SyntaxProtocol>: ReplacingChildData {
+      let parent: Parent
+      let newChild: Child
+      let keyPath: WritableKeyPath<Parent, Child?> & Sendable
+
+      var replacementRange: Range<AbsolutePosition> {
+        // need to upcast keyPath to strip Sendable for older Swift versions
+        let keyPath: WritableKeyPath<Parent, Child?> = keyPath
+        if let oldChild = parent[keyPath: keyPath] {
+          return oldChild.range
+        } else {
+          let newChild = parent.with(keyPath, newChild)[keyPath: keyPath]!
+          return newChild.position..<newChild.position
+        }
+      }
+    }
+
     /// Replace `oldNode` by `newNode`.
     case replace(oldNode: Syntax, newNode: Syntax)
     /// Replace the leading trivia on the given token
     case replaceLeadingTrivia(token: TokenSyntax, newTrivia: Trivia)
     /// Replace the trailing trivia on the given token
     case replaceTrailingTrivia(token: TokenSyntax, newTrivia: Trivia)
+    /// Replace the child node of the given parent node at the given replacement range with the given new child node
+    case replaceChild(data: any ReplacingChildData)
   }
 
   /// A description of what this Fix-It performs.
@@ -89,6 +130,12 @@ private extension FixIt.Change {
         range: token.endPositionBeforeTrailingTrivia..<token.endPosition,
         replacement: newTrivia.description
       )
+
+    case .replaceChild(let replacingChildData):
+      return SourceEdit(
+        range: replacingChildData.replacementRange,
+        replacement: replacingChildData.newChild.description
+      )
     }
   }
 }
diff --git a/Sources/SwiftSyntaxMacrosGenericTestSupport/Assertions.swift b/Sources/SwiftSyntaxMacrosGenericTestSupport/Assertions.swift
index 2023ed6d55f..ae22cd5c898 100644
--- a/Sources/SwiftSyntaxMacrosGenericTestSupport/Assertions.swift
+++ b/Sources/SwiftSyntaxMacrosGenericTestSupport/Assertions.swift
@@ -622,6 +622,15 @@ fileprivate extension FixIt.Change {
         range: start..<end,
         replacement: newTrivia.description
       )
+
+    case .replaceChild(let replacingChildData):
+      let range = replacingChildData.replacementRange
+      let start = expansionContext.position(of: range.lowerBound, anchoredAt: replacingChildData.parent)
+      let end = expansionContext.position(of: range.upperBound, anchoredAt: replacingChildData.parent)
+      return SourceEdit(
+        range: start..<end,
+        replacement: replacingChildData.newChild.description
+      )
     }
   }
 }
diff --git a/Tests/SwiftSyntaxMacroExpansionTest/PeerMacroTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/PeerMacroTests.swift
index f5a5ad812cf..48cea8d2027 100644
--- a/Tests/SwiftSyntaxMacroExpansionTest/PeerMacroTests.swift
+++ b/Tests/SwiftSyntaxMacroExpansionTest/PeerMacroTests.swift
@@ -49,24 +49,17 @@ final class PeerMacroTests: XCTestCase {
           newEffects = FunctionEffectSpecifiersSyntax(asyncSpecifier: .keyword(.async))
         }
 
-        let newSignature = funcDecl.signature.with(\.effectSpecifiers, newEffects)
-
         let diag = Diagnostic(
           node: Syntax(funcDecl.funcKeyword),
           message: SwiftSyntaxMacros.MacroExpansionErrorMessage(
             "can only add a completion-handler variant to an 'async' function"
           ),
           fixIts: [
-            FixIt(
-              message: SwiftSyntaxMacros.MacroExpansionFixItMessage(
-                "add 'async'"
-              ),
-              changes: [
-                FixIt.Change.replace(
-                  oldNode: Syntax(funcDecl.signature),
-                  newNode: Syntax(newSignature)
-                )
-              ]
+            .replaceChild(
+              message: SwiftSyntaxMacros.MacroExpansionFixItMessage("add 'async'"),
+              parent: funcDecl,
+              replacingChildAt: \.signature.effectSpecifiers,
+              with: newEffects
             )
           ]
         )