[Diagnostics] fix trivia merging and transfer

Sources/SwiftParserDiagnostics/DiagnosticExtensions.swift

@@ -52,13 +52,39 @@ extension FixIt {
 
 // MARK: - Make missing
 
+private extension FixIt.Change {
+  /// Transfers the leading and trailing trivia of `nodes` to the previous token
+  /// While doing this, it tries to be smart, merging trivia where it makes sense
+  /// and refusing to add e.g. a space after punctuation, where it usually
+  /// doesn't make sense.
+  static func transferTriviaAtSides(from nodes: [some SyntaxProtocol]) -> Self? {
+    guard let first = nodes.first, let last = nodes.last else {
+      preconditionFailure("nodes must not be empty")
+    }
+    let removedTriviaAtSides = first.leadingTrivia.mergingCommonSuffix(last.trailingTrivia)
+    guard !removedTriviaAtSides.isEmpty, let previousToken = first.previousToken(viewMode: .sourceAccurate) else {
+      return nil
+    }
+    let mergedTrivia = previousToken.trailingTrivia.mergingCommonPrefix(removedTriviaAtSides)
+    // We merge with the common prefix instead of the common suffix to preserve the original shape of the previous
+    // token's trailing trivia after merging.
+    guard !previousToken.tokenKind.isPunctuation || !mergedTrivia.allSatisfy(\.isSpaceOrTab) else {
+      // Punctuation is generally not followed by spaces in Swift.
+      // If this action would only add spaces to the punctuation, drop it.
+      // This generally yields better results.
+      return nil
+    }
+    return .replaceTrailingTrivia(token: previousToken, newTrivia: mergedTrivia)
+  }
+}
+
 extension FixIt.MultiNodeChange {
-  /// Replaced a present token with a missing node.
+  /// Replaced a present token with a missing token.
   ///
   /// If `transferTrivia` is `true`, the leading and trailing trivia of the
-  /// removed node will be transferred to the trailing trivia of the previous token.
+  /// removed token will be transferred to the trailing trivia of the previous token.
   static func makeMissing(_ token: TokenSyntax, transferTrivia: Bool = true) -> Self {
-    return makeMissing([token], transferTrivia: transferTrivia)
+    self.makeMissing([token], transferTrivia: transferTrivia)
   }
 
   /// Replace present tokens with missing tokens.
@@ -68,57 +94,33 @@ extension FixIt.MultiNodeChange {
   /// tokens.
   static func makeMissing(_ tokens: [TokenSyntax], transferTrivia: Bool = true) -> Self {
     precondition(tokens.allSatisfy(\.isPresent))
-    return .makeMissing(tokens.map(Syntax.init), transferTrivia: transferTrivia)
+    return self.makeMissing(tokens.map(Syntax.init), transferTrivia: transferTrivia)
   }
 
   /// If `transferTrivia` is `true`, the leading and trailing trivia of the
   /// removed node will be transferred to the trailing trivia of the previous token.
   static func makeMissing(_ node: (some SyntaxProtocol)?, transferTrivia: Bool = true) -> Self {
-    guard let node = node else {
+    guard let node else {
       return FixIt.MultiNodeChange(primitiveChanges: [])
     }
-    var changes = [FixIt.Change.replace(oldNode: Syntax(node), newNode: MissingMaker().rewrite(node, detach: true))]
-    if transferTrivia {
-      changes += FixIt.MultiNodeChange.transferTriviaAtSides(from: [node]).primitiveChanges
-    }
-    return FixIt.MultiNodeChange(primitiveChanges: changes)
-  }
-
-  /// Transfers the leading and trailing trivia of `nodes` to the previous token
-  /// While doing this, it tries to be smart, merging trivia where it makes sense
-  /// and refusing to add e.g. a space after punctuation, where it usually
-  /// doesn't make sense.
-  private static func transferTriviaAtSides(from nodes: [some SyntaxProtocol]) -> Self {
-    let removedTriviaAtSides = (nodes.first?.leadingTrivia ?? []).merging(nodes.last?.trailingTrivia ?? [])
-    if !removedTriviaAtSides.isEmpty, let previousToken = nodes.first?.previousToken(viewMode: .sourceAccurate) {
-      let mergedTrivia = previousToken.trailingTrivia.merging(removedTriviaAtSides)
-      if previousToken.tokenKind.isPunctuation, mergedTrivia.allSatisfy({ $0.isSpaceOrTab }) {
-        // Punctuation is generally not followed by spaces in Swift.
-        // If this action would only add spaces to the punctuation, drop it.
-        // This generally yields better results.
-        return FixIt.MultiNodeChange()
-      }
-      return FixIt.MultiNodeChange(.replaceTrailingTrivia(token: previousToken, newTrivia: mergedTrivia))
-    } else {
-      return FixIt.MultiNodeChange()
-    }
+    return self.makeMissing([node], transferTrivia: transferTrivia)
   }
 
-  /// Replace present nodes with their missing equivalents.
+  /// Replace present nodes with missing nodes.
   ///
   /// If `transferTrivia` is `true`, the leading trivia of the first node and
   /// the trailing trivia of the last node will be transferred to their adjecent
   /// tokens.
-  static func makeMissing(_ nodes: [Syntax], transferTrivia: Bool = true) -> Self {
+  static func makeMissing(_ nodes: [some SyntaxProtocol], transferTrivia: Bool = true) -> Self {
     precondition(!nodes.isEmpty)
     var changes = nodes.map {
       FixIt.Change.replace(
-        oldNode: $0,
+        oldNode: Syntax($0),
         newNode: MissingMaker().rewrite($0, detach: true)
       )
     }
-    if transferTrivia {
-      changes += FixIt.MultiNodeChange.transferTriviaAtSides(from: nodes).primitiveChanges
+    if transferTrivia, let transferredTrivia = FixIt.Change.transferTriviaAtSides(from: nodes) {
+      changes.append(transferredTrivia)
     }
     return FixIt.MultiNodeChange(primitiveChanges: changes)
   }

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -213,7 +213,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
       )
       changes.append(FixIt.MultiNodeChange.makeMissing(misplacedTokens, transferTrivia: false))
     } else {
-      changes += misplacedTokens.map { FixIt.MultiNodeChange.makeMissing($0) }
+      changes += [.makeMissing(misplacedTokens)]
       changes += correctAndMissingNodes.map { FixIt.MultiNodeChange.makePresent($0) }
     }
     var fixIts: [FixIt] = []

Sources/SwiftRefactor/CallToTrailingClosures.swift

@@ -78,7 +78,7 @@ extension FunctionCallExprSyntax {
       // Trailing comma won't exist any more, move its trivia to the end of
       // the closure instead
       if let comma = arg.trailingComma {
-        closure.trailingTrivia = closure.trailingTrivia.merging(triviaOf: comma)
+        closure.trailingTrivia = closure.trailingTrivia.mergingCommonSuffix(triviaOf: comma)
       }
       closures.append((arg, closure))
     }
@@ -88,12 +88,12 @@ extension FunctionCallExprSyntax {
     }
 
     // First trailing closure won't have label/colon. Transfer their trivia.
-    var trailingClosure = closures.first!.closure
+    let (firstOriginal, firstClosure) = closures.first!
+    var trailingClosure = firstClosure
     trailingClosure.leadingTrivia =
-      Trivia()
-      .merging(triviaOf: closures.first!.original.label)
-      .merging(triviaOf: closures.first!.original.colon)
-      .merging(closures.first!.closure.leadingTrivia)
+      (firstOriginal.label?.triviaByMergingCommonSuffix ?? [])
+      .mergingCommonSuffix(triviaOf: firstOriginal.colon)
+      .mergingCommonSuffix(firstClosure.leadingTrivia)
       .droppingLeadingWhitespace
     let additionalTrailingClosures = closures.dropFirst().map {
       MultipleTrailingClosureElementSyntax(
@@ -118,21 +118,20 @@ extension FunctionCallExprSyntax {
       // No left paren any more, right paren is handled below since it makes
       // sense to keep its trivia of the end of the call, regardless of whether
       // it was removed or not.
-      if let leftParen = leftParen {
-        trailingClosure.leadingTrivia = Trivia()
-          .merging(triviaOf: leftParen)
-          .merging(trailingClosure.leadingTrivia)
+      if let leftParen {
+        trailingClosure.leadingTrivia = leftParen.triviaByMergingCommonSuffix
+          .mergingCommonSuffix(trailingClosure.leadingTrivia)
       }
       // No right paren anymore. Attach its trivia to the end of the call.
-      if let rightParen = rightParen {
-        additionalTriviaAtEndOfCall = Trivia().merging(triviaOf: rightParen)
+      if let rightParen {
+        additionalTriviaAtEndOfCall = rightParen.triviaByMergingCommonSuffix
       }
     } else {
       let last = argList.last!
       // Move the trailing trivia of the closing parenthesis to the end of the call after the last trailing, instead of
       // keeping it in the middle of the call where the new closing parenthesis lives.
       // Also ensure that we don't drop trivia from any comma we remove.
-      converted.rightParen?.trailingTrivia = Trivia().merging(triviaOf: last.trailingComma)
+      converted.rightParen?.trailingTrivia = last.trailingComma?.triviaByMergingCommonSuffix ?? []
       additionalTriviaAtEndOfCall = rightParen?.trailingTrivia
       argList[argList.count - 1] = last.with(\.trailingComma, nil)
     }
@@ -145,7 +144,9 @@ extension FunctionCallExprSyntax {
     }
 
     if let additionalTriviaAtEndOfCall {
-      converted.trailingTrivia = converted.trailingTrivia.merging(additionalTriviaAtEndOfCall.droppingLeadingWhitespace)
+      converted.trailingTrivia = converted.trailingTrivia.mergingCommonSuffix(
+        additionalTriviaAtEndOfCall.droppingLeadingWhitespace
+      )
     }
 
     return converted

Sources/SwiftSyntax/Trivia.swift

@@ -69,6 +69,7 @@ public struct Trivia: Sendable {
 
   /// Creates a new ``Trivia`` by merging in the given trivia. Only includes one
   /// copy of a common prefix of `self` and `trivia`.
+  @available(*, deprecated, message: "Use mergingCommonPrefix(trivia) or mergingCommonSuffix(trivia) instead")
   public func merging(_ trivia: Trivia?) -> Trivia {
     guard let trivia else {
       return self
@@ -84,16 +85,69 @@ public struct Trivia: Sendable {
     return lhs.appending(rhs)
   }
 
+  /// Creates a new ``Trivia`` by merging in the given trivia. Only includes one
+  /// copy of the common prefix of `self` and `trivia`.
+  public func mergingCommonPrefix(_ trivia: Trivia?) -> Trivia {
+    guard let trivia else {
+      return self
+    }
+
+    let lhs = self.decomposed
+    let rhs = trivia.decomposed
+    let commonPrefix = zip(lhs, rhs).prefix(while: ==)
+    if commonPrefix.isEmpty {
+      return lhs + rhs
+    } else {
+      return lhs + Trivia(pieces: rhs.dropFirst(commonPrefix.count))
+    }
+  }
+
+  /// Creates a new ``Trivia`` by merging in the given trivia. Only includes one
+  /// copy of the common suffix of `self` and `trivia`.
+  public func mergingCommonSuffix(_ trivia: Trivia?) -> Trivia {
+    guard let trivia else {
+      return self
+    }
+
+    let lhs = self.decomposed
+    let rhs = trivia.decomposed
+    let commonSuffix = zip(lhs.reversed(), rhs.reversed()).prefix(while: ==)
+    if commonSuffix.isEmpty {
+      return lhs + rhs
+    } else {
+      return Trivia(pieces: lhs.dropLast(commonSuffix.count)) + rhs
+    }
+  }
+
   /// Creates a new ``Trivia`` by merging the leading and trailing ``Trivia``
   /// of `triviaOf` into the end of `self`. Only includes one copy of any
   /// common prefixes.
+  @available(*, deprecated, message: "Use mergingCommonPrefix(triviaOf:) or mergingCommonSuffix(triviaOf:) instead")
   public func merging(triviaOf node: (some SyntaxProtocol)?) -> Trivia {
     guard let node else {
       return self
     }
     return merging(node.leadingTrivia).merging(node.trailingTrivia)
   }
 
+  /// Creates a new ``Trivia`` by merging ``SyntaxProtocol/leadingTrivia`` and ``SyntaxProtocol/trailingTrivia`` of
+  /// `node` into the end of `self`. Only includes one copy of any common prefixes.
+  public func mergingCommonPrefix(triviaOf node: (some SyntaxProtocol)?) -> Trivia {
+    guard let node else {
+      return self
+    }
+    return self.mergingCommonPrefix(node.leadingTrivia).mergingCommonPrefix(node.trailingTrivia)
+  }
+
+  /// Creates a new ``Trivia`` by merging ``SyntaxProtocol/leadingTrivia`` and ``SyntaxProtocol/trailingTrivia`` of
+  /// `node` into the end of `self`. Only includes one copy of any common suffixes.
+  public func mergingCommonSuffix(triviaOf node: (some SyntaxProtocol)?) -> Trivia {
+    guard let node else {
+      return self
+    }
+    return self.mergingCommonSuffix(node.leadingTrivia).mergingCommonSuffix(node.trailingTrivia)
+  }
+
   /// Concatenates two collections of ``Trivia`` into one collection.
   public static func + (lhs: Trivia, rhs: Trivia) -> Trivia {
     return lhs.appending(rhs)
@@ -215,3 +269,17 @@ extension RawTriviaPiece: CustomDebugStringConvertible {
     TriviaPiece(raw: self).debugDescription
   }
 }
+
+public extension SyntaxProtocol {
+  /// Create a new ``Trivia`` by merging ``SyntaxProtocol/leadingTrivia`` and ``SyntaxProtocol/trailingTrivia`` of
+  /// this node, including only one copy of their common prefix.
+  var triviaByMergingCommonPrefix: Trivia {
+    self.leadingTrivia.mergingCommonPrefix(self.trailingTrivia)
+  }
+
+  /// Create a new ``Trivia`` by merging ``SyntaxProtocol/leadingTrivia`` and ``SyntaxProtocol/trailingTrivia`` of
+  /// this node, including only one copy of their common suffix.
+  var triviaByMergingCommonSuffix: Trivia {
+    self.leadingTrivia.mergingCommonSuffix(self.trailingTrivia)
+  }
+}

Tests/SwiftParserTest/DeclarationTests.swift

@@ -960,18 +960,18 @@ final class DeclarationTests: ParserTestCase {
           fixIts: ["move 'throws(any Error)' in front of '->'"]
         )
       ],
-      fixedSource: "func test() throws(any Error) ->  Int"
+      fixedSource: "func test() throws(any Error) -> Int"
     )
 
     assertParse(
-      "func test() -> 1️⃣throws(any Error Int",
+      "func test() -> 1️⃣throws(any Error/* */ Int",
       diagnostics: [
         DiagnosticSpec(
           message: "'throws(any Error' must precede '->'",
           fixIts: ["move 'throws(any Error' in front of '->'"]
         )
       ],
-      fixedSource: "func test() throws(any Error -> Int"
+      fixedSource: "func test() throws(any Error -> /* */ Int"
     )
 
     assertParse(
@@ -986,14 +986,14 @@ final class DeclarationTests: ParserTestCase {
     )
 
     assertParse(
-      "func test() -> 1️⃣throws (any Error) Int",
+      "func test() -> 1️⃣throws/**/ (any Error) Int",
       diagnostics: [
         DiagnosticSpec(
           message: "'throws' must precede '->'",
           fixIts: ["move 'throws' in front of '->'"]
         )
       ],
-      fixedSource: "func test() throws -> (any Error) Int"
+      fixedSource: "func test() throws -> /**/ (any Error) Int"
     )
   }
 

Tests/SwiftParserTest/translated/RecoveryTests.swift

@@ -774,7 +774,8 @@ final class RecoveryTests: ParserTestCase {
         ),
       ],
       fixedSource: """
-        for   <#pattern#> in <#expression#> {
+        for 
+          <#pattern#> in <#expression#> {
         }
         """
     )