Add fixed source to assertMacroExpansion

Author: kimdv

Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift

@@ -254,13 +254,18 @@ func assertDiagnostic(
 ///   - macros: The macros that should be expanded, provided as a dictionary
 ///     mapping macro names (e.g., `"stringify"`) to implementation types
 ///     (e.g., `StringifyMacro.self`).
+///   - applyFixIts: If specified, filters the Fix-Its that are applied to generate `fixedSource` to only those whose message occurs in this array. If `nil`, all Fix-Its from the diagnostics are applied.
+///   - fixedSource: If specified, asserts that the source code after applying Fix-Its matches this string.
 ///   - testModuleName: The name of the test module to use.
 ///   - testFileName: The name of the test file name to use.
+///   - indentationWidth: The indentation width used in the expansion.
 public func assertMacroExpansion(
   _ originalSource: String,
   expandedSource expectedExpandedSource: String,
   diagnostics: [DiagnosticSpec] = [],
   macros: [String: Macro.Type],
+  applyFixIts: [String]? = nil,
+  fixedSource expectedFixedSource: String? = nil,
   testModuleName: String = "TestModule",
   testFileName: String = "test.swift",
   indentationWidth: Trivia = .spaces(4),
@@ -317,4 +322,16 @@ public func assertMacroExpansion(
       assertDiagnostic(actualDiag, in: context, expected: expectedDiag)
     }
   }
+
+  // Applying Fix-Its
+  if let expectedFixedSource = expectedFixedSource {
+    let fixedTree = FixItApplier.applyFixes(in: context.diagnostics, withMessages: applyFixIts, to: origSourceFile)
+    let fixedTreeDescription = fixedTree.description
+    assertStringsEqualWithDiff(
+      fixedTreeDescription.trimmingTrailingWhitespace(),
+      expectedFixedSource.trimmingTrailingWhitespace(),
+      file: file,
+      line: line
+    )
+  }
 }

Sources/_SwiftSyntaxTestSupport/FixItApplier.swift

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+import SwiftDiagnostics
+
+public class FixItApplier: SyntaxRewriter {
+  fileprivate struct Edit {
+    let startUtf8Offset: Int
+    let endUtf8Offset: Int
+    let replacement: String
+  }
+
+  var changes: [FixIt.Change]
+
+  private var locationConverter: SourceLocationConverter?
+  private var edits: [Edit] = []
+
+  init(diagnostics: [Diagnostic], withMessages messages: [String]?) {
+    let messages = messages ?? diagnostics.compactMap { $0.fixIts.first?.message.message }
+
+    self.changes =
+      diagnostics
+      .flatMap { $0.fixIts }
+      .filter {
+        return messages.contains($0.message.message)
+      }
+      .flatMap { $0.changes }
+
+    super.init(viewMode: .all)
+  }
+
+  public override func visitPre(_ node: Syntax) {
+    if locationConverter == nil && !node.hasParent {
+      locationConverter = SourceLocationConverter(fileName: "", tree: node)
+
+      for change in self.changes {
+        switch change {
+        case .replace(oldNode: let oldNode, newNode: let newNode):
+          let oldStartLocation = oldNode.startLocation(converter: locationConverter!)
+          let oldEndLocation = oldNode.startLocation(converter: locationConverter!)
+
+          edits.append(Edit(startUtf8Offset: oldStartLocation.offset, endUtf8Offset: oldEndLocation.offset, replacement: newNode.description))
+
+        default:
+          break
+        }
+      }
+    }
+  }
+
+  public override func visitAny(_ node: Syntax) -> Syntax? {
+    return nil
+  }
+
+  public override func visit(_ node: TokenSyntax) -> TokenSyntax {
+    var modifiedNode = node
+    for change in changes {
+      switch change {
+      case .replaceLeadingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.with(\.leadingTrivia, newTrivia)
+      case .replaceTrailingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.with(\.trailingTrivia, newTrivia)
+      default:
+        break
+      }
+    }
+    return modifiedNode
+  }
+
+  /// If `messages` is `nil`, applies all Fix-Its in `diagnostics` to `tree` and returns the fixed syntax tree.
+  /// If `messages` is not `nil`, applies only Fix-Its whose message is in `messages`.
+  public static func applyFixes<T: SyntaxProtocol>(in diagnostics: [Diagnostic], withMessages messages: [String]?, to tree: T) -> Syntax {
+    let applier = FixItApplier(diagnostics: diagnostics, withMessages: messages)
+    return applier.rewrite(tree)
+  }
+}

Tests/SwiftParserTest/Assertions.swift

@@ -276,58 +276,6 @@ struct DiagnosticSpec {
   }
 }
 
-class FixItApplier: SyntaxRewriter {
-  var changes: [FixIt.Change]
-
-  init(diagnostics: [Diagnostic], withMessages messages: [String]?) {
-    let messages = messages ?? diagnostics.compactMap { $0.fixIts.first?.message.message }
-
-    self.changes =
-      diagnostics
-      .flatMap { $0.fixIts }
-      .filter {
-        return messages.contains($0.message.message)
-      }
-      .flatMap { $0.changes }
-
-    super.init(viewMode: .all)
-  }
-
-  public override func visitAny(_ node: Syntax) -> Syntax? {
-    for change in changes {
-      switch change {
-      case .replace(oldNode: let oldNode, newNode: let newNode) where oldNode.id == node.id:
-        return newNode
-      default:
-        break
-      }
-    }
-    return nil
-  }
-
-  override func visit(_ node: TokenSyntax) -> TokenSyntax {
-    var modifiedNode = node
-    for change in changes {
-      switch change {
-      case .replaceLeadingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
-        modifiedNode = node.with(\.leadingTrivia, newTrivia)
-      case .replaceTrailingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
-        modifiedNode = node.with(\.trailingTrivia, newTrivia)
-      default:
-        break
-      }
-    }
-    return modifiedNode
-  }
-
-  /// If `messages` is `nil`, applies all Fix-Its in `diagnostics` to `tree` and returns the fixed syntax tree.
-  /// If `messages` is not `nil`, applies only Fix-Its whose message is in `messages`.
-  public static func applyFixes<T: SyntaxProtocol>(in diagnostics: [Diagnostic], withMessages messages: [String]?, to tree: T) -> Syntax {
-    let applier = FixItApplier(diagnostics: diagnostics, withMessages: messages)
-    return applier.rewrite(tree)
-  }
-}
-
 /// Assert that `location` is the same as that of `locationMarker` in `tree`.
 func assertLocation<T: SyntaxProtocol>(
   _ location: SourceLocation,