Fix trivia handling issues in editor placeholder expansion

Sources/SwiftBasicFormat/BasicFormat.swift

@@ -69,6 +69,16 @@ open class BasicFormat: SyntaxRewriter {
     super.init(viewMode: viewMode)
   }
 
+  /// Clears all stateful data from this `BasicFormat`.
+  ///
+  /// This needs to be called between multiple `rewrite` calls to a syntax tree.
+  func reset() {
+    indentationStack = [indentationStack.first!]
+    anchorPoints = [:]
+    previousToken = nil
+    stringLiteralNestingLevel = 0
+  }
+
   // MARK: - Updating indentation level
 
   public func increaseIndentationLevel(to userDefinedIndentation: Trivia? = nil) {
@@ -145,7 +155,7 @@ open class BasicFormat: SyntaxRewriter {
 
   /// Find the indentation of the nearest ancestor whose first token is an
   /// anchor point (see `anchorPoints`).
-  private func anchorPointIndentation(for token: TokenSyntax) -> Trivia? {
+  private func anchorPointIndentation(for token: TokenSyntax) -> Trivia {
     var ancestor: Syntax = Syntax(token)
     while let parent = ancestor.parent {
       ancestor = parent
@@ -155,7 +165,7 @@ open class BasicFormat: SyntaxRewriter {
         return anchorPointIndentation
       }
     }
-    return nil
+    return Trivia()
   }
 
   // MARK: - Customization points
@@ -575,23 +585,24 @@ open class BasicFormat: SyntaxRewriter {
       trailingTrivia += .space
     }
 
-    var leadingTriviaIndentation = self.currentIndentationLevel
-    var trailingTriviaIndentation = self.currentIndentationLevel
+    let leadingTriviaIndentation: Trivia
+    let trailingTriviaIndentation: Trivia
 
     // If the trivia contains user-defined indentation, find their anchor point
     // and indent the token relative to that anchor point.
+    //
     // Always indent string literals relative to their anchor point because
     // their indentation has structural meaning and we just want to maintain
     // what the user wrote.
-    if leadingTrivia.containsIndentation(isOnNewline: previousTokenWillEndWithNewline) || isInsideStringLiteral,
-      let anchorPointIndentation = self.anchorPointIndentation(for: token)
-    {
-      leadingTriviaIndentation = anchorPointIndentation
+    if leadingTrivia.containsIndentation(isOnNewline: previousTokenWillEndWithNewline) || isInsideStringLiteral {
+      leadingTriviaIndentation = anchorPointIndentation(for: token)
+    } else {
+      leadingTriviaIndentation = currentIndentationLevel
     }
-    if combinedTrailingTrivia.containsIndentation(isOnNewline: previousTokenWillEndWithNewline) || isInsideStringLiteral,
-      let anchorPointIndentation = self.anchorPointIndentation(for: token)
-    {
-      trailingTriviaIndentation = anchorPointIndentation
+    if combinedTrailingTrivia.containsIndentation(isOnNewline: previousTokenWillEndWithNewline) {
+      trailingTriviaIndentation = anchorPointIndentation(for: token)
+    } else {
+      trailingTriviaIndentation = currentIndentationLevel
     }
 
     leadingTrivia = leadingTrivia.indented(

Sources/SwiftBasicFormat/Syntax+Extensions.swift

@@ -19,7 +19,7 @@ extension TokenSyntax {
   /// find the indentation of the line this token occurs on.
   private var indentation: Trivia? {
     let previous = self.previousToken(viewMode: .sourceAccurate)
-    return ((previous?.trailingTrivia ?? []) + leadingTrivia).indentation(isOnNewline: false)
+    return ((previous?.trailingTrivia ?? []) + leadingTrivia).indentation(isOnNewline: previous == nil)
   }
 
   /// Returns the indentation of the line this token occurs on

Sources/SwiftBasicFormat/SyntaxProtocol+Formatted.swift

@@ -15,6 +15,7 @@ import SwiftSyntax
 public extension SyntaxProtocol {
   /// Build a syntax node from this `Buildable` and format it with the given format.
   func formatted(using format: BasicFormat = BasicFormat()) -> Syntax {
+    format.reset()
     return format.rewrite(self)
   }
 }

Sources/SwiftRefactor/CallToTrailingClosures.swift

@@ -48,7 +48,8 @@ public struct CallToTrailingClosures: SyntaxRefactoringProvider {
   // TODO: Rather than returning nil, we should consider throwing errors with
   // appropriate messages instead.
   public static func refactor(syntax call: FunctionCallExprSyntax, in context: Context = Context()) -> FunctionCallExprSyntax? {
-    return call.convertToTrailingClosures(from: context.startAtArgument)?.formatted().as(FunctionCallExprSyntax.self)
+    let converted = call.convertToTrailingClosures(from: context.startAtArgument)
+    return converted?.formatted().as(FunctionCallExprSyntax.self)
   }
 }
 
@@ -69,7 +70,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 = closure.with(\.trailingTrivia, closure.trailingTrivia.merging(triviaOf: comma))
+        closure.trailingTrivia = closure.trailingTrivia.merging(triviaOf: comma)
       }
       closures.append((arg, closure))
     }
@@ -80,13 +81,12 @@ extension FunctionCallExprSyntax {
 
     // First trailing closure won't have label/colon. Transfer their trivia.
     var trailingClosure = closures.first!.closure
-      .with(
-        \.leadingTrivia,
-        Trivia()
-          .merging(triviaOf: closures.first!.original.label)
-          .merging(triviaOf: closures.first!.original.colon)
-          .merging(closures.first!.closure.leadingTrivia)
-      )
+    trailingClosure.leadingTrivia =
+      Trivia()
+      .merging(triviaOf: closures.first!.original.label)
+      .merging(triviaOf: closures.first!.original.colon)
+      .merging(closures.first!.closure.leadingTrivia)
+      .droppingLeadingWhitespace
     let additionalTrailingClosures = closures.dropFirst().map {
       MultipleTrailingClosureElementSyntax(
         label: $0.original.label ?? .wildcardToken(),
@@ -97,53 +97,61 @@ extension FunctionCallExprSyntax {
 
     var converted = self.detached
 
+    // Trivia that should be attached to the end of the converted call.
+    var additionalTriviaAtEndOfCall: Trivia? = nil
+
     // Remove parens if there's no non-closure arguments left and remove the
     // last comma otherwise. Makes sure to keep the trivia of any removed node.
     var argList = Array(arguments.dropLast(closures.count))
     if argList.isEmpty {
-      converted =
-        converted
-        .with(\.leftParen, nil)
-        .with(\.rightParen, nil)
+      converted.leftParen = nil
+      converted.rightParen = nil
 
       // 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 = trailingClosure.with(
-          \.leadingTrivia,
-          Trivia()
-            .merging(triviaOf: leftParen)
-            .merging(trailingClosure.leadingTrivia)
-        )
+        trailingClosure.leadingTrivia = Trivia()
+          .merging(triviaOf: leftParen)
+          .merging(trailingClosure.leadingTrivia)
+      }
+      // No right paren anymore. Attach its trivia to the end of the call.
+      if let rightParen = rightParen {
+        additionalTriviaAtEndOfCall = Trivia().merging(triviaOf: rightParen)
       }
     } else {
       let last = argList.last!
-      if let comma = last.trailingComma {
-        converted =
-          converted
-          .with(\.rightParen, TokenSyntax.rightParenToken(trailingTrivia: Trivia().merging(triviaOf: comma)))
-      }
-      argList[argList.count - 1] =
-        last
-        .with(\.trailingComma, nil)
+      // 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)
+      additionalTriviaAtEndOfCall = rightParen?.trailingTrivia
+      argList[argList.count - 1] = last.with(\.trailingComma, nil)
     }
 
     // Update arguments and trailing closures
-    converted =
-      converted
-      .with(\.arguments, LabeledExprListSyntax(argList))
-      .with(\.trailingClosure, trailingClosure)
+    converted.arguments = LabeledExprListSyntax(argList)
+    converted.trailingClosure = trailingClosure
     if !additionalTrailingClosures.isEmpty {
-      converted = converted.with(\.additionalTrailingClosures, MultipleTrailingClosureElementListSyntax(additionalTrailingClosures))
+      converted.additionalTrailingClosures = MultipleTrailingClosureElementListSyntax(additionalTrailingClosures)
     }
 
-    // The right paren either doesn't exist any more, or is before all the
-    // trailing closures. Moves its trivia to the end of the converted call.
-    if let rightParen = rightParen {
-      converted = converted.with(\.trailingTrivia, converted.trailingTrivia.merging(triviaOf: rightParen))
+    if let additionalTriviaAtEndOfCall {
+      converted.trailingTrivia = converted.trailingTrivia.merging(additionalTriviaAtEndOfCall.droppingLeadingWhitespace)
     }
 
     return converted
   }
 }
+
+fileprivate extension Trivia {
+  var droppingLeadingWhitespace: Trivia {
+    return Trivia(pieces: self.drop(while: \.isWhitespace))
+  }
+}
+
+fileprivate extension Sequence {
+  func dropSuffix(while predicate: (Element) -> Bool) -> [Element] {
+    self.reversed().drop(while: predicate).reversed()
+  }
+}

Sources/SwiftRefactor/ExpandEditorPlaceholder.swift

@@ -70,7 +70,17 @@ import SwiftSyntaxBuilder
 /// anything here
 /// ```
 struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
-  static func textRefactor(syntax token: TokenSyntax, in context: Void) -> [SourceEdit] {
+  struct Context {
+    let indentationWidth: Trivia?
+    let initialIndentation: Trivia
+
+    init(indentationWidth: Trivia? = nil, initialIndentation: Trivia = []) {
+      self.indentationWidth = indentationWidth
+      self.initialIndentation = initialIndentation
+    }
+  }
+
+  static func textRefactor(syntax token: TokenSyntax, in context: Context = Context()) -> [SourceEdit] {
     guard let placeholder = EditorPlaceholderData(token: token) else {
       return []
     }
@@ -81,7 +91,18 @@ struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
       expanded = String(text)
     case let .typed(text, type):
       if let functionType = type.as(FunctionTypeSyntax.self) {
-        expanded = functionType.closureExpansion.formatted().description
+        let basicFormat = BasicFormat(
+          indentationWidth: context.indentationWidth,
+          initialIndentation: context.initialIndentation
+        )
+        var formattedExpansion = functionType.closureExpansion.formatted(using: basicFormat).description
+        // Strip the initial indentation from the placeholder itself. We only introduced the initial indentation to
+        // format consecutive lines. We don't want it at the front of the initial line because it replaces an expression
+        // that might be in the middle of a line.
+        if formattedExpansion.hasPrefix(context.initialIndentation.description) {
+          formattedExpansion = String(formattedExpansion.dropFirst(context.initialIndentation.description.count))
+        }
+        expanded = formattedExpansion
       } else {
         expanded = String(text)
       }
@@ -122,13 +143,25 @@ struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
 ///
 /// Expansion on `closure1` and `normalArg` is the same as `ExpandSingleEditorPlaceholder`.
 public struct ExpandEditorPlaceholder: EditRefactoringProvider {
-  public static func textRefactor(syntax token: TokenSyntax, in context: Void) -> [SourceEdit] {
+  public struct Context {
+    public let indentationWidth: Trivia?
+
+    public init(indentationWidth: Trivia? = nil) {
+      self.indentationWidth = indentationWidth
+    }
+  }
+
+  public static func textRefactor(syntax token: TokenSyntax, in context: Context = Context()) -> [SourceEdit] {
     guard let placeholder = token.parent?.as(DeclReferenceExprSyntax.self),
       placeholder.baseName.isEditorPlaceholder,
       let arg = placeholder.parent?.as(LabeledExprSyntax.self),
       let argList = arg.parent?.as(LabeledExprListSyntax.self),
       let call = argList.parent?.as(FunctionCallExprSyntax.self),
-      let expandedTrailingClosures = ExpandEditorPlaceholdersToTrailingClosures.expandTrailingClosurePlaceholders(in: call, ifIncluded: arg)
+      let expandedTrailingClosures = ExpandEditorPlaceholdersToTrailingClosures.expandTrailingClosurePlaceholders(
+        in: call,
+        ifIncluded: arg,
+        indentationWidth: context.indentationWidth
+      )
     else {
       return ExpandSingleEditorPlaceholder.textRefactor(syntax: token)
     }
@@ -159,8 +192,16 @@ public struct ExpandEditorPlaceholder: EditRefactoringProvider {
 /// }
 /// ```
 public struct ExpandEditorPlaceholdersToTrailingClosures: SyntaxRefactoringProvider {
-  public static func refactor(syntax call: FunctionCallExprSyntax, in context: Void = ()) -> FunctionCallExprSyntax? {
-    return Self.expandTrailingClosurePlaceholders(in: call, ifIncluded: nil)
+  public struct Context {
+    public let indentationWidth: Trivia?
+
+    public init(indentationWidth: Trivia? = nil) {
+      self.indentationWidth = indentationWidth
+    }
+  }
+
+  public static func refactor(syntax call: FunctionCallExprSyntax, in context: Context = Context()) -> FunctionCallExprSyntax? {
+    return Self.expandTrailingClosurePlaceholders(in: call, ifIncluded: nil, indentationWidth: context.indentationWidth)
   }
 
   /// If the given argument is `nil` or one of the last arguments that are all
@@ -170,9 +211,10 @@ public struct ExpandEditorPlaceholdersToTrailingClosures: SyntaxRefactoringProvi
   /// Otherwise return nil.
   fileprivate static func expandTrailingClosurePlaceholders(
     in call: FunctionCallExprSyntax,
-    ifIncluded arg: LabeledExprSyntax?
+    ifIncluded arg: LabeledExprSyntax?,
+    indentationWidth: Trivia?
   ) -> FunctionCallExprSyntax? {
-    guard let expanded = call.expandTrailingClosurePlaceholders(ifIncluded: arg) else {
+    guard let expanded = call.expandTrailingClosurePlaceholders(ifIncluded: arg, indentationWidth: indentationWidth) else {
       return nil
     }
 
@@ -260,7 +302,8 @@ extension FunctionCallExprSyntax {
   /// closures based on the function types provided by each editor placeholder.
   /// Otherwise return nil.
   fileprivate func expandTrailingClosurePlaceholders(
-    ifIncluded: LabeledExprSyntax?
+    ifIncluded: LabeledExprSyntax?,
+    indentationWidth: Trivia?
   ) -> (expr: FunctionCallExprSyntax, numClosures: Int)? {
     var includedArg = false
     var argsToExpand = 0
@@ -283,9 +326,14 @@ extension FunctionCallExprSyntax {
       return nil
     }
 
+    let lineIndentation = self.firstToken(viewMode: .sourceAccurate)?.indentationOfLine ?? []
+
     var expandedArgs = [LabeledExprSyntax]()
     for arg in arguments.suffix(argsToExpand) {
-      let edits = ExpandSingleEditorPlaceholder.textRefactor(syntax: arg.expression.cast(DeclReferenceExprSyntax.self).baseName)
+      let edits = ExpandSingleEditorPlaceholder.textRefactor(
+        syntax: arg.expression.cast(DeclReferenceExprSyntax.self).baseName,
+        in: ExpandSingleEditorPlaceholder.Context(indentationWidth: indentationWidth, initialIndentation: lineIndentation)
+      )
       guard edits.count == 1, let edit = edits.first, !edit.replacement.isEmpty else {
         return nil
       }

Tests/SwiftBasicFormatTest/BasicFormatTests.swift

@@ -610,4 +610,39 @@ final class BasicFormatTest: XCTestCase {
     )
   }
 
+  func testFormatClosureWithInitialIndentation() throws {
+    assertFormatted(
+      tree: ClosureExprSyntax(
+        statements: CodeBlockItemListSyntax([
+          CodeBlockItemSyntax(item: CodeBlockItemSyntax.Item(IntegerLiteralExprSyntax(integerLiteral: 2)))
+        ])
+      ),
+      expected: """
+            {
+                2
+            }
+        """,
+      using: BasicFormat(initialIndentation: .spaces(4))
+    )
+  }
+
+  func testIndentedStandaloneClosureRoundTrips() throws {
+    let source = """
+          foo {
+              "abc"
+          }
+      """
+    assertFormatted(source: source, expected: source)
+  }
+
+  func testIndentedStandaloneClosureRoundTrips2() throws {
+    let source = """
+          foo {
+              if true {
+                  print("test")
+              }
+          }
+      """
+    assertFormatted(source: source, expected: source)
+  }
 }