Keep EOL comments on same line as when condition

When adding blank lines between when-entries, or when adding braces around the when-condition, then do not wrap the EOL comment after the when-condition

Closes #3209
diff --git a/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditions.kt b/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditions.kt
index 4b488a4..979a0d7 100644
--- a/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditions.kt
+++ b/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditions.kt
@@ -15,8 +15,11 @@
 import com.pinterest.ktlint.rule.engine.core.api.indentWithoutNewlinePrefix
 import com.pinterest.ktlint.rule.engine.core.api.isPartOfComment20
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpace20
+import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpaceWithNewline20
 import com.pinterest.ktlint.rule.engine.core.api.lastChildLeafOrSelf20
+import com.pinterest.ktlint.rule.engine.core.api.leavesForwardsIncludingSelf
 import com.pinterest.ktlint.rule.engine.core.api.nextLeaf
+import com.pinterest.ktlint.rule.engine.core.api.prevCodeLeaf
 import com.pinterest.ktlint.rule.engine.core.api.prevCodeSibling20
 import com.pinterest.ktlint.rule.engine.core.api.prevSibling
 import com.pinterest.ktlint.rule.engine.core.api.upsertWhitespaceBeforeMe
@@ -77,7 +80,11 @@
             .drop(1)
             .forEach { whenEntry ->
                 whenEntry
-                    .findWhitespaceAfterPreviousCodeSibling()
+                    .prevCodeLeaf
+                    ?.leavesForwardsIncludingSelf
+                    ?.takeWhile { !it.isWhiteSpaceWithNewline20 }
+                    ?.last()
+                    ?.nextLeaf
                     ?.takeUnless { it.containsBlankLine() }
                     ?.let { whitespaceBeforeWhenEntry ->
                         emitAndApprove(
diff --git a/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracing.kt b/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracing.kt
index c9a0212..cef94a0 100644
--- a/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracing.kt
+++ b/ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracing.kt
@@ -16,16 +16,18 @@
 import com.pinterest.ktlint.rule.engine.core.api.editorconfig.EditorConfig
 import com.pinterest.ktlint.rule.engine.core.api.editorconfig.INDENT_SIZE_PROPERTY
 import com.pinterest.ktlint.rule.engine.core.api.editorconfig.INDENT_STYLE_PROPERTY
+import com.pinterest.ktlint.rule.engine.core.api.firstChildLeafOrSelf20
 import com.pinterest.ktlint.rule.engine.core.api.ifAutocorrectAllowed
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpace20
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpaceWithNewline20
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling20
 import com.pinterest.ktlint.rule.engine.core.api.parent
+import com.pinterest.ktlint.rule.engine.core.api.prevCodeSibling20
 import com.pinterest.ktlint.rule.engine.core.api.prevSibling20
 import com.pinterest.ktlint.ruleset.standard.StandardRule
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiWhiteSpaceImpl
+import org.jetbrains.kotlin.psi.psiUtil.leaves
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 
 /**
@@ -125,26 +127,37 @@
     private fun ASTNode.surroundWithBraces() {
         require(elementType == ARROW)
         val whenEntryIndent = indentConfig.parentIndentOf(this).removePrefix("\n")
-        val whenEntry =
-            "${whenEntryIndent}true -> {" +
-                // Replace the whitespaces (possibly this could be a proper indent) at the beginning of the body with an indent. In case
-                // the body was already a multiline statement, then the second and following lines should already be properly indented.
-                indentConfig.childIndentOf(this) +
-                siblings()
-                    .dropWhile { it.isWhiteSpace20 }
-                    .joinToString(separator = "") { it.text } +
-                "\n$whenEntryIndent}"
-        val blockExpression = createBlockExpression(whenEntry)
-        parent?.removeRange(nextSibling20!!, null)
-        prevSibling20!!
+        val whenEntry = parent!!
+        // Find the anchor node, e.g. the last node before the current when-entry which is not altered
+        val siblingBeforeWhenEntry = whenEntry.prevSibling20!!
+        // Find the first leaf after the when-entry (including the EOL comment after it) that not has to be enclosed within the braces
+        val stopLeaf =
+            siblingBeforeWhenEntry
+                .siblings()
+                .takeWhile { !it.isWhiteSpaceWithNewline20 }
+                .last()
+                .nextSibling20!!
+                .firstChildLeafOrSelf20
+        val whenEntryNode =
+            createWhenEntryNode(
+                "${whenEntryIndent}${prevCodeSibling20!!.text} -> {" +
+                    // Replace the whitespaces (possibly this could be a proper indent) at the beginning of the body with an indent. In case
+                    // the body was already a multiline statement, then the second and following lines should already be properly indented.
+                    indentConfig.childIndentOf(this) +
+                    leaves()
+                        .dropWhile { it.isWhiteSpace20 }
+                        .takeWhile { it != stopLeaf }
+                        .joinToString(separator = "") { it.text } +
+                    "\n$whenEntryIndent}",
+            )
+        // Remove the old when entry and if applicable the EOL-comment after it
+        whenEntry.removeRange(whenEntry, stopLeaf)
+        siblingBeforeWhenEntry
             .parent
-            ?.run {
-                addChild(PsiWhiteSpaceImpl(" "), null)
-                addChild(blockExpression!!, null)
-            }
+            ?.run { addChild(whenEntryNode!!, stopLeaf) }
     }
 
-    private fun createBlockExpression(whenEntry: String) =
+    private fun createWhenEntryNode(whenEntry: String) =
         KtlintKotlinCompiler
             .createASTNodeFromText(
                 """
@@ -154,7 +167,6 @@
                 """.trimMargin(),
             )?.findChildByType(WHEN)
             ?.findChildByType(WHEN_ENTRY)
-            ?.findChildByType(BLOCK)
 }
 
 public val WHEN_ENTRY_BRACING_RULE_ID: RuleId = WhenEntryBracing().ruleId
diff --git a/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditionsTest.kt b/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditionsTest.kt
index f98e079..4ebfc89 100644
--- a/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditionsTest.kt
+++ b/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/BlankLineBetweenWhenConditionsTest.kt
@@ -283,16 +283,36 @@
     }
 
     @Test
-    fun `Issue 3209 - Given a when-statement with single line when-conditions followed by comments then`() {
+    fun `Issue 3209 - Given a when-statement with single line when-conditions followed by EOL comments, and one multiline when-entry then add blank lines but keep EOL comments on same line`() {
         val code =
             """
             val foo =
                 when (bar) {
                     BAR1 -> "bar1" // bar 1
                     BAR2 -> "bar2" // bar 2
-                    else -> null
+                    else -> {
+                        null
+                    }
                 }
             """.trimIndent()
-        blankLineAfterWhenConditionRuleAssertThat(code).hasNoLintViolations()
+        val formattedCode =
+            """
+            val foo =
+                when (bar) {
+                    BAR1 -> "bar1" // bar 1
+
+                    BAR2 -> "bar2" // bar 2
+
+                    else -> {
+                        null
+                    }
+                }
+            """.trimIndent()
+        @Suppress("ktlint:standard:max-line-length")
+        blankLineAfterWhenConditionRuleAssertThat(code)
+            .hasLintViolations(
+                LintViolation(4, 1, "Add a blank line between all when-conditions in case at least one multiline when-condition is found in the statement"),
+                LintViolation(5, 1, "Add a blank line between all when-conditions in case at least one multiline when-condition is found in the statement"),
+            ).isFormattedAs(formattedCode)
     }
 }
diff --git a/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracingTest.kt b/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracingTest.kt
index d970314..1954042 100644
--- a/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracingTest.kt
+++ b/ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/rules/WhenEntryBracingTest.kt
@@ -195,4 +195,45 @@
                 LintViolation(6, 17, "Body of when entry should be surrounded by braces if any when entry body is surrounded by braces or has a multiline body"),
             ).isFormattedAs(formattedCode)
     }
+
+    @Test
+    fun `Given a when-statement with a when entry condition followed by an EOL comment on same line then keep that comment inside the block`() {
+        val code =
+            """
+            val foo =
+                when (bar) {
+                    BAR1 -> "bar1" // bar 1 comment
+                    // comment before BAR2
+                    BAR2 ->
+                        "bar2" // bar 2 comment
+                    else -> null // else comment
+                }
+            """.trimIndent()
+        val formattedCode =
+            """
+            val foo =
+                when (bar) {
+                    BAR1 -> {
+                        "bar1" // bar 1 comment
+                    }
+                    // comment before BAR2
+                    BAR2 -> {
+                        "bar2" // bar 2 comment
+                    }
+                    else -> {
+                        null // else comment
+                    }
+                }
+            """.trimIndent()
+        @Suppress("ktlint:standard:max-line-length")
+        whenEntryBracingRuleAssertThat(code)
+            .addAdditionalRuleProvider {
+                // Ensures that the first when entry is also wrapped to a multiline body
+                StatementWrappingRule()
+            }.hasLintViolations(
+                LintViolation(3, 17, "Body of when entry should be surrounded by braces if any when entry body is surrounded by braces or has a multiline body"),
+                LintViolation(6, 13, "Body of when entry should be surrounded by braces if any when entry body is surrounded by braces or has a multiline body"),
+                LintViolation(7, 17, "Body of when entry should be surrounded by braces if any when entry body is surrounded by braces or has a multiline body"),
+            ).isFormattedAs(formattedCode)
+    }
 }