From 3f5d15a9c7a966227c1f52172f55a0138addf8b7 Mon Sep 17 00:00:00 2001
From: TTOzzi <devttozzi@gmail.com>
Date: Fri, 23 Aug 2024 23:39:36 +0900
Subject: [PATCH] Fix extraction of trailing comments only when necessary
 during afterToken addition

---
 .../PrettyPrint/TokenStreamCreator.swift      | 12 ++-
 .../PrettyPrint/CommaTests.swift              | 96 +++++++++++++++++++
 2 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift b/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
index dd3dd1929..6e090e402 100644
--- a/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
+++ b/Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift
@@ -2929,8 +2929,16 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
         var shouldExtractTrailingComment = false
         if wasLineComment && !hasAppendedTrailingComment {
           switch afterToken {
-          case .break, .printerControl: shouldExtractTrailingComment = true
-          default: break
+          case let .break(kind, _, _):
+            if case let .close(mustBreak) = kind {
+              shouldExtractTrailingComment = mustBreak
+            } else {
+              shouldExtractTrailingComment = true
+            }
+          case .printerControl: 
+            shouldExtractTrailingComment = true
+          default: 
+            break
           }
         }
         if shouldExtractTrailingComment {
diff --git a/Tests/SwiftFormatTests/PrettyPrint/CommaTests.swift b/Tests/SwiftFormatTests/PrettyPrint/CommaTests.swift
index 4264370c8..70c05dfa7 100644
--- a/Tests/SwiftFormatTests/PrettyPrint/CommaTests.swift
+++ b/Tests/SwiftFormatTests/PrettyPrint/CommaTests.swift
@@ -143,6 +143,102 @@ final class CommaTests: PrettyPrintTestCase {
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 40, configuration: configuration)
   }
   
+  func testArrayWithCommentCommasPresentEnabled() {
+    let input =
+      """
+      let MyCollection = [
+        1,
+        2 // some comment
+      ]
+
+      """
+    
+    let expected =
+      """
+      let MyCollection = [
+        1,
+        2,  // some comment
+      ]
+      
+      """
+    
+    var configuration = Configuration.forTesting
+    configuration.multiElementCollectionTrailingCommas = true
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 40, configuration: configuration)
+  }
+  
+  func testArrayWithCommentCommasPresentDisabled() {
+    let input =
+      """
+      let MyCollection = [
+        1,
+        2 // some comment
+      ]
+
+      """
+    
+    let expected =
+      """
+      let MyCollection = [
+        1,
+        2  // some comment
+      ]
+      
+      """
+    
+    var configuration = Configuration.forTesting
+    configuration.multiElementCollectionTrailingCommas = false
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 40, configuration: configuration)
+  }
+  
+  func testArrayWithTernaryOperatorAndCommentCommasPresentEnabled() {
+    let input =
+      """
+      let MyCollection = [
+        1,
+        true ? 1 : 2 // some comment
+      ]
+
+      """
+    
+    let expected =
+      """
+      let MyCollection = [
+        1,
+        true ? 1 : 2,  // some comment
+      ]
+      
+      """
+    
+    var configuration = Configuration.forTesting
+    configuration.multiElementCollectionTrailingCommas = true
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 40, configuration: configuration)
+  }
+  
+  func testArrayWithTernaryOperatorAndCommentCommasPresentDisabled() {
+    let input =
+      """
+      let MyCollection = [
+        1,
+        true ? 1 : 2 // some comment
+      ]
+
+      """
+    
+    let expected =
+      """
+      let MyCollection = [
+        1,
+        true ? 1 : 2  // some comment
+      ]
+      
+      """
+    
+    var configuration = Configuration.forTesting
+    configuration.multiElementCollectionTrailingCommas = false
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 40, configuration: configuration)
+  }
+
   func testDictionaryCommasAbsentEnabled() {
     let input =
       """