Replace the ReplaceTrivia rewriter with direct trivia mutations.

Previously, we were using a `SyntaxRewriter` to replace trivia on
nodes. Historically, this originated from a time (I believe) when
trivia wasn't directly mutable on nodes. The vast majority of these
uses were replacing the extents: the leading trivia of the first
token of a node or the trailing trivia of the last token. This can
be done more easily now by just mutating those properties of the
nodes.

Even in cases where we were doing something slightly more advanced,
it's clearer to do in-place mutation and reconstruction; having
this specialized rewriter no longer held its weight.

Author: allevato

Sources/SwiftFormatRules/AddModifierRewriter.swift

@@ -24,8 +24,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.bindingSpecifier)
       return DeclSyntax(result)
     }
     // If variable already has an accessor keyword, skip (do not overwrite)
@@ -40,8 +39,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.funcKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -53,8 +51,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.associatedtypeKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -66,8 +63,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.classKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -79,8 +75,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.enumKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -92,8 +87,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.protocolKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -105,8 +99,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.structKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -118,8 +111,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.typealiasKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -131,8 +123,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.initKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -144,8 +135,7 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
     // Check for modifiers, and, if none, insert the modifier and relocate trivia from the displaced
     // token.
     guard var modifiers = node.modifiers else {
-      let nodeWithModifier = node.with(\.modifiers, .init([modifierKeyword]))
-      let result = nodeByRelocatingTrivia(in: nodeWithModifier) { $0.modifiers }
+      let result = setOnlyModifier(in: node, keywordKeypath: \.subscriptKeyword)
       return DeclSyntax(result)
     }
     guard modifiers.accessLevelModifier == nil else { return DeclSyntax(node) }
@@ -161,24 +151,16 @@ fileprivate final class AddModifierRewriter: SyntaxRewriter {
   /// this method does nothing and returns the given node as-is.
   /// - Parameter node: A node that was updated to include a new modifier.
   /// - Parameter modifiersProvider: A closure that returns all modifiers for the given node.
-  private func nodeByRelocatingTrivia<NodeType: DeclSyntaxProtocol>(
+  private func setOnlyModifier<NodeType: DeclSyntaxProtocol & WithModifiersSyntax>(
     in node: NodeType,
-    for modifiersProvider: (NodeType) -> DeclModifierListSyntax?
+    keywordKeypath: WritableKeyPath<NodeType, TokenSyntax>
   ) -> NodeType {
-    guard let modifier = modifiersProvider(node)?.firstAndOnly,
-      let movingLeadingTrivia = modifier.nextToken(viewMode: .sourceAccurate)?.leadingTrivia
-    else {
-      // Otherwise, there's no trivia that needs to be relocated so the node is fine.
-      return node
-    }
-    let nodeWithTrivia = replaceTrivia(
-      on: node,
-      token: modifier.firstToken(viewMode: .sourceAccurate),
-      leadingTrivia: movingLeadingTrivia)
-    return replaceTrivia(
-      on: nodeWithTrivia,
-      token: modifiersProvider(nodeWithTrivia)?.first?.nextToken(viewMode: .sourceAccurate),
-      leadingTrivia: [])
+    var node = node
+    var modifier = modifierKeyword
+    modifier.leadingTrivia = node[keyPath: keywordKeypath].leadingTrivia
+    node[keyPath: keywordKeypath].leadingTrivia = []
+    node.modifiers = .init([modifier])
+    return node
   }
 }
 

Sources/SwiftFormatRules/DoNotUseSemicolons.swift

@@ -59,17 +59,11 @@ public final class DoNotUseSemicolons: SyntaxFormatRule {
       if previousHadSemicolon, let firstToken = newItem.firstToken(viewMode: .sourceAccurate),
         !firstToken.leadingTrivia.containsNewlines
       {
-        let leadingTrivia = .newlines(1) + firstToken.leadingTrivia
-        newItem = replaceTrivia(
-          on: newItem,
-          token: firstToken,
-          leadingTrivia: leadingTrivia
-        )
+        newItem.leadingTrivia = .newlines(1) + firstToken.leadingTrivia
       }
 
       // If there's a semicolon, diagnose and remove it.
       if let semicolon = item.semicolon {
-
         // Exception: do not remove the semicolon if it is separating a 'do' statement from a
         // 'while' statement.
         if Syntax(item).as(CodeBlockItemSyntax.self)?

Sources/SwiftFormatRules/FullyIndirectEnum.swift

@@ -45,8 +45,7 @@ public final class FullyIndirectEnum: SyntaxFormatRule {
       }
 
       let newCase = caseMember.with(\.modifiers, modifiers.remove(name: "indirect"))
-      let formattedCase = formatCase(
-        unformattedCase: newCase, leadingTrivia: firstModifier.leadingTrivia)
+      let formattedCase = rearrangeLeadingTrivia(firstModifier.leadingTrivia, on: newCase)
       return member.with(\.decl, DeclSyntax(formattedCase))
     }
 
@@ -55,15 +54,13 @@ public final class FullyIndirectEnum: SyntaxFormatRule {
     // line breaks/comments/indentation.
     let firstTok = node.firstToken(viewMode: .sourceAccurate)!
     let leadingTrivia: Trivia
-    let newEnumDecl: EnumDeclSyntax
+    var newEnumDecl = node
 
     if firstTok.tokenKind == .keyword(.enum) {
       leadingTrivia = firstTok.leadingTrivia
-      newEnumDecl = replaceTrivia(
-        on: node, token: node.firstToken(viewMode: .sourceAccurate), leadingTrivia: [])
+      newEnumDecl.leadingTrivia = []
     } else {
       leadingTrivia = []
-      newEnumDecl = node
     }
 
     let newModifier = DeclModifierSyntax(
@@ -94,19 +91,23 @@ public final class FullyIndirectEnum: SyntaxFormatRule {
   }
 
   /// Transfers given leading trivia to the first token in the case declaration.
-  private func formatCase(
-    unformattedCase: EnumCaseDeclSyntax,
-    leadingTrivia: Trivia?
+  private func rearrangeLeadingTrivia(
+    _ leadingTrivia: Trivia,
+    on enumCaseDecl: EnumCaseDeclSyntax
   ) -> EnumCaseDeclSyntax {
-    if let modifiers = unformattedCase.modifiers, let first = modifiers.first {
-      return replaceTrivia(
-        on: unformattedCase, token: first.firstToken(viewMode: .sourceAccurate), leadingTrivia: leadingTrivia
-      )
+    var formattedCase = enumCaseDecl
+
+    if var modifiers = formattedCase.modifiers, var firstModifier = modifiers.first {
+      // If the case has modifiers, attach the leading trivia to the first one.
+      firstModifier.leadingTrivia = leadingTrivia
+      modifiers[modifiers.startIndex] = firstModifier
+      formattedCase.modifiers = modifiers
     } else {
-      return replaceTrivia(
-        on: unformattedCase, token: unformattedCase.caseKeyword, leadingTrivia: leadingTrivia
-      )
+      // Otherwise, attach the trivia to the `case` keyword itself.
+      formattedCase.caseKeyword.leadingTrivia = leadingTrivia
     }
+
+    return formattedCase
   }
 }
 

Sources/SwiftFormatRules/ModifierListSyntax+Convenience.swift

@@ -54,29 +54,22 @@ extension DeclModifierListSyntax {
   mutating func triviaPreservingInsert(
     _ modifier: DeclModifierSyntax, at index: SyntaxChildrenIndex
   ) {
-    let modifier = replaceTrivia(
-      on: modifier,
-      token: modifier.name,
-      trailingTrivia: .spaces(1))
+    var modifier = modifier
+    modifier.trailingTrivia = [.spaces(1)]
 
     guard index == self.startIndex else {
       self.insert(modifier, at: index)
       return
     }
-    guard let firstMod = first, let firstTok = firstMod.firstToken(viewMode: .sourceAccurate) else {
+    guard var firstMod = first, let firstTok = firstMod.firstToken(viewMode: .sourceAccurate) else {
       self.insert(modifier, at: index)
       return
     }
 
-    let formattedMod = replaceTrivia(
-      on: modifier,
-      token: modifier.firstToken(viewMode: .sourceAccurate),
-      leadingTrivia: firstTok.leadingTrivia)
-    self[self.startIndex] = replaceTrivia(
-      on: firstMod,
-      token: firstTok,
-      leadingTrivia: [],
-      trailingTrivia: .spaces(1))
-    self.insert(formattedMod, at: self.startIndex)
+    modifier.leadingTrivia = firstTok.leadingTrivia
+    firstMod.leadingTrivia = []
+    firstMod.trailingTrivia = [.spaces(1)]
+    self[self.startIndex] = firstMod
+    self.insert(modifier, at: self.startIndex)
   }
 }

Sources/SwiftFormatRules/NoAccessLevelOnExtensionDeclaration.swift

@@ -48,10 +48,8 @@ public final class NoAccessLevelOnExtensionDeclaration: SyntaxFormatRule {
         leftBrace: node.memberBlock.leftBrace,
         members: addMemberAccessKeywords(memDeclBlock: node.memberBlock, keyword: accessKeywordToAdd),
         rightBrace: node.memberBlock.rightBrace)
-      let newKeyword = replaceTrivia(
-        on: node.extensionKeyword,
-        token: node.extensionKeyword,
-        leadingTrivia: accessKeyword.leadingTrivia)
+      var newKeyword = node.extensionKeyword
+      newKeyword.leadingTrivia = accessKeyword.leadingTrivia
       let result = node.with(\.memberBlock, newMembers)
         .with(\.modifiers, modifiers.remove(name: accessKeyword.name.text))
         .with(\.extensionKeyword, newKeyword)
@@ -62,10 +60,8 @@ public final class NoAccessLevelOnExtensionDeclaration: SyntaxFormatRule {
       diagnose(
         .removeRedundantAccessKeyword(name: node.extendedType.description),
         on: accessKeyword)
-      let newKeyword = replaceTrivia(
-        on: node.extensionKeyword,
-        token: node.extensionKeyword,
-        leadingTrivia: accessKeyword.leadingTrivia)
+      var newKeyword = node.extensionKeyword
+      newKeyword.leadingTrivia = accessKeyword.leadingTrivia
       let result = node.with(\.modifiers, modifiers.remove(name: accessKeyword.name.text))
         .with(\.extensionKeyword, newKeyword)
       return DeclSyntax(result)
@@ -82,10 +78,9 @@ public final class NoAccessLevelOnExtensionDeclaration: SyntaxFormatRule {
     keyword: DeclModifierSyntax
   ) -> MemberBlockItemListSyntax {
     var newMembers: [MemberBlockItemSyntax] = []
-    let formattedKeyword = replaceTrivia(
-      on: keyword,
-      token: keyword.name,
-      leadingTrivia: [])
+
+    var formattedKeyword = keyword
+    formattedKeyword.leadingTrivia = []
 
     for memberItem in memDeclBlock.members {
       let member = memberItem.decl

Sources/SwiftFormatRules/NoEmptyTrailingClosureParentheses.swift

@@ -29,24 +29,24 @@ public final class NoEmptyTrailingClosureParentheses: SyntaxFormatRule {
     {
       return super.visit(node)
     }
-    guard let name = node.calledExpression.lastToken(viewMode: .sourceAccurate)?.with(\.leadingTrivia, []).with(\.trailingTrivia, []) else {
+    guard let name = node.calledExpression.lastToken(viewMode: .sourceAccurate) else {
       return super.visit(node)
     }
 
-    diagnose(.removeEmptyTrailingParentheses(name: "\(name)"), on: node)
+    diagnose(.removeEmptyTrailingParentheses(name: "\(name.trimmedDescription)"), on: node)
 
     // Need to visit `calledExpression` before creating a new node so that the location data (column
     // and line numbers) is available.
-    guard let rewrittenCalledExpr = ExprSyntax(rewrite(Syntax(node.calledExpression))) else {
+    guard var rewrittenCalledExpr = ExprSyntax(rewrite(Syntax(node.calledExpression))) else {
       return super.visit(node)
     }
-    let formattedExp = replaceTrivia(
-      on: rewrittenCalledExpr,
-      token: rewrittenCalledExpr.lastToken(viewMode: .sourceAccurate),
-      trailingTrivia: .spaces(1))
-    let formattedClosure = visit(trailingClosure).as(ClosureExprSyntax.self)
-    let result = node.with(\.leftParen, nil).with(\.rightParen, nil).with(\.calledExpression, formattedExp)
-      .with(\.trailingClosure, formattedClosure)
+    rewrittenCalledExpr.trailingTrivia = [.spaces(1)]
+
+    var result = node
+    result.leftParen = nil
+    result.rightParen = nil
+    result.calledExpression = rewrittenCalledExpr
+    result.trailingClosure = rewrite(trailingClosure).as(ClosureExprSyntax.self)
     return ExprSyntax(result)
   }
 }

Sources/SwiftFormatRules/NoParensAroundConditions.swift

@@ -46,16 +46,13 @@ public final class NoParensAroundConditions: SyntaxFormatRule {
 
     guard
       let visitedTuple = visit(tuple).as(TupleExprSyntax.self),
-      let visitedExpr = visitedTuple.elements.first?.expression
+      var visitedExpr = visitedTuple.elements.first?.expression
     else {
       return expr
     }
-    return replaceTrivia(
-      on: visitedExpr,
-      token: visitedExpr.lastToken(viewMode: .sourceAccurate),
-      leadingTrivia: visitedTuple.leftParen.leadingTrivia,
-      trailingTrivia: visitedTuple.rightParen.trailingTrivia
-    )
+    visitedExpr.leadingTrivia = visitedTuple.leftParen.leadingTrivia
+    visitedExpr.trailingTrivia = visitedTuple.rightParen.trailingTrivia
+    return visitedExpr
   }
 
   public override func visit(_ node: IfExprSyntax) -> ExprSyntax {

Sources/SwiftFormatRules/OneVariableDeclarationPerLine.swift

@@ -163,8 +163,7 @@ private struct VariableDeclSplitter<Node: SyntaxProtocol> {
     // We intentionally don't try to infer the indentation for subsequent
     // lines because the pretty printer will re-indent them correctly; we just
     // need to ensure that a newline is inserted before new decls.
-    varDecl = replaceTrivia(
-      on: varDecl, token: varDecl.firstToken(viewMode: .sourceAccurate), leadingTrivia: .newlines(1))
+    varDecl.leadingTrivia = [.newlines(1)]
     fixedUpTrivia = true
   }
 

Sources/SwiftFormatRules/OrderedImports.swift

@@ -359,13 +359,10 @@ fileprivate func convertToCodeBlockItems(lines: [Line]) -> [CodeBlockItemSyntax]
     func append(codeBlockItem: CodeBlockItemSyntax) {
       // Comments and newlines are always located in the leading trivia of an AST node, so we need
       // not deal with trailing trivia.
-      output.append(
-        replaceTrivia(
-          on: codeBlockItem,
-          token: codeBlockItem.firstToken(viewMode: .sourceAccurate),
-          leadingTrivia: Trivia(pieces: triviaBuffer)
-        )
-      )
+      var codeBlockItem = codeBlockItem
+      codeBlockItem.leadingTrivia = Trivia(pieces: triviaBuffer)
+      output.append(codeBlockItem)
+
       triviaBuffer = []
       triviaBuffer += line.trailingTrivia
     }