MultiLineIfElseRule: place end of line comments in block (#945)

pinterest · Oct 19, 2020 · 2a4b537 · 2a4b537
1 parent 69cc0f7
commit 2a4b537
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 2 deletions.
diff --git a/...rimental/src/main/kotlin/com/pinterest/ktlint/ruleset/experimental/MultiLineIfElseRule.kt b/...rimental/src/main/kotlin/com/pinterest/ktlint/ruleset/experimental/MultiLineIfElseRule.kt
@@ -5,13 +5,16 @@ import com.pinterest.ktlint.core.ast.ElementType.ELSE
 import com.pinterest.ktlint.core.ast.ElementType.ELSE_KEYWORD
 import com.pinterest.ktlint.core.ast.ElementType.LBRACE
 import com.pinterest.ktlint.core.ast.ElementType.RBRACE
+import com.pinterest.ktlint.core.ast.ElementType.RPAR
 import com.pinterest.ktlint.core.ast.ElementType.THEN
+import com.pinterest.ktlint.core.ast.isWhiteSpaceWithNewline
 import com.pinterest.ktlint.core.ast.prevLeaf
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiWhiteSpaceImpl
 import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.psiUtil.leaves
 
 /**
  * https://kotlinlang.org/docs/reference/coding-conventions.html#formatting-control-flow-statements
@@ -40,17 +43,21 @@ class MultiLineIfElseRule : Rule("multiline-if-else") {
     }
 
     private fun autocorrect(node: ASTNode) {
-        val bodyIndent = node.treePrev.text
+        val prevLeaves =
+            node.leaves(forward = false).takeWhile { it.elementType !in listOf(RPAR, ELSE_KEYWORD) }.toList().reversed()
+        val nextLeaves = node.leaves(forward = true).takeWhile { !it.isWhiteSpaceWithNewline() }.toList()
         val rightBraceIndent = node.treeParent
             .prevLeaf { it is PsiWhiteSpace && it.textContains('\n') }?.text.orEmpty()
             .let { "\n${it.substringAfterLast("\n")}" }
+
         (node.treePrev as LeafPsiElement).rawReplaceWithText(" ")
         KtBlockExpression(null).apply {
             val previousChild = node.firstChildNode
             node.replaceChild(node.firstChildNode, this)
             addChild(LeafPsiElement(LBRACE, "{"))
-            addChild(PsiWhiteSpaceImpl(bodyIndent))
+            prevLeaves.forEach(::addChild)
             addChild(previousChild)
+            nextLeaves.forEach(::addChild)
             addChild(PsiWhiteSpaceImpl(rightBraceIndent))
             addChild(LeafPsiElement(RBRACE, "}"))
         }

diff --git a/...ntal/src/test/kotlin/com/pinterest/ktlint/ruleset/experimental/MultiLineIfElseRuleTest.kt b/...ntal/src/test/kotlin/com/pinterest/ktlint/ruleset/experimental/MultiLineIfElseRuleTest.kt
@@ -358,6 +358,64 @@ class MultiLineIfElseRuleTest {
         )
     }
 
+    @Test
+    fun testWithEolComments() {
+        val actual = format(
+            """
+            fun test() {
+                val s = if (x > 0)
+                // comment1
+                    "a"
+                else
+                // comment2
+                    "b"
+            }
+            """.trimIndent()
+        )
+        assertThat(actual).isEqualTo(
+            """
+            fun test() {
+                val s = if (x > 0) {
+                // comment1
+                    "a"
+                } else {
+                // comment2
+                    "b"
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testWithEolComments2() {
+        val actual = format(
+            """
+            fun test() {
+                val s = if (x > 0)
+                    // comment1
+                    "a" // comment2
+                else
+                    // comment3
+                    "b" // comment4
+            }
+            """.trimIndent()
+        )
+        assertThat(actual).isEqualTo(
+            """
+            fun test() {
+                val s = if (x > 0) {
+                    // comment1
+                    "a" // comment2
+                } else {
+                    // comment3
+                    "b" // comment4
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
     private fun assertOK(kotlinScript: String) {
         assertThat(format(kotlinScript)).isEqualTo(kotlinScript)
         assertThat(lint(kotlinScript)).isEqualTo(emptyList<LintError>())