Fix multiline string indentation

Author: kimdv

CodeGeneration/Sources/generate-swiftsyntax/templates/basicformat/BasicFormatFile.swift

@@ -93,10 +93,15 @@ let basicFormatFile = SourceFileSyntax(leadingTrivia: generateCopyrightHeader(fo
           trailingTrivia += .space
         }
         if let keyPath = getKeyPath(Syntax(node)), requiresLeadingNewline(keyPath), !(leadingTrivia.first?.isNewline ?? false) {
-          leadingTrivia = .newline + leadingTrivia
+          if !shouldOmitNewline(node) {
+            leadingTrivia = .newline + leadingTrivia
+          }
+        }
+        var isOnNewline: Bool = (lastRewrittenToken?.trailingTrivia.pieces.last?.isNewline == true)
+        if case .stringSegment(let text) = lastRewrittenToken?.tokenKind {
+          isOnNewline = isOnNewline || (text.last?.isNewline == true)
         }
-        leadingTrivia = leadingTrivia.indented(indentation: indentation)
-        trailingTrivia = trailingTrivia.indented(indentation: indentation)
+        leadingTrivia = leadingTrivia.indented(indentation: indentation, isOnNewline: isOnNewline)
         let rewritten = TokenSyntax(
           node.tokenKind,
           leadingTrivia: leadingTrivia,
@@ -110,6 +115,22 @@ let basicFormatFile = SourceFileSyntax(leadingTrivia: generateCopyrightHeader(fo
       """
     )
 
+    DeclSyntax(
+      """
+      open func shouldOmitNewline(_ node: TokenSyntax) -> Bool {
+        var ancestor: Syntax = Syntax(node)
+        while let parent = ancestor.parent {
+            ancestor = parent
+            if ancestor.is(ExpressionSegmentSyntax.self) {
+                return true
+            }
+        }
+
+        return false
+      }
+      """
+    )
+
     try FunctionDeclSyntax("open func shouldIndent(_ keyPath: AnyKeyPath) -> Bool") {
       try SwitchExprSyntax("switch keyPath") {
         for node in SYNTAX_NODES where !node.isBase {

Sources/SwiftBasicFormat/Trivia+Indented.swift

@@ -13,25 +13,33 @@
 import SwiftSyntax
 
 extension Trivia {
-  func indented(indentation: TriviaPiece) -> Trivia {
+  func indented(indentation: TriviaPiece, isOnNewline: Bool = false) -> Trivia {
     var indentedPieces: [TriviaPiece] = []
     for (index, piece) in self.enumerated() {
-      let nextPiece = index < pieces.count - 1 ? pieces[index + 1] : nil
-      indentedPieces.append(piece)
-      if piece.isNewline {
-        switch (nextPiece, indentation) {
-        case (.spaces(let nextPieceSpaces)?, .spaces(let indentationSpaces)):
+      let previousPieceIsNewline: Bool
+      if index == 0 {
+        previousPieceIsNewline = isOnNewline
+      } else {
+        previousPieceIsNewline = pieces[index - 1].isNewline
+      }
+      if previousPieceIsNewline {
+        switch (piece, indentation) {
+        case (.spaces(let nextPieceSpaces), .spaces(let indentationSpaces)):
           if nextPieceSpaces < indentationSpaces {
             indentedPieces.append(.spaces(indentationSpaces - nextPieceSpaces))
           }
-        case (.tabs(let nextPieceTabs)?, .tabs(let indentationTabs)):
+        case (.tabs(let nextPieceTabs), .tabs(let indentationTabs)):
           if nextPieceTabs < indentationTabs {
             indentedPieces.append(.tabs(indentationTabs - nextPieceTabs))
           }
         default:
           indentedPieces.append(indentation)
         }
       }
+      indentedPieces.append(piece)
+    }
+    if self.pieces.last?.isNewline == true {
+      indentedPieces.append(indentation)
     }
     return Trivia(pieces: indentedPieces)
   }

Sources/SwiftBasicFormat/generated/BasicFormat.swift

@@ -55,10 +55,15 @@ open class BasicFormat: SyntaxRewriter {
       trailingTrivia += .space
     }
     if let keyPath = getKeyPath(Syntax(node)), requiresLeadingNewline(keyPath), !(leadingTrivia.first?.isNewline ?? false) {
-      leadingTrivia = .newline + leadingTrivia
+      if !shouldOmitNewline(node) {
+        leadingTrivia = .newline + leadingTrivia
+      }
     }
-    leadingTrivia = leadingTrivia.indented(indentation: indentation)
-    trailingTrivia = trailingTrivia.indented(indentation: indentation)
+    var isOnNewline: Bool = (lastRewrittenToken?.trailingTrivia.pieces.last?.isNewline == true)
+    if case .stringSegment(let text) = lastRewrittenToken?.tokenKind {
+      isOnNewline = isOnNewline || (text.last?.isNewline == true)
+    }
+    leadingTrivia = leadingTrivia.indented(indentation: indentation, isOnNewline: isOnNewline)
     let rewritten = TokenSyntax(
         node.tokenKind, 
         leadingTrivia: leadingTrivia, 
@@ -71,6 +76,18 @@ open class BasicFormat: SyntaxRewriter {
     return rewritten
   }
 
+  open func shouldOmitNewline(_ node: TokenSyntax) -> Bool {
+    var ancestor: Syntax = Syntax(node)
+    while let parent = ancestor.parent {
+      ancestor = parent
+      if ancestor.is(ExpressionSegmentSyntax.self) {
+          return true
+      }
+    }
+
+    return false
+  }
+  
   open func shouldIndent(_ keyPath: AnyKeyPath) -> Bool {
     switch keyPath {
     case \AccessorBlockSyntax.accessors:

Tests/SwiftSyntaxBuilderTest/StringLiteralTests.swift

@@ -282,11 +282,11 @@ final class StringLiteralTests: XCTestCase {
       buildable,
       #"""
       assertionFailure("""
-            Error validating child at index \(index) of \(nodeKind):
-            Node did not satisfy any node choice requirement.
-            Validation failures:
-            \(nonNilErrors.map({ "- \($0.description)" }).joined(separator: "\n"))
-            """, file: file, line: line)
+          Error validating child at index \(index) of \(nodeKind):
+          Node did not satisfy any node choice requirement.
+          Validation failures:
+          \(nonNilErrors.map({ "- \($0.description)" }).joined(separator: "\n"))
+          """, file: file, line: line)
       """#
     )
   }
@@ -309,10 +309,10 @@ final class StringLiteralTests: XCTestCase {
       #"""
       if true {
           assertionFailure("""
-                  Error validating child at index
-                  Node did not satisfy any node choice requirement.
-                  Validation failures:
-                  """)
+              Error validating child at index
+              Node did not satisfy any node choice requirement.
+              Validation failures:
+              """)
       }
       """#
     )
@@ -339,10 +339,10 @@ final class StringLiteralTests: XCTestCase {
       if true {
           assertionFailure(
               """
-                  Error validating child at index
-                  Node did not satisfy any node choice requirement.
-                  Validation failures:
-                  """
+              Error validating child at index
+              Node did not satisfy any node choice requirement.
+              Validation failures:
+              """
           )
       }
       """#