When making a token missing, transfer its leading triva to the next token

This makes sure we don’t drop newlines or comments when removing tokens.

Author: ahoppen

Sources/SwiftDiagnostics/FixIt.swift

@@ -26,21 +26,39 @@ public protocol FixItMessage {
 
 /// A Fix-It that can be applied to resolve a diagnostic.
 public struct FixIt {
+  public struct Changes: ExpressibleByArrayLiteral {
+    public var changes: [Change]
+
+    public init(changes: [Change]) {
+      self.changes = changes
+    }
+
+    public init(arrayLiteral elements: FixIt.Change...) {
+      self.init(changes: elements)
+    }
+
+    public init(combining: [Changes]) {
+      self.init(changes: combining.flatMap(\.changes))
+    }
+  }
+
   public enum Change {
     /// Replace `oldNode` by `newNode`.
     case replace(oldNode: Syntax, newNode: Syntax)
-    /// Remove the trailing trivia of the given token
-    case removeTrailingTrivia(TokenSyntax)
+    /// 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)
   }
 
   /// A description of what this Fix-It performs.
   public let message: FixItMessage
 
   /// The changes that need to be performed when the Fix-It is applied.
-  public let changes: [Change]
+  public let changes: Changes
 
-  public init(message: FixItMessage, changes: [Change]) {
-    assert(!changes.isEmpty, "A Fix-It must have at least one change associated with it")
+  public init(message: FixItMessage, changes: Changes) {
+    assert(!changes.changes.isEmpty, "A Fix-It must have at least one change associated with it")
     self.message = message
     self.changes = changes
   }

Sources/SwiftParser/Diagnostics/DiagnosticExtensions.swift

@@ -14,42 +14,72 @@ import SwiftDiagnostics
 import SwiftSyntax
 
 extension FixIt {
-  init(message: StaticParserFixIt, changes: [Change]) {
+  init(message: StaticParserFixIt, changes: Changes) {
     self.init(message: message as FixItMessage, changes: changes)
   }
+
+  public init(message: FixItMessage, changes: [Changes]) {
+    self.init(message: message, changes: FixIt.Changes(combining: changes))
+  }
+
+  init(message: StaticParserFixIt, changes: [Changes]) {
+    self.init(message: message as FixItMessage, changes: FixIt.Changes(combining: changes))
+  }
 }
 
-extension FixIt.Change {
+extension FixIt.Changes {
   /// Replaced a present node with a missing node
-  static func makeMissing(node: TokenSyntax) -> FixIt.Change {
+  static func makeMissing(node: TokenSyntax) -> Self {
     assert(node.presence == .present)
-    return .replace(
-      oldNode: Syntax(node),
-      newNode: Syntax(TokenSyntax(node.tokenKind, leadingTrivia: [], trailingTrivia: [], presence: .missing))
-    )
+    var changes = [
+      FixIt.Change.replace(
+        oldNode: Syntax(node),
+        newNode: Syntax(TokenSyntax(node.tokenKind, leadingTrivia: [], trailingTrivia: [], presence: .missing))
+      )
+    ]
+    if !node.leadingTrivia.isEmpty, let nextToken = node.nextToken(viewMode: .sourceAccurate) {
+      changes.append(.replaceLeadingTrivia(token: nextToken, newTrivia: node.leadingTrivia))
+    }
+    return FixIt.Changes(changes: changes)
   }
 
-  static func makePresent<T: SyntaxProtocol>(node: T) -> FixIt.Change {
-    return .replace(
+  static func makePresent<T: SyntaxProtocol>(node: T) -> Self {
+    return [.replace(
       oldNode: Syntax(node),
       newNode: PresentMaker().visit(Syntax(node))
-    )
+    )]
+  }
+
+  /// Make a token present. If `leadingTrivia` or `trailingTrivia` is specified,
+  /// override the default leading/trailing trivia inferred from `BasicFormat`.
+  static func makePresent(
+    node: TokenSyntax,
+    leadingTrivia: Trivia? = nil,
+    trailingTrivia: Trivia? = nil
+  ) -> Self {
+    var presentNode = PresentMaker().visit(Syntax(node)).as(TokenSyntax.self)!
+    if let leadingTrivia = leadingTrivia {
+      presentNode.leadingTrivia = leadingTrivia
+    }
+    if let trailingTrivia = trailingTrivia {
+      presentNode.trailingTrivia = trailingTrivia
+    }
+    return [.replace(
+      oldNode: Syntax(node),
+      newNode: Syntax(presentNode)
+    )]
   }
-}
 
-extension Array where Element == FixIt.Change {
   /// Makes the `token` present, moving it in front of the previous token's trivia.
-  static func makePresentBeforeTrivia(token: TokenSyntax) -> [FixIt.Change] {
+  static func makePresentBeforeTrivia(token: TokenSyntax) -> Self {
     if let previousToken = token.previousToken(viewMode: .sourceAccurate) {
       let presentToken = PresentMaker().visit(token).withTrailingTrivia(previousToken.trailingTrivia)
       return [
-        .removeTrailingTrivia(previousToken),
+        .replaceTrailingTrivia(token: previousToken, newTrivia: []),
         .replace(oldNode: Syntax(token), newNode: presentToken),
       ]
     } else {
-      return [
-        .makePresent(node: token)
-      ]
+      return .makePresent(node: token)
     }
   }
 }

Sources/SwiftParser/Diagnostics/MissingNodesError.swift

@@ -256,7 +256,7 @@ extension ParseDiagnosticsGenerator {
 
     let fixIt = FixIt(
       message: InsertTokenFixIt(missingNodes: missingNodes),
-      changes: missingNodes.map(FixIt.Change.makePresent)
+      changes: missingNodes.map(FixIt.Changes.makePresent)
     )
 
     var notes: [Note] = []

Tests/SwiftParserTest/Assertions.swift

@@ -120,7 +120,7 @@ class FixItApplier: SyntaxRewriter {
   var changes: [FixIt.Change]
 
   init(diagnostics: [Diagnostic]) {
-    self.changes = diagnostics.flatMap({ $0.fixIts }).flatMap({ $0.changes })
+    self.changes = diagnostics.flatMap({ $0.fixIts }).flatMap({ $0.changes.changes })
   }
 
   public override func visitAny(_ node: Syntax) -> Syntax? {
@@ -136,15 +136,18 @@ class FixItApplier: SyntaxRewriter {
   }
 
   override func visit(_ node: TokenSyntax) -> Syntax {
+    var modifiedNode = node
     for change in changes {
       switch change {
-      case .removeTrailingTrivia(let changeNode) where changeNode.id == node.id:
-        return Syntax(node.withTrailingTrivia([]))
+      case .replaceLeadingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.withLeadingTrivia(newTrivia)
+      case .replaceTrailingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.withTrailingTrivia(newTrivia)
       default:
         break
       }
     }
-    return Syntax(node)
+    return Syntax(modifiedNode)
   }
 
   /// Applies all Fix-Its in `diagnostics` to `tree` and returns the fixed syntax tree.