MultiLineIfElseRule: don't insert unnecessary empty lines

pinterest · Jun 27, 2020 · c7bdae0 · c7bdae0
1 parent f62f7b3
commit c7bdae0
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 9 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
@@ -1,14 +1,12 @@
 package com.pinterest.ktlint.ruleset.experimental
 
 import com.pinterest.ktlint.core.Rule
-import com.pinterest.ktlint.core.ast.ElementType.BLOCK
 import com.pinterest.ktlint.core.ast.ElementType.ELSE
 import com.pinterest.ktlint.core.ast.ElementType.ELSE_KEYWORD
-import com.pinterest.ktlint.core.ast.ElementType.IF
 import com.pinterest.ktlint.core.ast.ElementType.LBRACE
 import com.pinterest.ktlint.core.ast.ElementType.RBRACE
 import com.pinterest.ktlint.core.ast.ElementType.THEN
-import com.pinterest.ktlint.core.ast.parent
+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
@@ -43,12 +41,9 @@ class MultiLineIfElseRule : Rule("multiline-if-else") {
 
     private fun autocorrect(node: ASTNode) {
         val bodyIndent = node.treePrev.text
-        val rightBraceIndent = when {
-            // in case of else if, get the indentation from the first if
-            node.treeParent.treeParent.elementType == ELSE -> node.parent({ it.elementType == IF && it.treeParent.elementType == BLOCK })!!.treePrev.text
-            node.treeParent.treePrev is PsiWhiteSpace -> node.treeParent.treePrev.text
-            else -> "\n"
-        }
+        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

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
@@ -251,6 +251,113 @@ class MultiLineIfElseRuleTest {
         )
     }
 
+    @Test
+    fun testWithEmptyLineBeforeIfExpression() {
+        val ifElseWithoutCurlyBrace =
+            """
+            fun test(): Int {
+                val b = foo()
+
+                if (b)
+                    return 1
+                else
+                    return 2
+            }
+            """.trimIndent()
+        assertThat(format(ifElseWithoutCurlyBrace)).isEqualTo(
+            """
+            fun test(): Int {
+                val b = foo()
+
+                if (b) {
+                    return 1
+                } else {
+                    return 2
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testInReturnExpression() {
+        val ifElseWithoutCurlyBrace =
+            """
+            fun test(i: Int, j: Int): Int {
+                return if (i == 1)
+                    if (j == 1)
+                        1
+                    else
+                        2
+                else if (i == 2)
+                    if (j == 1)
+                        3
+                    else
+                        4
+                else
+                    if (j == 1)
+                        5
+                    else
+                        6
+            }
+            """.trimIndent()
+        assertThat(format(ifElseWithoutCurlyBrace)).isEqualTo(
+            """
+            fun test(i: Int, j: Int): Int {
+                return if (i == 1) {
+                    if (j == 1) {
+                        1
+                    } else {
+                        2
+                    }
+                } else if (i == 2) {
+                    if (j == 1) {
+                        3
+                    } else {
+                        4
+                    }
+                } else {
+                    if (j == 1) {
+                        5
+                    } else {
+                        6
+                    }
+                }
+            }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testInLambdaExpression() {
+        val ifElseWithoutCurlyBrace =
+            """
+            fun test(s: String?): Int {
+                val i = s.let {
+                    if (it == "")
+                        1
+                    else
+                        2
+                } ?: 0
+                return i
+            }
+            """.trimIndent()
+        assertThat(format(ifElseWithoutCurlyBrace)).isEqualTo(
+            """
+            fun test(s: String?): Int {
+                val i = s.let {
+                    if (it == "") {
+                        1
+                    } else {
+                        2
+                    }
+                } ?: 0
+                return i
+            }
+            """.trimIndent()
+        )
+    }
+
     private fun assertOK(kotlinScript: String) {
         assertThat(format(kotlinScript)).isEqualTo(kotlinScript)
         assertThat(lint(kotlinScript)).isEqualTo(emptyList<LintError>())