From 18a26beb425041d8327dfc0e1cd3dfabc0113d16 Mon Sep 17 00:00:00 2001
From: Manny Lopez <manuel.est.lopez@gmail.com>
Date: Mon, 10 Jun 2024 21:52:39 -0700
Subject: [PATCH] Support comments when removing spaces between chained
 functions (#1723)

---
 Sources/Rules.swift               | 10 ++++++----
 Tests/RulesTests+Linebreaks.swift | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/Sources/Rules.swift b/Sources/Rules.swift
index d4f03134c..008d04bff 100644
--- a/Sources/Rules.swift
+++ b/Sources/Rules.swift
@@ -1188,10 +1188,12 @@ public struct _FormatRules {
     ) { formatter in
         formatter.forEach(.operator(".", .infix)) { i, _ in
             let endOfLine = formatter.endOfLine(at: i)
-            if let nextIndex = formatter.index(of: .nonSpaceOrLinebreak, after: endOfLine, if: {
-                $0 == .operator(".", .infix)
-            }) {
-                let startOfLine = formatter.startOfLine(at: nextIndex)
+            if let nextIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: endOfLine),
+               formatter.tokens[nextIndex] == .operator(".", .infix),
+               // Make sure to preserve any code comment between the two lines
+               let nextTokenOrComment = formatter.index(of: .nonSpaceOrLinebreak, after: endOfLine)
+            {
+                let startOfLine = formatter.startOfLine(at: nextTokenOrComment)
                 formatter.removeTokens(in: endOfLine + 1 ..< startOfLine)
             }
         }
diff --git a/Tests/RulesTests+Linebreaks.swift b/Tests/RulesTests+Linebreaks.swift
index f78cc1e26..67a8eb0d7 100644
--- a/Tests/RulesTests+Linebreaks.swift
+++ b/Tests/RulesTests+Linebreaks.swift
@@ -394,6 +394,23 @@ class LinebreakTests: RulesTests {
         testFormatting(for: input, [output1, output2], rules: [FormatRules.blankLinesBetweenChainedFunctions])
     }
 
+    func testBlankLinesWithCommentsBetweenChainedFunctions() {
+        let input = """
+        [0, 1, 2]
+            .map { $0 * 2 }
+
+            // Multiplies by 3
+            .map { $0 * 3 }
+        """
+        let output = """
+        [0, 1, 2]
+            .map { $0 * 2 }
+            // Multiplies by 3
+            .map { $0 * 3 }
+        """
+        testFormatting(for: input, output, rule: FormatRules.blankLinesBetweenChainedFunctions)
+    }
+
     // MARK: - blankLineAfterImports
 
     func testBlankLineAfterImport() {