From 9626e004e7cdb1327118e75dcf31dadfbf885c72 Mon Sep 17 00:00:00 2001
From: Tony Allevato <allevato@google.com>
Date: Thu, 19 Oct 2023 13:13:37 -0400
Subject: [PATCH] Fix multi-line string wrapping in `@available` attributes.

---
 .../PrettyPrint/TokenStreamCreator.swift      | 29 +++++++++-
 .../PrettyPrint/AttributeTests.swift          | 54 +++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift b/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
index 0137e3a2e..85ccb71d0 100644
--- a/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
+++ b/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
@@ -1801,8 +1801,23 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: AvailabilityLabeledArgumentSyntax) -> SyntaxVisitorContinueKind {
     before(node.label, tokens: .open)
-    after(node.colon, tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
-    after(node.value.lastToken(viewMode: .sourceAccurate), tokens: .close)
+
+    let tokensAfterColon: [Token]
+    let endTokens: [Token]
+
+    if case .string(let string) = node.value,
+      string.openingQuote.tokenKind == .multilineStringQuote
+    {
+      tokensAfterColon =
+        [.break(.open(kind: .block), newlines: .elective(ignoresDiscretionary: true))]
+      endTokens = [.break(.close(mustBreak: false), size: 0), .close]
+    } else {
+      tokensAfterColon = [.break(.continue, newlines: .elective(ignoresDiscretionary: true))]
+      endTokens = [.close]
+    }
+
+    after(node.colon, tokens: tokensAfterColon)
+    after(node.value.lastToken(viewMode: .sourceAccurate), tokens: endTokens)
     return .visitChildren
   }
 
@@ -2382,6 +2397,16 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visit(_ node: SimpleStringLiteralExprSyntax) -> SyntaxVisitorContinueKind {
+    if node.openingQuote.tokenKind == .multilineStringQuote {
+      after(node.openingQuote, tokens: .break(.same, size: 0, newlines: .hard(count: 1)))
+      if !node.segments.isEmpty {
+        before(node.closingQuote, tokens: .break(.same, newlines: .hard(count: 1)))
+      }
+    }
+    return .visitChildren
+  }
+
   override func visit(_ node: StringSegmentSyntax) -> SyntaxVisitorContinueKind {
     // Looks up the correct break kind based on prior context.
     func breakKind() -> BreakKind {
diff --git a/Tests/SwiftFormatTests/PrettyPrint/AttributeTests.swift b/Tests/SwiftFormatTests/PrettyPrint/AttributeTests.swift
index 57a12a3a7..06aa46c8f 100644
--- a/Tests/SwiftFormatTests/PrettyPrint/AttributeTests.swift
+++ b/Tests/SwiftFormatTests/PrettyPrint/AttributeTests.swift
@@ -356,4 +356,58 @@ final class AttributeTests: PrettyPrintTestCase {
 
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 32)
   }
+
+  func testMultilineStringLiteralInCustomAttribute() {
+    let input =
+      #"""
+      @CustomAttribute(message: """
+      This is a
+      multiline
+      string
+      """)
+      public func f() {}
+      """#
+
+    let expected =
+      #"""
+      @CustomAttribute(
+        message: """
+          This is a
+          multiline
+          string
+          """)
+      public func f() {}
+
+      """#
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 100)
+  }
+
+  func testMultilineStringLiteralInAvailableAttribute() {
+    let input =
+      #"""
+      @available(*, deprecated, message: """
+      This is a
+      multiline
+      string
+      """)
+      public func f() {}
+      """#
+
+    let expected =
+      #"""
+      @available(
+        *, deprecated,
+        message: """
+          This is a
+          multiline
+          string
+          """
+      )
+      public func f() {}
+
+      """#
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 100)
+  }
 }