Fix trivia handling issues in editor placeholder expansion

rdar://123287930

Author: ahoppen

Release Notes/600.md

@@ -53,6 +53,10 @@
   - Description: `macroSpecs` can have additional specifications like conformances provided by member or extension macro that can be used for macro expansion.
   - Issue: https://github.com/apple/swift-syntax/issues/2031
   - Pull Request: https://github.com/apple/swift-syntax/pull/2327
+  
+- `BasicFormat.reset`
+  - Description: `BasicFormat` stores internal state while it formats a syntax tree. If the same `BasicFormat` instance is used to format multiple syntax trees, the initial state nees to be reset between calls to `BasicFormat.rewrite`. `SyntaxProtocol.formatted()` automatically handles this.
+  - Pull Request: https://github.com/apple/swift-syntax/pull/2511
 
 ## API Behavior Changes
 

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.
+  public 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)
   }
 }
 
@@ -80,10 +81,12 @@ extension FunctionCallExprSyntax {
 
     // First trailing closure won't have label/colon. Transfer their trivia.
     var trailingClosure = closures.first!.closure
-    trailingClosure.leadingTrivia = Trivia()
+    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(),
@@ -94,6 +97,9 @@ 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))
@@ -109,11 +115,17 @@ extension FunctionCallExprSyntax {
           .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.rightParen = TokenSyntax.rightParenToken(trailingTrivia: Trivia().merging(triviaOf: comma))
-      }
+      // 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)
     }
 
@@ -124,12 +136,22 @@ extension FunctionCallExprSyntax {
       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.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)
+  }
 }

Tests/SwiftRefactorTest/CallToTrailingClosureTests.swift

@@ -217,6 +217,36 @@ final class CallToTrailingClosuresTest: XCTestCase {
 
     try assertRefactorCall(baseline, expected: expected)
   }
+
+  func testClosureWithArgumentLabel() throws {
+    try assertRefactorCall(
+      """
+      foo(arg: 1, closure: { someInt in
+
+      })
+      """,
+      expected: """
+        foo(arg: 1) { someInt in
+
+        }
+        """
+    )
+  }
+
+  func testCallHasInitialIndentationIndentation() throws {
+    try assertRefactorCall(
+      """
+          foo(arg: 1, closure: { someInt in
+              "abc"
+          })
+      """,
+      expected: """
+            foo(arg: 1) { someInt in
+                "abc"
+            }
+        """
+    )
+  }
 }
 
 fileprivate func assertRefactorCall(